sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

nv: wait for wpr to reset (#13282)

Browse Source 
* nv: wait for wpr to reet

* fix

* comment

* wai

* f

* fix
This commit is contained in:
nimlgen
2025-11-15 20:00:49 +08:00
committed by GitHub
parent 22c08b470c
commit d37e1fe065
2 changed files with 31 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/support/nv/ip.py
+8
View File
@@ -76,6 +76,10 @@ class NVRpcQueue:
raise RuntimeError(f"Timeout waiting for RPC response for command {cmd}")
class NV_FLCN(NV_IP):
def wait_for_reset(self):
wait_cond(lambda _: self.nvdev.NV_PGC6_AON_SECURE_SCRATCH_GROUP_05_PRIV_LEVEL_MASK.read_bitfields()['read_protection_level0'] == 1 and
self.nvdev.NV_PGC6_AON_SECURE_SCRATCH_GROUP_05[0].read() & 0xff == 0xff, "waiting for reset")
def init_sw(self):
self.nvdev.include("src/common/inc/swref/published/ampere/ga102/dev_gsp.h")
self.nvdev.include("src/common/inc/swref/published/ampere/ga102/dev_falcon_v4.h")
@@ -262,6 +266,10 @@ class NV_FLCN(NV_IP):
self.nvdev.NV_PFALCON_FALCON_RM.with_base(base).write(self.nvdev.chip_id)
class NV_FLCN_COT(NV_IP):
def wait_for_reset(self):
self.nvdev.include("src/common/inc/swref/published/blackwell/gb202/dev_therm.h")
wait_cond(lambda _: self.nvdev.NV_THERM_I2CS_SCRATCH.read() == 0xff, "waiting for reset")
def init_sw(self):
self.nvdev.include("src/common/inc/swref/published/ampere/ga102/dev_gsp.h")
self.nvdev.include("src/common/inc/swref/published/hopper/gh100/dev_falcon_v4.h")
tinygrad/runtime/support/nv/nvdev.py
+23 -19
View File
@@ -74,20 +74,8 @@ class NVDev(PCIDevImplBase):
self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
self.smi_dev, self.is_booting = False, True
self._early_init()
# UVM depth HW level VA bits
# 0 PDE4 56:56 (hopper+)
# 1 PDE3 55:47
# 2 PDE2 46:38
# 3 PDE1 (or 512M PTE) 37:29
# 4 PDE0 (dual 64k/4k PDE, or 2M PTE) 28:21
# 5 PTE_64K / PTE_4K 20:16 / 20:12
bits, shifts = (56, [12, 21, 29, 38, 47, 56]) if self.mmu_ver == 3 else (48, [12, 21, 29, 38, 47])
self.mm = NVMemoryManager(self, self.vram_size, boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=bits, va_shifts=shifts, va_base=0,
palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]], reserve_ptable=not self.large_bar)
self.flcn:NV_FLCN|NV_FLCN_COT = NV_FLCN_COT(self) if self.fmc_boot else NV_FLCN(self)
self.gsp:NV_GSP = NV_GSP(self)
self._early_ip_init()
self._early_mmu_init()
# Turn the booting early, gsp client is loaded from the clean.
self.is_booting = False
@@ -104,7 +92,7 @@ class NVDev(PCIDevImplBase):
if NV_DEBUG >= 4: print(f"wreg: {hex(addr)} = {hex(value)}")
def rreg(self, addr:int) -> int: return self.mmio[addr // 4]
def _early_init(self):
def _early_ip_init(self):
self.reg_names:set[str] = set()
self.reg_offsets:dict[str, tuple[int, int]] = {}
@@ -115,15 +103,20 @@ class NVDev(PCIDevImplBase):
self.fw_name = {"GB2": "GB202", "AD1": "AD102", "GA1": "GA102"}[self.chip_name[:3]]
self.mmu_ver, self.fmc_boot = (3, True) if self.chip_details['architecture'] >= 0x1a else (2, False)
self.flcn:NV_FLCN|NV_FLCN_COT = NV_FLCN_COT(self) if self.fmc_boot else NV_FLCN(self)
self.gsp:NV_GSP = NV_GSP(self)
self.include("src/common/inc/swref/published/turing/tu102/dev_fb.h")
self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island.h")
self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island_addendum.h")
if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0:
if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
self.pci_dev.reset()
time.sleep(0.5)
time.sleep(0.1) # wait until device can respond again
self.flcn.wait_for_reset()
def _early_mmu_init(self):
self.include("src/common/inc/swref/published/turing/tu102/dev_vm.h")
self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island.h")
self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island_addendum.h")
# MMU Init
self.reg_names.update(mmu_pd_names:=[f'NV_MMU_VER{self.mmu_ver}_PTE', f'NV_MMU_VER{self.mmu_ver}_PDE', f'NV_MMU_VER{self.mmu_ver}_DUAL_PDE'])
@@ -136,6 +129,17 @@ class NVDev(PCIDevImplBase):
self.vram, self.mmio = self.pci_dev.map_bar(1), self.pci_dev.map_bar(0, fmt='I')
self.large_bar = self.vram.nbytes >= self.vram_size
# UVM depth HW level VA bits
# 0 PDE4 56:56 (hopper+)
# 1 PDE3 55:47
# 2 PDE2 46:38
# 3 PDE1 (or 512M PTE) 37:29
# 4 PDE0 (dual 64k/4k PDE, or 2M PTE) 28:21
# 5 PTE_64K / PTE_4K 20:16 / 20:12
bits, shifts = (56, [12, 21, 29, 38, 47, 56]) if self.mmu_ver == 3 else (48, [12, 21, 29, 38, 47])
self.mm = NVMemoryManager(self, self.vram_size, boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=bits, va_shifts=shifts, va_base=0,
palloc_ranges=[(x, x) for x in [512 << 20, 2 << 20, 4 << 10]], reserve_ptable=not self.large_bar)
def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
view, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
view[:sz] = bytes(struct)
@@ -158,7 +162,7 @@ class NVDev(PCIDevImplBase):
regs_off = {'NV_PFALCON_FALCON': 0x0, 'NV_PGSP_FALCON': 0x0, 'NV_PSEC_FALCON': 0x0, 'NV_PRISCV_RISCV': 0x1000, 'NV_PGC6_AON': 0x0, 'NV_PFSP': 0x0,
'NV_PGC6_BSI': 0x0, 'NV_PFALCON_FBIF': 0x600, 'NV_PFALCON2_FALCON': 0x1000, 'NV_PBUS': 0x0, 'NV_PFB': 0x0, 'NV_PMC': 0x0, 'NV_PGSP_QUEUE': 0x0,
'NV_VIRTUAL_FUNCTION':0xb80000}
'NV_VIRTUAL_FUNCTION':0xb80000, "NV_THERM": 0x0}
for raw in self._download(file).splitlines():
if not raw.startswith("#define "): continue