StarPilot/tinygrad_repo/tinygrad/runtime/support/mlx/mlxdev.py

from __future__ import annotations
import struct, random, socket, 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
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 ipv4_to_gid(ip): return bytes(10) + b'\xff\xff' + socket.inet_aton(ip)

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})")