diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index 5aefa06b62..0332bea173 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -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") diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index c004c43b20..132eebc44f 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -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