vf MI350X passthrough mode driver

extra/amdpci/hive_reset.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
 if __name__ == "__main__":
-  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0, 0x75b0])])
   for gpu in gpus:
     drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
     if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":

tinygrad/runtime/ops_amd.py

@@ -842,7 +842,7 @@ class KFDIface:
 
 class PCIIface(PCIIfaceBase):
   def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0,0x75b0)),), vram_bar=0,
       va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
     self._compute_props()
 
@@ -1096,6 +1096,11 @@ class AMDDevice(HCQCompiled):
 
   def on_device_hang(self): self.iface.on_device_hang()
 
+  def finalize(self):
+    try: super().finalize()
+    finally:
+      if self.is_am(): self.iface.dev_impl.release_vf_access()
+
   def device_props(self): return self.iface.props
 
   def hw_copy_queues(self): return [(f"SDMA:{i}", functools.partial(unwrap(self.hw_copy_queue_t), queue_idx=i)) for i in self.sdma_queues]

tinygrad/runtime/support/am/amdev.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, hashlib, array
+import ctypes, collections, dataclasses, functools, hashlib, array, time, contextlib
 from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am, fw
@@ -149,6 +149,14 @@ class AMDev:
     self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
 
+    # MI350X VFs start with most MMIO and VRAM access gated by the host PF. Ask the PF for full access only when discovery isn't readable yet.
+    self.is_vf = bool(self.mmio[0xde5] & 1) # RCC_IOV_FUNC_IDENTIFIER.FUNC_IDENTIFIER
+    self.vf_access_acquired, self.vf_initialized = False, False
+    if self.is_vf:
+      self._vf_mailbox_request(6, 7, data2=2, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_DATA -> IDH_REQ_GPU_INIT_DATA_READY
+      self._vf_mailbox_request(1, 1, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_ACCESS -> IDH_READY_TO_ACCESS_GPU
+      self.vf_access_acquired = True
+
     self._run_discovery()
     self._build_regs()
 
@@ -165,21 +173,21 @@ class AMDev:
     self.is_booting = True # During boot only boot memory can be allocated. This flag is to validate this.
     self.init_sw(smi_dev=False)
 
-    self.partial_boot = (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
+    self.partial_boot = not self.is_vf and (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
     if self.partial_boot and (self.reg("regSCRATCH_REG6").read() != 0 or self.reg(self.gmc.pf_status_reg("GC")).read() != 0):
       if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
       self.partial_boot = False
 
-    # Init hw for IP blocks where it is needed
+    # Init hw for IP blocks where it is needed. PSP and SMU are PF-owned on a VF and must not be reset or reloaded by the guest.
     if not self.partial_boot:
-      if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+      if not self.is_vf and self.psp.is_sos_alive() and self.smu.is_smu_alive():
         self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
         self.smu.mode1_reset()
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-      self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
+      self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))
 
     # Booting done
     self.is_booting = False
@@ -187,13 +195,17 @@ class AMDev:
     # Re-initialize main blocks
     self.init_hw(self.gfx, self.sdma)
 
-    if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
-      self.smu.set_power_limit(max_power)
-      self.smu.set_clocks(level=None)
-    else: self.smu.set_clocks(level=-1) # last level, max perf.
-    for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
-    self.reg("regSCRATCH_REG7").write(AMDev.Version)
-    self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    if not self.is_vf:
+      if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
+        self.smu.set_power_limit(max_power)
+        self.smu.set_clocks(level=None)
+      else: self.smu.set_clocks(level=-1) # last level, max perf.
+    if not self.is_vf:
+      for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
+    if not self.is_vf:
+      self.reg("regSCRATCH_REG7").write(AMDev.Version)
+      self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    self.vf_initialized = self.is_vf
     if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
 
   def init_sw(self, smi_dev=False):
@@ -202,7 +214,8 @@ class AMDev:
     # Memory manager & firmware
     self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
       va_bits=48, first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base, reserve_ptable=not self.large_bar,
-      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
+      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)],
+      paddr_base=(1 << 20) if self.is_vf else 0)
     self.fw = AMFirmware(self)
 
     # Initialize IP blocks
@@ -224,10 +237,52 @@ class AMDev:
 
   def fini(self):
     if DEBUG >= 2: print(f"am {self.devfmt}: Finalizing")
-    for ip in [self.sdma, self.gfx]: ip.fini_hw()
-    self.smu.set_clocks(level=0)
-    self.ih.interrupt_handler()
-    self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    try:
+      for ip in [self.sdma, self.gfx]: ip.fini_hw()
+      if not self.is_vf: self.smu.set_clocks(level=0)
+      self.ih.interrupt_handler()
+      if not self.is_vf: self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    finally: self.release_vf_access()
+
+  def release_vf_access(self):
+    if not getattr(self, "vf_access_acquired", False): return
+    # tinygrad retains IDH_REQ_GPU_INIT_ACCESS for direct MMIO/VRAM access, so always release that same lease.
+    with contextlib.suppress(Exception): self._vf_mailbox_request(2, None) # IDH_REL_GPU_INIT_ACCESS
+    self.vf_access_acquired = False
+
+  def __del__(self):
+    # Constructor failures do not reach HCQ finalization; return a partially acquired VF init lease to the PF.
+    self.release_vf_access()
+
+  def _vf_mailbox_request(self, req:int, event:int|None, data1=0, data2=0, data3=0, retries=1, event_timeout=2.0):
+    # Navi VF/PF mailbox protocol from the kernel's mxgpu_nv driver. This requests access only; it never requests a GPU or PCI reset.
+    mmio8, control, trn, rcv = self.mmio.view(fmt='B'), 0xe5e * 4, 0xe56, 0xe5a
+    if mmio8[control+1] & 1: mmio8[control+1] = 2 # acknowledge a stale PF event before transmitting a new request
+    for retry in range(retries):
+      deadline = time.monotonic() + 1
+      while True:
+        mmio8[control] = 0 # clear TRN_MSG_VALID and wait for the old PF acknowledgement to drop
+        if not (mmio8[control] & 2): break
+        if time.monotonic() > deadline: raise TimeoutError("VF mailbox acknowledgement did not clear")
+        time.sleep(0.001)
+
+      for i, val in enumerate((req, data1, data2, data3)): self.mmio[trn+i] = val
+      mmio8[control] = 1
+      deadline = time.monotonic() + 0.5
+      while not (mmio8[control] & 2):
+        if time.monotonic() > deadline: raise TimeoutError(f"VF mailbox request {req:#x} was not acknowledged")
+        time.sleep(0.005)
+      mmio8[control] = 0
+
+      if event is None: return
+      deadline = time.monotonic() + event_timeout
+      while time.monotonic() <= deadline:
+        if self.mmio[rcv] == event:
+          mmio8[control+1] = 2 # acknowledge RCV_MSG_VALID
+          return
+        time.sleep(0.01)
+      if DEBUG >= 2 and retry+1 < retries: print(f"am {self.devfmt}: retrying VF mailbox request {req:#x} ({retry+1}/{retries})")
+    raise TimeoutError(f"VF mailbox request {req:#x} did not receive event {event:#x}")
 
   def recover(self, force=False) -> bool:
     if not force and not self.is_err_state: return False

tinygrad/runtime/support/am/ip.py

@@ -251,9 +251,10 @@ class AM_GFX(AM_IP):
     self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]
 
   def init_hw(self):
-    # Wait for RLC autoload to complete
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
-              value=True, msg="RLC autoload timeout")
+    # Wait for RLC autoload to complete on architectures that expose the bootload status register.
+    if hasattr(self.adev, "regRLC_RLCS_BOOTLOAD_STATUS"):
+      wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
+                value=True, msg="RLC autoload timeout")
 
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
     if self.adev.partial_boot: return self.reset_mec()
@@ -297,8 +298,8 @@ class AM_GFX(AM_IP):
 
     self._enable_mec()
 
-    # Set 1 partition
-    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
+    # Set 1 partition on bare metal. A VF must use the spatial partition assigned by its host PF.
+    if self.xccs > 1 and not self.adev.is_vf: self.adev.psp._spatial_partition_cmd(1)
 
   def fini_hw(self): self._dequeue_hqds()
 
@@ -486,7 +487,7 @@ class AM_IH(AM_IP):
     if athub_err or cntlr_err:
       print(f"am {self.adev.devfmt}: fatal hardware error detected: {'RAS_ATHUB_ERR_EVENT ' if athub_err else ''}{'RAS_CNTLR' if cntlr_err else ''}")
 
-      acas = self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
+      acas = [] if self.adev.is_vf else self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
       for regs in acas:
         acatyp = 'Uncorrectable' if (regs[1] >> 61) & 1 and (regs[1] >> 57) & 1 else 'Correctable'
         hwname = f'{self.adev.hwid_names.get((regs[5] >> 32) & 0xFFF, "")} ({(regs[5] >> 32) & 0xFFF:#03x})'

tinygrad/runtime/support/memory.py

@@ -173,14 +173,16 @@ class MemoryManager:
   va_allocator: ClassVar[TLSFAllocator|None] = None
 
   def __init__(self, dev, vram_size:int, boot_size:int, pt_t, va_bits:int, va_shifts:list[int], va_base:int,
-               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False):
+               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False, paddr_base:int=0):
     self.dev, self.vram_size, self.va_shifts, self.va_base, lvl_msb = dev, vram_size, va_shifts, va_base, va_shifts + [va_bits + 1]
     self.pte_covers, self.pte_cnt = [1 << x for x in va_shifts][::-1], [1 << (lvl_msb[i+1] - lvl_msb[i]) for i in range(len(lvl_msb) - 1)][::-1]
     self.pt_t, self.palloc_ranges, self.level_cnt, self.va_bits, self.reserve_ptable = pt_t, palloc_ranges, len(va_shifts), va_bits, reserve_ptable
 
-    self.boot_allocator = TLSFAllocator(boot_size, base=0)
-    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0, base=self.boot_allocator.size)
-    self.pa_allocator = TLSFAllocator(vram_size - (off_sz:=self.boot_allocator.size + self.ptable_allocator.size), base=off_sz)
+    self.boot_allocator = TLSFAllocator(boot_size, base=paddr_base)
+    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0,
+                                          base=paddr_base + self.boot_allocator.size)
+    off_sz = paddr_base + self.boot_allocator.size + self.ptable_allocator.size
+    self.pa_allocator = TLSFAllocator(vram_size - off_sz, base=off_sz)
     self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)
 
   def _frag_size(self, va, sz, must_cover=True):