tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 11:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

nv: support all gb20x and small bar (#12721)

Browse Source
This commit is contained in:
nimlgen
2025-10-16 15:37:54 +08:00
committed by GitHub
parent 1d1e1d9d88
commit b8cd66c7a2
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/support/nv/nvdev.py
+5 -3
View File
@@ -87,7 +87,7 @@ class NVDev(PCIDevImplBase):
# 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]])
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)
@@ -114,6 +114,7 @@ class NVDev(PCIDevImplBase):
self.chip_id = self.reg("NV_PMC_BOOT_0").read()
self.chip_details = self.reg("NV_PMC_BOOT_42").read_bitfields()
self.chip_name = {0x17: "GA1", 0x19: "AD1", 0x1b: "GB2"}[self.chip_details['architecture']] + f"{self.chip_details['implementation']:02d}"
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.include("src/common/inc/swref/published/turing/tu102/dev_fb.h")
@@ -133,6 +134,7 @@ class NVDev(PCIDevImplBase):
self.pte_t, self.pde_t, self.dual_pde_t = tuple([self.__dict__[name] for name in mmu_pd_names])
self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20
self.large_bar = self.vram.nbytes >= self.vram_size
def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
@@ -146,8 +148,8 @@ class NVDev(PCIDevImplBase):
def extract_fw(self, file:str, dname:str) -> bytes:
# Extracts the firmware binary from the given header
tname = file.replace("kgsp", "kgspGet")
text = self._download(f"src/nvidia/generated/g_bindata_{tname}_{self.chip_name}.c")
info, sl = text[text[:text.index(dnm:=f'{file}_{self.chip_name}_{dname}')].rindex("COMPRESSION:"):][:16], text[text.index(dnm) + len(dnm) + 7:]
text = self._download(f"src/nvidia/generated/g_bindata_{tname}_{self.fw_name}.c")
info, sl = text[text[:text.index(dnm:=f'{file}_{self.fw_name}_{dname}')].rindex("COMPRESSION:"):][:16], text[text.index(dnm) + len(dnm) + 7:]
image = bytes.fromhex(sl[:sl.find("};")].strip().replace("0x", "").replace(",", "").replace(" ", "").replace("\n", ""))
return gzip.decompress(struct.pack("<4BL2B", 0x1f, 0x8b, 8, 0, 0, 0, 3) + image) if "COMPRESSION: YES" in info else image