mi350p: aqua vanjaram bringup for the raw PCI driver

Minimum stable bringup, all root-caused:
- disable ASPM on the PCI path: link L1 across retimers makes GPU reads oscillate to 0xffffffff
- probe live AIDs via MMHUB FB_LOCATION; dead ones read 0xffffffff and indirect writes to them corrupt the fabric
- PSP DRAM training before bootloader loads on gc 9.5.0 (BL hangs on load-KDB otherwise)
- aqua reset semantics: full boot over a live state can kill the fabric, partial boot + reset_mec is the deepest safe reset

tinygrad/runtime/support/am/amdev.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, hashlib, array
+import os, re, ctypes, collections, dataclasses, functools, hashlib, array
 from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw, to_mv
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am, fw
@@ -145,8 +145,18 @@ class AMMemoryManager(MemoryManager):
 class AMDev:
   Version = 0xA0000008
 
+  @staticmethod
+  def _disable_aspm(pcibus:str): # link L1 across retimers makes reads oscillate to 0xffffffff (GPU looks dead)
+    chain = [pcibus] + re.findall(r"[0-9a-f]{4}:[0-9a-f]{2}:[0-9a-f]{2}\.[0-7]", os.path.dirname(os.path.realpath(f"/sys/bus/pci/devices/{pcibus}")))
+    for bus in chain:
+      fd, cap = os.open(f"/sys/bus/pci/devices/{bus}/config", os.O_RDWR), 0x34
+      while (ptr := os.pread(fd, 1, cap)[0] & 0xfc) and os.pread(fd, 1, ptr)[0] != 0x10: cap = ptr + 1 # PCIe cap: lnkctl at +0x10
+      if ptr: os.pwrite(fd, (int.from_bytes(os.pread(fd, 2, ptr + 0x10)) & ~3).to_bytes(2, 'little'), ptr + 0x10)
+      os.close(fd)
+
   def __init__(self, pci_dev:PCIDevice, reset_mode=False):
     self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    self._disable_aspm(self.devfmt)
     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')
 
     self._run_discovery()
@@ -170,6 +180,9 @@ class AMDev:
       if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
       self.partial_boot = False
 
+    # aqua (gc 9.5.0): full boot over live state can kill the fabric (power cycle recovers); partial boot+reset_mec is the deepest safe reset
+    if self.ip_ver[am.GC_HWIP] == (9,5,0) and self.reg("regSCRATCH_REG7").read() == AMDev.Version: self.partial_boot = True
+
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
       if self.psp.is_sos_alive() and self.smu.is_smu_alive():
@@ -342,3 +355,6 @@ class AMDev:
     for prefix, hwip in mods:
       self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))
     self.__dict__.update(import_asic_regs('mp', (11, 0, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
+
+    # Live AIDs: probe via MMHUB FB_LOCATION (dead ones read 0xffffffff; indirect writes to them corrupt the fabric)
+    self.aids = [i for i in range(len(self.regs_offset[am.MMHUB_HWIP])) if self.reg("regMMMC_VM_FB_LOCATION_BASE").read(inst=i) != 0xffffffff] or [0]

tinygrad/runtime/support/am/ip.py

@@ -32,7 +32,7 @@ class AM_SOC(AM_IP):
       # fence doorbells for harvested xccs (0xff & ~xcc_mask in the kernel); a fully-unharvested chip keeps the previous 0x0
       live_xccs = sum(1 << i for i in self.adev.regs_offset[am.GC_HWIP] if i not in self.adev.harvested[am.GC_HWIP] and i < 8)
       self.adev.regXCC_DOORBELL_FENCE.write(0xff & ~live_xccs)
-      for aid in range(1, self.adev.gmc.vmhubs):
+      for aid in self.adev.aids[1:]:
         self.adev.indirect_wreg_pcie(self.adev.regXCC_DOORBELL_FENCE.addr[0], self.adev.regXCC_DOORBELL_FENCE.encode(shub_slv_mode=1), aid=aid)
       self.adev.regBIFC_GFX_INT_MONITOR_MASK.write(0x7ff)
       self.adev.regBIFC_DOORBELL_ACCESS_EN_PF.write(0xfffff)
@@ -81,9 +81,11 @@ class AM_GMC(AM_IP):
     # MM hub is inited before any tlb flushes and is still valid during partial_boot, so set it to true
     self.hub_initted = {"MM": True, "GC": False}
 
+    self.mm_insts = self.adev.aids if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)} else list(range(self.vmhubs)) # dead mmhubs hang us
+
     self.pf_status_reg = lambda ip: f"reg{ip}VM_L2_PROTECTION_FAULT_STATUS{'_LO32' if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else ''}"
 
-  def init_hw(self): self.init_hub("MM", inst_cnt=self.vmhubs)
+  def init_hw(self): self.init_hub("MM", insts=self.mm_insts)
 
   def flush_hdp(self): self.adev.wreg(self.adev.reg("regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL").read() // 4, 0x0)
   def flush_tlb(self, ip:Literal["MM", "GC"], vmid, flush_type=0):
@@ -92,7 +94,7 @@ class AM_GMC(AM_IP):
     # Can't issue TLB invalidation if the hub isn't initialized.
     if not self.hub_initted[ip]: return
 
-    for inst in range(self.adev.gmc.vmhubs if ip == "MM" else self.adev.gfx.xccs):
+    for inst in (self.adev.gmc.mm_insts if ip == "MM" else range(self.adev.gfx.xccs)):
       if ip == "MM": wait_cond(lambda: self.adev.regMMVM_INVALIDATE_ENG17_SEM.read(inst=inst) & 0x1, value=1, msg="mm flush_tlb timeout")
 
       self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_REQ").write(flush_type=flush_type, per_vmid_invalidate_req=(1 << vmid), invalidate_l2_ptes=1,
@@ -117,9 +119,9 @@ class AM_GMC(AM_IP):
     self.adev.reg(f"reg{ip}VM_CONTEXT{vmid}_CNTL").write(0x1800000, **fault_flags, **en_def_flags, enable_context=1,
       page_table_depth=((2 if self.trans_futher else 3) - page_table.lv), page_table_block_size=9 if self.trans_futher else 0, inst=inst)
 
-  def init_hub(self, ip:Literal["MM", "GC"], inst_cnt:int):
+  def init_hub(self, ip:Literal["MM", "GC"], insts):
     # Init system apertures
-    for inst in range(inst_cnt):
+    for inst in insts:
       self.adev.reg(f"reg{ip}MC_VM_AGP_BASE").write(0, inst=inst)
       self.adev.reg(f"reg{ip}MC_VM_AGP_BOT").write(0xffffffffffff >> 24, inst=inst) # disable AGP
       self.adev.reg(f"reg{ip}MC_VM_AGP_TOP").write(0, inst=inst)
@@ -258,7 +260,7 @@ class AM_GFX(AM_IP):
     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)
+    self.adev.gmc.init_hub("GC", insts=range(self.xccs))
     if self.adev.partial_boot: return self.reset_mec()
 
     self._config_mec()
@@ -406,7 +408,11 @@ class AM_GFX(AM_IP):
         if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
           self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
           self.adev.regSPI_COMPUTE_QUEUE_RESET.write(0x1, inst=xcc)
-          if not self.adev.is_err_state: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+          if not self.adev.is_err_state:
+            try: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+            # kernel tolerates this too; a wedged wave can survive RESET_WAVES
+            except TimeoutError:
+              if DEBUG >= 2: print(f"am {self.adev.devfmt}: HQD dequeue timeout xcc{xcc} q{q}, continuing")
     self._grbm_select()
 
 class AM_IH(AM_IP):
@@ -517,7 +523,7 @@ class AM_SDMA(AM_IP):
         **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)
 
     if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      for aid_id in range(self.adev.gmc.vmhubs):
+      for aid_id in self.adev.aids:
         for dev_inst, (port, awid, offset, awaddr) in enumerate([(1, 0xe, 0xe, 0x1), (2, 0x8, 0x8, 0x2), (5, 0x9, 0x9, 0x8), (6, 0xa, 0xa, 0x9)]):
           entry = dev_inst + 1 + 4 * aid_id
           self.adev.reg(f"regDOORBELL0_CTRL_ENTRY_{entry}").write(**{f"bif_doorbell{entry}_range_size_entry": 20,
@@ -590,6 +596,12 @@ class AM_PSP(AM_IP):
       (am.PSP_FW_TYPE_PSP_RAS_DRV, am.PSP_BL__LOAD_RASDRV), (am.PSP_FW_TYPE_PSP_SOS, am.PSP_BL__LOAD_SOSDRV)]
 
     if not self.is_sos_alive():
+      # chips with two-stage DRAM training (aqua/gc 9.5+) need the BL to train vram before loads
+      if self.adev.ip_ver[am.GC_HWIP] in {(9,5,0)}:
+        self._wait_for_bootloader()
+        self.adev.reg(f"{self.reg_pref}_36").write(self.adev.mm.palloc(am.GDDR6_MEM_TRAINING_DATA_SIZE_IN_BYTES, boot=True) >> 20)
+        self.adev.reg(f"{self.reg_pref}_35").write(am.PSP_BL__DRAM_LONG_TRAIN)
+        wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_35").read() & 0x80000000, value=0x80000000, msg="mem training timeout", timeout_ms=300000)
       for fw, compid in sos_components: self._bootloader_load_component(fw, compid)
       wait_cond(self.is_sos_alive, value=True, msg="sOS failed to start")