mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
247 lines
15 KiB
Python
247 lines
15 KiB
Python
from __future__ import annotations
|
|
import struct, random, ctypes, functools, itertools
|
|
from tinygrad.helpers import getenv, wait_cond, round_up, next_power2, ceildiv, DEBUG, hi32, lo32, to_be32, to_be64
|
|
from tinygrad.runtime.support.memory import BumpAllocator
|
|
from tinygrad.runtime.support.system import PCIDevice, ipv4_to_gid
|
|
from tinygrad.runtime.autogen import mlx5, pci
|
|
|
|
MLX_DEBUG = getenv("MLX_DEBUG", 0)
|
|
|
|
MLX5_CMD_STRUCTS = {v: (getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_in_bits", None),
|
|
getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_out_bits", None)) for n, v in mlx5.__dict__.items() if n.startswith("MLX5_CMD_OP_")}
|
|
MLX5_CMD_STRUCTS[mlx5.MLX5_CMD_OP_ACCESS_REG] = (mlx5.struct_mlx5_ifc_access_register_in_bits, mlx5.struct_mlx5_ifc_access_register_out_bits)
|
|
|
|
def udp_sport(lqpn, rqpn):
|
|
v = (lqpn * rqpn ^ ((lqpn * rqpn) >> 20) ^ ((lqpn * rqpn) >> 40)) & 0xFFFFF
|
|
return ((v & 0x3FFF) ^ ((v & 0xFC000) >> 14)) | 0xC000
|
|
|
|
def ifc_get(buf, bit_off, width):
|
|
byte_off, bit_in, n = bit_off // 8, bit_off % 8, (bit_off % 8 + width + 7) // 8
|
|
return (int.from_bytes(buf[byte_off:byte_off + n], 'big') >> (n * 8 - bit_in - width)) & ((1 << width) - 1)
|
|
|
|
def ifc_set(buf, bit_off, width, value):
|
|
byte_off, bit_in, n = bit_off // 8, bit_off % 8, (bit_off % 8 + width + 7) // 8
|
|
shift, val = n * 8 - bit_in - width, int.from_bytes(buf[byte_off:byte_off + n], 'big')
|
|
buf[byte_off:byte_off + n] = ((val & ~(((1 << width) - 1) << shift)) | ((value & ((1 << width) - 1)) << shift)).to_bytes(n, 'big')
|
|
|
|
@functools.cache
|
|
def ifc_fields(ifc_struct): return {name: (off, ctypes.sizeof(typ)) for name, typ, off in ifc_struct._real_fields_ if not name.startswith('reserved')}
|
|
|
|
def ifc_subfield(ifc_struct, field_name):
|
|
for name, typ, off in ifc_struct._real_fields_:
|
|
if name == field_name: return typ, ifc_fields(ifc_struct)[field_name][0]
|
|
raise KeyError(f"no field '{field_name}' in {ifc_struct}")
|
|
|
|
def fill_ifc(buf, ifc_struct, base=0, **kw):
|
|
fields = ifc_fields(ifc_struct)
|
|
for name, val in kw.items():
|
|
if isinstance(val, dict):
|
|
sub_struct, sub_off = ifc_subfield(ifc_struct, name)
|
|
fill_ifc(buf, sub_struct, base=base + sub_off, **val)
|
|
else: ifc_set(buf, base + fields[name][0], fields[name][1], val)
|
|
|
|
def ifc_decode(buf, ifc_struct, base=0):return {name: ifc_get(buf, base + off, width) for name, (off, width) in ifc_fields(ifc_struct).items()}
|
|
|
|
class MLXCmdQueue:
|
|
def __init__(self, dev):
|
|
self.dev, self._tok = dev, itertools.count(1)
|
|
|
|
cmd_l = dev.iseg_r('cmdq_addr_l_sz') & 0xFF
|
|
self.log_stride, self.max_reg_cmds = cmd_l & 0xF, (1 << ((cmd_l >> 4) & 0xF)) - 1
|
|
|
|
stride = next_power2(ctypes.sizeof(mlx5.struct_mlx5_cmd_prot_block))
|
|
self.queue, self.queue_paddrs = dev.pci_dev.alloc_sysmem(0x1000 + 1024 * stride)
|
|
self.mboxes = [(off:=0x1000 + i * stride, self.queue_paddrs[1 + (i * stride) // 0x1000] + (off % 0x1000)) for i in range(1024)]
|
|
|
|
dev.iseg_w('cmdq_addr_h', hi32(self.queue_paddrs[0]))
|
|
dev.iseg_w('cmdq_addr_l_sz', lo32(self.queue_paddrs[0]) | cmd_l)
|
|
|
|
def create_mbox_chain(self, base, tok, data):
|
|
n = ceildiv(len(data), chunk_sz:=mlx5.MLX5_CMD_DATA_BLOCK_SIZE)
|
|
for i in range(n):
|
|
off, _ = self.mboxes[base + i]
|
|
blk = mlx5.struct_mlx5_cmd_prot_block(data=(ctypes.c_ubyte*chunk_sz).from_buffer_copy(data[i*chunk_sz:(i+1)*chunk_sz].ljust(chunk_sz, b'\x00')),
|
|
next=to_be64(self.mboxes[base+i+1][1]) if i < n-1 else 0, block_num=to_be32(i), token=tok)
|
|
self.queue[off:off + ctypes.sizeof(mlx5.struct_mlx5_cmd_prot_block)] = bytes(blk)
|
|
return (self.mboxes[base][0], self.mboxes[base][1], n)
|
|
|
|
def exec(self, opcode, op_mod=0, payload=b'', raw=False, **kw):
|
|
in_struct, out_struct = MLX5_CMD_STRUCTS[opcode]
|
|
out_sz = max(0, ctypes.sizeof(out_struct) - 16) if out_struct else 0
|
|
tok, slot = (next(self._tok) % 255) + 1, self.max_reg_cmds if opcode == mlx5.MLX5_CMD_OP_MANAGE_PAGES else 0
|
|
|
|
# serialize input
|
|
inp_sz = max(16, ceildiv(max((off + w for off, w in ifc_fields(in_struct).values()), default=0), 8))
|
|
fill_ifc(inp:=bytearray(inp_sz + len(payload)), in_struct, opcode=opcode, op_mod=op_mod, **kw)
|
|
if payload: inp[inp_sz:] = payload
|
|
|
|
# prepare mailboxes and build command layout
|
|
_, in_ptr, n_in = self.create_mbox_chain(0, tok, inp[16:])
|
|
_, out_ptr, n_out = self.create_mbox_chain(n_in, tok, bytes(out_sz))
|
|
cmd = mlx5.struct_mlx5_cmd_layout(type=mlx5.MLX5_PCI_CMD_XPORT, inlen=to_be32(len(inp)), in_ptr=to_be64(in_ptr),
|
|
_in=(ctypes.c_uint32*4)(*(int.from_bytes(inp[i:i+4], 'little') for i in range(0, 16, 4))),
|
|
out_ptr=to_be64(out_ptr), outlen=to_be32(16 + out_sz), token=tok, status_own=mlx5.CMD_OWNER_HW)
|
|
cmd_bytes = bytearray(bytes(cmd))
|
|
cmd_bytes[mlx5.struct_mlx5_cmd_layout.sig.offset] = (~functools.reduce(lambda a, b: a ^ b, cmd_bytes)) & 0xFF # type: ignore[attr-defined]
|
|
|
|
# submit and wait for completion
|
|
slot_view = self.queue.view(slot << self.log_stride, len(cmd_bytes))
|
|
slot_view[:] = cmd_bytes
|
|
self.dev.iseg_w('cmd_dbell', 1 << slot)
|
|
wait_cond(lambda: slot_view[mlx5.struct_mlx5_cmd_layout.status_own.offset] & mlx5.CMD_OWNER_HW, value=0, # type: ignore[attr-defined]
|
|
msg=f"cmd 0x{opcode:04x}")
|
|
|
|
# check status and read output
|
|
assert slot_view[mlx5.struct_mlx5_cmd_layout.status_own.offset] >> 1 == 0, f"cmd 0x{opcode:04x} delivery error" # type: ignore[attr-defined]
|
|
|
|
out_view = slot_view.view(mlx5.struct_mlx5_cmd_layout.out.offset, 16 + out_sz) # type: ignore[attr-defined]
|
|
status, syndrome = struct.unpack('>I', out_view[0:4])[0] >> 24, struct.unpack('>I', out_view[4:8])[0]
|
|
assert status == 0, f"cmd 0x{opcode:04x} failed status=0x{status:x} syn=0x{syndrome:08x}"
|
|
|
|
ret = bytearray(out_view[:16])
|
|
ret += b''.join(bytes(self.queue[va:va + mlx5.MLX5_CMD_DATA_BLOCK_SIZE]) for va, _ in self.mboxes[n_in:n_in+n_out])[:out_sz]
|
|
return ret if raw else (ifc_decode(ret, out_struct) if out_struct else ret)
|
|
|
|
class MLXDev:
|
|
def __init__(self, pci_dev:PCIDevice, ip:str=getenv("MLX_IP", "10.0.0.1")):
|
|
self.pci_dev, self.devfmt, self.bar = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
|
|
|
|
fw_rev, cmdif_sub = self.iseg_r('fw_rev'), self.iseg_r('cmdif_rev_fw_sub')
|
|
if DEBUG >= 2: print(f"mlx5 {self.devfmt}: firmware {fw_rev >> 16}.{fw_rev & 0xFFFF}.{cmdif_sub & 0xFFFF}")
|
|
assert (cmdif_sub >> 16) == 5, f"unsupported mlx version: {cmdif_sub >> 16}"
|
|
|
|
self.init_hw(ip)
|
|
|
|
def rreg(self, off): return to_be32(self.bar[off // 4])
|
|
def wreg(self, off, val): self.bar[off // 4] = to_be32(val)
|
|
def iseg_r(self, field): return self.rreg(getattr(mlx5.struct_mlx5_init_seg, field).offset)
|
|
def iseg_w(self, field, val): self.wreg(getattr(mlx5.struct_mlx5_init_seg, field).offset, val)
|
|
|
|
def init_hw(self, ip):
|
|
wait_cond(lambda: self.iseg_r('initializing') & 0x80000000, value=0, msg="FW init timeout")
|
|
self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
|
|
self.cmd = MLXCmdQueue(self)
|
|
wait_cond(lambda: self.iseg_r('initializing') & 0x80000000, value=0, msg="FW init timeout")
|
|
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_ENABLE_HCA)
|
|
if self.cmd.exec(mlx5.MLX5_CMD_OP_QUERY_ISSI)['supported_issi_dw0'] & 2:
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_SET_ISSI, current_issi=1)
|
|
|
|
self.provide_pages(mlx5.MLX5_BOOT_PAGES)
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_ACCESS_REG, register_id=mlx5.MLX5_REG_HOST_ENDIANNESS, payload=bytearray(16))
|
|
|
|
self.init_hca()
|
|
|
|
self.uar = self.cmd.exec(mlx5.MLX5_CMD_OP_ALLOC_UAR)['uar']
|
|
self.uar_view = self.pci_dev.map_bar(0, off=self.uar * 0x1000, size=0x1000, fmt='Q')
|
|
|
|
vport = self.cmd.exec(mlx5.MLX5_CMD_OP_QUERY_NIC_VPORT_CONTEXT, raw=True)
|
|
nvc_struct, nvc_off = ifc_subfield(mlx5.struct_mlx5_ifc_query_nic_vport_context_out_bits, 'nic_vport_context')
|
|
mac_struct, mac_off = ifc_subfield(nvc_struct, 'permanent_address')
|
|
self.mac = ifc_get(vport, nvc_off + mac_off + 16, 48)
|
|
|
|
# enable roce
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_MODIFY_NIC_VPORT_CONTEXT, field_select=dict(roce_en=1), nic_vport_context=dict(roce_en=1))
|
|
|
|
dbr_mem, self.dbr_paddrs = self.pci_dev.alloc_sysmem(0x1000)
|
|
self.dbr = dbr_mem.view(fmt='I')
|
|
self.dbr_alloc = BumpAllocator(0x1000, wrap=False)
|
|
|
|
self.pd = self.cmd.exec(mlx5.MLX5_CMD_OP_ALLOC_PD)['pd']
|
|
res = self.cmd.exec(mlx5.MLX5_CMD_OP_CREATE_MKEY, memory_key_mkey_entry=dict(access_mode_1_0=0, rw=1, rr=1, lw=1, lr=1, qpn=0xFFFFFF,
|
|
mkey_7_0=(key_lo:=0x22), length64=1, pd=self.pd))
|
|
self.mkey = (res['mkey_index'] << 8) | key_lo
|
|
|
|
self.local_gid = ipv4_to_gid(ip)
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_SET_ROCE_ADDRESS, roce_address=dict(roce_version=2, source_l3_address=int.from_bytes(self.local_gid, 'big'),
|
|
roce_l3_type=0, source_mac_47_32=hi32(self.mac), source_mac_31_0=lo32(self.mac)), roce_address_index=0, vhca_port_num=1)
|
|
|
|
if DEBUG >= 2: print(f"mlx5 {self.devfmt}: booted mac={self.mac.to_bytes(6,'big').hex(':')} mkey=0x{self.mkey:x}")
|
|
|
|
def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
|
|
n = len(paddrs)
|
|
mtt = struct.pack(f'>{round_up(n, 2)}Q', *paddrs, *([0] * (round_up(n, 2) - n)))
|
|
if MLX_DEBUG >= 1: print(f"mlx5 {self.devfmt}: register_mem pages={n} page_sz={1 << log_page_size} mtt_bytes={len(mtt)}")
|
|
self.provide_pages(mlx5.MLX5_INIT_PAGES)
|
|
res = self.cmd.exec(mlx5.MLX5_CMD_OP_CREATE_MKEY, translations_octword_actual_size=ceildiv(n, 2), payload=mtt,
|
|
memory_key_mkey_entry=dict(access_mode_1_0=1, lr=1, lw=1, rr=1, rw=1, pd=self.pd, qpn=0xFFFFFF, mkey_7_0=(key_lo:=0x33),
|
|
start_addr=paddrs[0], len=size, log_page_size=log_page_size, translations_octword_size=ceildiv(n, 2)))
|
|
return (res['mkey_index'] << 8) | key_lo
|
|
|
|
def unregister_mem(self, mkey:int): self.cmd.exec(mlx5.MLX5_CMD_OP_DESTROY_MKEY, mkey_index=mkey >> 8)
|
|
|
|
def provide_pages(self, mode):
|
|
if (npages:=self.cmd.exec(mlx5.MLX5_CMD_OP_QUERY_PAGES, op_mod=mode)['num_pages']) <= 0: return
|
|
if MLX_DEBUG >= 1: print(f"mlx5 {self.devfmt}: provide_pages mode={mode}, {npages} pages")
|
|
mem, paddrs = self.pci_dev.alloc_sysmem(npages * 0x1000)
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_MANAGE_PAGES, op_mod=mlx5.MLX5_PAGES_GIVE, input_num_entries=npages, payload=struct.pack(f'>{npages}Q', *paddrs))
|
|
|
|
def hca_query_cap(self, cap_type, cap_struct, mode):
|
|
raw = bytearray(self.cmd.exec(mlx5.MLX5_CMD_OP_QUERY_HCA_CAP, op_mod=(cap_type << 1) | mode, raw=True)[16:16+4096])
|
|
return raw, ifc_decode(raw, cap_struct)
|
|
|
|
def hca_set_cap(self, cap_type, cap_struct, raw, **kwargs):
|
|
fill_ifc(cap:=bytearray(raw), cap_struct, **kwargs)
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_SET_HCA_CAP, op_mod=cap_type << 1, capability=int.from_bytes(cap[:4096].ljust(4096, b'\x00'), 'big'))
|
|
|
|
def init_hca(self):
|
|
gen_caps, gen_cur = self.hca_query_cap(mlx5.MLX5_CAP_GENERAL, mlx5.struct_mlx5_ifc_cmd_hca_cap_bits, mode=1)
|
|
self.hca_set_cap(mlx5.MLX5_CAP_GENERAL, mlx5.struct_mlx5_ifc_cmd_hca_cap_bits, gen_caps,
|
|
pkey_table_size=0, cmdif_checksum=0, log_uar_page_sz=0, log_max_qp=18, roce=1)
|
|
|
|
roce_cur_raw, roce_cur = self.hca_query_cap(mlx5.MLX5_CAP_ROCE, mlx5.struct_mlx5_ifc_roce_cap_bits, mode=1)
|
|
self.hca_set_cap(mlx5.MLX5_CAP_ROCE, mlx5.struct_mlx5_ifc_roce_cap_bits, roce_cur_raw, sw_r_roce_src_udp_port=1)
|
|
|
|
self.provide_pages(mlx5.MLX5_INIT_PAGES)
|
|
self.cmd.exec(mlx5.MLX5_CMD_OP_INIT_HCA, sw_owner_id=random.getrandbits(128))
|
|
|
|
_, self.caps = self.hca_query_cap(mlx5.MLX5_CAP_GENERAL, mlx5.struct_mlx5_ifc_cmd_hca_cap_bits, 1)
|
|
|
|
if MLX_DEBUG >= 4: print(f"mlx5 {self.devfmt}: HCA initialized with gen_caps={gen_cur} roce_caps={roce_cur}")
|
|
|
|
class MLXQP:
|
|
def __init__(self, dev:MLXDev, log_sq_size=4, log_rq_size=4, log_eq_size=7, log_cq_size=7):
|
|
self.dev, self.cq_size, self.log_sq_size, self.log_rq_size, self.head = dev, 1 << log_cq_size, log_sq_size, log_rq_size, 0
|
|
|
|
self.cq_dbr, self.qp_dbr = dev.dbr_alloc.alloc(8, alignment=8), dev.dbr_alloc.alloc(8, alignment=8)
|
|
|
|
# create EQ, CQ
|
|
self.eq_mem, self.eq_paddrs, self.eq_info = self.create_queue(mlx5.MLX5_CMD_OP_CREATE_EQ, log_eq_size, entry_sz=64, owner_off=31,
|
|
eq_context_entry=dict(log_eq_size=log_eq_size, uar_page=dev.uar, log_page_size=0))
|
|
|
|
self.cq_mem, self.cq_paddrs, self.cq_info = self.create_queue(mlx5.MLX5_CMD_OP_CREATE_CQ, log_cq_size, entry_sz=64, owner_off=63,
|
|
cq_context=dict(log_cq_size=log_cq_size, uar_page=dev.uar, c_eqn_or_apu_element=self.eq_info['eq_number'],
|
|
dbr_addr=dev.dbr_paddrs[0] + self.cq_dbr, log_page_size=0))
|
|
|
|
# create QP, buffer is RQ (16B stride) + SQ (64B stride)
|
|
self.sq_offset = (1 << log_rq_size) << 4
|
|
self.qp_buf, self.qp_paddrs, self.qp_info = self.create_queue(mlx5.MLX5_CMD_OP_CREATE_QP, log_sq_size, entry_sz=64,
|
|
owner_off=0, extra_sz=self.sq_offset,
|
|
qpc=dict(st=0, pm_state=3, pd=dev.pd, cqn_snd=self.cq_info['cqn'], cqn_rcv=self.cq_info['cqn'], log_msg_max=30, log_rq_size=log_rq_size,
|
|
log_rq_stride=0, log_sq_size=log_sq_size, rlky=1, uar_page=dev.uar, log_page_size=0, dbr_addr=dev.dbr_paddrs[0] + self.qp_dbr))
|
|
|
|
# transition to INIT
|
|
self.qp_op(mlx5.MLX5_CMD_OP_RST2INIT_QP, qpc_args=dict(log_ack_req_freq=8), addr_args=dict(pkey_index=0, vhca_port_num=1))
|
|
|
|
for i in range(self.cq_size): self.cq_mem[i * 64 + 63] = 0x01 # init owner bits so poll_cq waits for real CQEs
|
|
if MLX_DEBUG >= 1: print(f"mlx5: QP 0x{self.qp_info['qpn']:x} (EQ={self.eq_info['eq_number']} CQ=0x{self.cq_info['cqn']:x})")
|
|
|
|
def create_queue(self, opcode, log_size, entry_sz, owner_off, extra_sz=0, **ctx_kw):
|
|
mem, paddrs = self.dev.pci_dev.alloc_sysmem((n := ceildiv((1 << log_size) * entry_sz + extra_sz, 0x1000)) * 0x1000)
|
|
return mem, paddrs, self.dev.cmd.exec(opcode, payload=struct.pack(f'>{n}Q', *paddrs), **ctx_kw)
|
|
|
|
def qp_op(self, opcode, qpc_args=None, addr_args=None, **kwargs):
|
|
qpc_args = dict(st=0, pm_state=3, pd=self.dev.pd, cqn_snd=self.cq_info['cqn'], cqn_rcv=self.cq_info['cqn'], **(qpc_args or {}))
|
|
self.dev.cmd.exec(opcode, qpn=self.qp_info['qpn'], qpc=(qpc_args or {}) | {'primary_address_path': addr_args or {}}, **kwargs)
|
|
|
|
def connect(self, remote:MLXQP):
|
|
self.qp_op(mlx5.MLX5_CMD_OP_INIT2RTR_QP, opt_param_mask=0x1A,
|
|
qpc_args=dict(mtu=5, log_msg_max=self.dev.caps['log_max_msg'], remote_qpn=remote.qp_info['qpn'], log_ack_req_freq=8,
|
|
log_rra_max=3, rre=1, rwe=1, min_rnr_nak=1, next_rcv_psn=0),
|
|
addr_args=dict(pkey_index=0, src_addr_index=0, hop_limit=64, udp_sport=udp_sport(self.qp_info['qpn'], remote.qp_info['qpn']), vhca_port_num=1,
|
|
rmac_47_32=hi32(remote.dev.mac), rmac_31_0=lo32(remote.dev.mac), rgid_rip=int.from_bytes(remote.dev.local_gid, 'big')))
|
|
self.qp_op(mlx5.MLX5_CMD_OP_RTR2RTS_QP, qpc_args=dict(log_ack_req_freq=8, next_send_psn=0, log_sra_max=3, retry_count=7, rnr_retry=7),
|
|
addr_args=dict(ack_timeout=14, vhca_port_num=1))
|
|
|
|
if MLX_DEBUG >= 1: print(f"mlx5: QP 0x{self.qp_info['qpn']:x} connected (remote=0x{remote.qp_info['qpn']:x})")
|