github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-07 22:23:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

system: add base_class to pci_scan_bus (#13845)

Browse Source 
* system: add base_class to pci_scan_bus

* fix
This commit is contained in:
nimlgen
2025-12-27 13:22:21 +03:00
committed by GitHub
parent fac137779e
commit 276159cb87
2 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/runtime/ops_nv.py
View File

@@ -508,7 +508,7 @@ class PCIIface(PCIIfaceBase):
# PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used.
if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)")
super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
base_class=0x03, bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
if not OSX: System.reserve_hugepages(64)
self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)

9
tinygrad/runtime/support/system.py
View File

@@ -93,11 +93,12 @@ class _System:
if data is not None: sysmem_view[:len(data)] = data
return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
def pci_scan_bus(self, target_vendor:int, target_devices:list[tuple[int, list[int]]]) -> list[str]:
def pci_scan_bus(self, target_vendor:int, target_devices:list[tuple[int, list[int]]], base_class:int|None=None) -> list[str]:
result = []
for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
vendor = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/vendor").read(), 16)
device = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/device").read(), 16)
if base_class is not None and int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/class").read(), 16) >> 16 != base_class: continue
if vendor == target_vendor and any((device & mask) in devlist for mask, devlist in target_devices): result.append(pcibus)
return sorted(result)
@@ -247,9 +248,9 @@ class LNXPCIIfaceBase:
dev_impl:PCIDevImplBase
gpus:ClassVar[list[str]] = []
def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size):
def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size, base_class:int|None=None):
if len((cls:=type(self)).gpus) == 0:
cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices))
cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices, base_class))
# Acquire va range to avoid collisions.
FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED, 0)
@@ -290,7 +291,7 @@ class LNXPCIIfaceBase:
self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
class APLPCIIfaceBase(LNXPCIIfaceBase):
def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size, base_class:int|None=None):
self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(dev.__class__.__name__[:2], pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
assert (read_vendor:=self.pci_dev.read_config(0x00, 2)) == vendor, f"Vendor ID mismatch: expected {vendor:#x}, got {read_vendor:#x}"
def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")