mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 15:33:45 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
164 lines
8.5 KiB
Python
164 lines
8.5 KiB
Python
from __future__ import annotations
|
|
import time, functools, tinygrad.runtime.autogen.nv_regs
|
|
from tinygrad.helpers import getenv, DEBUG, getbits, round_up
|
|
from tinygrad.runtime.autogen import pci
|
|
from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
|
|
from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
|
|
from tinygrad.runtime.support.system import PCIDevice
|
|
from tinygrad.runtime.support.hcq import MMIOInterface
|
|
|
|
NV_DEBUG = getenv("NV_DEBUG", 0)
|
|
|
|
class NVReg:
|
|
def __init__(self, nvdev, base, off, fields=None): self.nvdev, self.base, self.off, self.fields = nvdev, base, off, fields
|
|
|
|
def __getitem__(self, idx:int): return NVReg(self.nvdev, self.base, self.off(idx), fields=self.fields)
|
|
|
|
def add_field(self, name:str, start:int, end:int): self.fields[name] = (start, end)
|
|
def with_base(self, base:int): return NVReg(self.nvdev, base + self.base, self.off, self.fields)
|
|
|
|
def read(self): return self.nvdev.rreg(self.base + self.off)
|
|
def read_bitfields(self) -> dict[str, int]: return self.decode(self.read())
|
|
|
|
def write(self, _ini_val:int=0, **kwargs): self.nvdev.wreg(self.base + self.off, _ini_val | self.encode(**kwargs))
|
|
|
|
def update(self, **kwargs): self.write(self.read() & ~self.mask(*kwargs.keys()), **kwargs)
|
|
|
|
def mask(self, *names):
|
|
return functools.reduce(int.__or__, ((((1 << (self.fields[nm][1]-self.fields[nm][0] + 1)) - 1) << self.fields[nm][0]) for nm in names), 0)
|
|
|
|
def encode(self, **kwargs) -> int: return functools.reduce(int.__or__, (value << self.fields[name][0] for name,value in kwargs.items()), 0)
|
|
def decode(self, val: int) -> dict: return {name:getbits(val, start, end) for name,(start,end) in self.fields.items()}
|
|
|
|
class NVPageTableEntry:
|
|
def __init__(self, nvdev, paddr, lv): self.nvdev, self.paddr, self.lv, self.entries = nvdev, paddr, lv, nvdev.vram.view(paddr, 0x1000, fmt='Q')
|
|
|
|
def _is_dual_pde(self) -> bool: return self.lv == self.nvdev.mm.level_cnt - 2
|
|
|
|
def set_entry(self, entry_id:int, paddr:int, table=False, uncached=False, aspace=AddrSpace.PHYS, snooped=False, frag=0, valid=True):
|
|
if not table:
|
|
x = self.nvdev.pte_t.encode(valid=valid, address_sys=paddr >> 12, aperture=2 if aspace is AddrSpace.SYS else 0, kind=6,
|
|
**({'pcf': int(uncached)} if self.nvdev.mmu_ver == 3 else {'vol': uncached}))
|
|
else:
|
|
pde = self.nvdev.dual_pde_t if self._is_dual_pde() else self.nvdev.pde_t
|
|
small, sys = ("_small" if self._is_dual_pde() else ""), "" if self.nvdev.mmu_ver == 3 else "_sys"
|
|
x = pde.encode(is_pte=False, **{f'aperture{small}': 1 if valid else 0, f'address{small}{sys}': paddr >> 12},
|
|
**({f'pcf{small}': 0b10} if self.nvdev.mmu_ver == 3 else {'no_ats': 1}))
|
|
|
|
if self._is_dual_pde(): self.entries[2*entry_id], self.entries[2*entry_id+1] = x & 0xffffffffffffffff, x >> 64
|
|
else: self.entries[entry_id] = x
|
|
|
|
def entry(self, entry_id:int) -> int:
|
|
return (self.entries[2*entry_id+1]<<64) | self.entries[2*entry_id] if self._is_dual_pde() else self.entries[entry_id]
|
|
|
|
def read_fields(self, entry_id:int) -> dict:
|
|
if self.is_page(entry_id): return self.nvdev.pte_t.decode(self.entry(entry_id))
|
|
return (self.nvdev.dual_pde_t if self._is_dual_pde() else self.nvdev.pde_t).decode(self.entry(entry_id))
|
|
|
|
def is_page(self, entry_id) -> bool: return (self.entry(entry_id) & 1 == 1) if self.lv < self.nvdev.mm.level_cnt - 1 else True
|
|
def supports_huge_page(self, paddr:int): return self.lv >= self.nvdev.mm.level_cnt - 3 and paddr % self.nvdev.mm.pte_covers[self.lv] == 0
|
|
|
|
def valid(self, entry_id):
|
|
if self.is_page(entry_id): return self.read_fields(entry_id)['valid']
|
|
return self.read_fields(entry_id)['aperture_small' if self._is_dual_pde() else 'aperture'] != 0
|
|
|
|
def address(self, entry_id:int) -> int:
|
|
small, sys = ("_small" if self._is_dual_pde() else ""), "_sys" if self.nvdev.mmu_ver == 2 or self.lv == self.nvdev.mm.level_cnt - 1 else ""
|
|
return self.read_fields(entry_id)[f'address{small}{sys}'] << 12
|
|
|
|
class NVMemoryManager(MemoryManager):
|
|
va_allocator = TLSFAllocator((1 << 44), base=0x1000000000)
|
|
|
|
def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))
|
|
|
|
class NVDev:
|
|
def __init__(self, pci_dev:PCIDevice):
|
|
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, self.is_err_state = False, True, False
|
|
self._early_ip_init()
|
|
self._early_mmu_init()
|
|
|
|
# No booting state, gsp client is reinited every run.
|
|
self.is_booting = False
|
|
|
|
for ip in [self.flcn, self.gsp]: ip.init_sw()
|
|
for ip in [self.flcn, self.gsp]: ip.init_hw()
|
|
|
|
def fini(self):
|
|
for ip in [self.gsp, self.flcn]: ip.fini_hw()
|
|
|
|
def reg(self, reg:str) -> NVReg: return self.__dict__[reg]
|
|
def wreg(self, addr:int, value:int):
|
|
self.mmio[addr // 4] = value
|
|
if NV_DEBUG >= 4: print(f"wreg: {hex(addr)} = {hex(value)}")
|
|
def rreg(self, addr:int) -> int: return self.mmio[addr // 4]
|
|
|
|
def _early_ip_init(self):
|
|
self.reg_names:set[str] = set()
|
|
self.reg_offsets:dict[str, tuple[int, int]] = {}
|
|
|
|
self.include("nv_ref", "")
|
|
self.include("dev_fb", "tu102")
|
|
self.include("dev_gc6_island", "ga102")
|
|
|
|
if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0:
|
|
self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
|
|
if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
|
|
self.pci_dev.reset()
|
|
time.sleep(0.1) # wait until device can respond again
|
|
|
|
self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
|
|
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.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.flcn.wait_for_reset()
|
|
|
|
def _early_mmu_init(self):
|
|
self.include("dev_vm", "tu102")
|
|
|
|
# MMU Init
|
|
self.include("dev_mmu", "gh100" if self.mmu_ver == 3 else "tu102")
|
|
self.pte_t, self.pde_t, self.dual_pde_t = [self.__dict__[name] for name in [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']]
|
|
|
|
self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20
|
|
|
|
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])
|
|
|
|
# tail vram reserved for falcon structs
|
|
self.mm = NVMemoryManager(self, self.vram_size - (64 << 20), 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_mem(self, size:int, data:bytes|None=None, contiguous:bool=False, sysmem:bool|None=None) -> tuple[MMIOInterface,int|None,list[int]]:
|
|
sz = round_up(size, 0x1000)
|
|
if sysmem is True or (sysmem is None and not self.large_bar):
|
|
view, sysaddr = self.pci_dev.alloc_sysmem(size, 0, contiguous=contiguous)
|
|
paddr = None
|
|
else:
|
|
paddr = self.mm.palloc(sz, boot=False)
|
|
view = self.vram.view(paddr, sz)
|
|
sysaddr = [self.pci_dev.bar_info(1)[0] + paddr + i * 0x1000 for i in range(sz // 0x1000)]
|
|
if data is not None: view[:size] = data
|
|
return view, paddr, sysaddr
|
|
|
|
def include(self, name:str, arch:str):
|
|
for k,v in getattr(getattr(tinygrad.runtime.autogen.nv_regs, name), arch or 'regs').items():
|
|
self.__dict__[k] = NVReg(self, *v) if isinstance(v, tuple) else v
|