diff --git a/extra/amdpci/hive_reset.py b/extra/amdpci/hive_reset.py index 4e72550aa2..6b2f515bb1 100755 --- a/extra/amdpci/hive_reset.py +++ b/extra/amdpci/hive_reset.py @@ -12,7 +12,7 @@ if __name__ == "__main__": drv_path = f"/sys/bus/pci/devices/{gpu}/driver" if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu": raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...") - pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus] + pcidevs = [PCIDevice("AM", gpu) for gpu in gpus] amdevs = [] with Context(DEBUG=2): for pcidev in pcidevs: diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 54f9002066..40f758b51a 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -825,7 +825,7 @@ class KFDIface: class PCIIface(PCIIfaceBase): def __init__(self, dev, dev_id): - super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x74a1, 0x744c, 0x7480, 0x7550, 0x7590, 0x75a0])], bars=[0, 2, 5], vram_bar=0, + super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x74a1, 0x744c, 0x7480, 0x7550, 0x7590, 0x75a0])], vram_bar=0, va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev) self._compute_props() @@ -892,10 +892,10 @@ class PCIIface(PCIIfaceBase): class USBIface(PCIIface): def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called - self.dev, self.pci_dev = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}", bars=[0, 2, 5]) + self.dev, self.pci_dev = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}") self.dev_impl = AMDev(self.pci_dev) self._compute_props() - self.pci_dev.usb._pci_cacheable += [(self.pci_dev.bar_info[2].addr, self.pci_dev.bar_info[2].size)] # doorbell region is cacheable + self.pci_dev.usb._pci_cacheable += [self.pci_dev.bar_info(2)] # doorbell region is cacheable # special regions self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)] diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index e34704e1f5..1d84948c89 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -540,8 +540,7 @@ class PCIIface(PCIIfaceBase): if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)") if not OSX: System.reserve_hugepages(64) super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])], - base_class=0x03, bars=[0, 1, 3], vram_bar=1, dev_impl_t=NVDev, - va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size) + base_class=0x03, vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size, dev_impl_t=NVDev) self.root, self.gpu_instance = 0xc1000000, 0 self.rm_alloc(0, nv_gpu.NV01_ROOT, nv_gpu.NV0000_ALLOC_PARAMETERS()) diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index 84e77baf95..9441fad1fd 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -575,7 +575,7 @@ class NV_GSP(NV_IP): def bdf_as_int(s): return 0x000 if s.startswith("usb") or s.startswith("remote") else (int(s[5:7],16)<<8) | (int(s[8:10],16)<<3) | int(s[-1],16) pcidev = self.nvdev.pci_dev - data = nv.GspSystemInfo(gpuPhysAddr=pcidev.bar_info[0].addr, gpuPhysFbAddr=pcidev.bar_info[1].addr, gpuPhysInstAddr=pcidev.bar_info[3].addr, + data = nv.GspSystemInfo(gpuPhysAddr=pcidev.bar_info(0)[0], gpuPhysFbAddr=pcidev.bar_info(1)[0], gpuPhysInstAddr=pcidev.bar_info(3)[0], pciConfigMirrorBase=[0x88000, 0x92000][self.nvdev.fmc_boot], pciConfigMirrorSize=0x1000, nvDomainBusDeviceFunc=bdf_as_int(self.nvdev.devfmt), bIsPassthru=1, PCIDeviceID=pcidev.read_config(pci.PCI_VENDOR_ID, 4), PCISubDeviceID=pcidev.read_config(pci.PCI_SUBSYSTEM_VENDOR_ID, 4), PCIRevisionID=pcidev.read_config(pci.PCI_REVISION_ID, 1), maxUserVa=0x7ffffffff000) diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index af7b3eb1ad..d9fbe11c30 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -9,9 +9,6 @@ from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface MAP_FIXED, MAP_FIXED_NOREPLACE = 0x10, 0x100000 MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400 -@dataclasses.dataclass(frozen=True) -class PCIBarInfo: addr:int; size:int # noqa: E702 - class _System: def write_sysfs(self, path:str, value:str, msg:str, expected:str|None=None): if FileIOInterface(path, os.O_RDONLY).read().splitlines()[0] != (expected or value): @@ -76,14 +73,13 @@ class _System: return sorted([val for vndr, device, val in all_devs if vndr == vendor and any((device & mask) in devlist for mask, devlist in devices)]) - def pci_probe_device(self, devpref:str, dev_id:int, vendor:int, devices:list[tuple[int, list[int]]], bars:list[int], - resize_bars:list[int]|None=None, base_class:int|None=None): + def pci_probe_device(self, devpref:str, dev_id:int, vendor:int, devices:list[tuple[int, list[int]]], base_class:int|None=None): gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices, base_class)) if not gpus: raise RuntimeError("No supported GPUs found") - if OSX: return APLRemotePCIDevice(devpref, f'usb4:{dev_id}', bars) - return PCIDevice(devpref, gpus[dev_id], bars=bars, resize_bars=resize_bars) + if OSX: return APLRemotePCIDevice(devpref, f'usb4:{dev_id}') + return PCIDevice(devpref, gpus[dev_id]) - def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, PCIBarInfo]: + def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]: for bus in range(gpu_bus): # All 3 values must be written at the same time. buses = (0 << 0) | ((bus+1) << 8) | ((gpu_bus) << 16) @@ -125,7 +121,7 @@ class _System: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] & 0xffffffff, size=4) if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] >> 32, size=4) - bars[bar_off // 4] = PCIBarInfo(mem_space_addr[bar_mem], bar_size) + bars[bar_off // 4] = (mem_space_addr[bar_mem], bar_size) mem_space_addr[bar_mem] += round_up(bar_size, 2 << 20) bar_off += 8 if bar_64 else 4 @@ -152,7 +148,7 @@ System = _System() # *** PCI Devices class PCIDevice: - def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None): + def __init__(self, devpref:str, pcibus:str): self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock") self.pcibus, self.irq_poller = pcibus, None @@ -162,11 +158,6 @@ class PCIDevice: if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"): FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver/unbind", os.O_WRONLY).write(self.pcibus) - for i in resize_bars or []: - if FileIOInterface.exists(rpath:=f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize"): - try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1)) - except OSError as e: raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled.") from e - if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci") FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus) @@ -188,10 +179,6 @@ class PCIDevice: else: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/enable", os.O_RDWR).write("1") self.cfg_fd = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/config", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC) - self.bar_fds = {b: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{b}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC) for b in bars} - - res = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines() - self.bar_info = {j:PCIBarInfo(int(s,16), int(e,16)-int(s,16)+1) for j,(s,e,_) in enumerate(l.split() for l in res)} def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]: assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB" @@ -199,42 +186,60 @@ class PCIDevice: va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0) sysmem_view, paddrs = MMIOInterface(va, size), [(x, mmap.PAGESIZE) for x in System.system_paddrs(va, size)] return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)] + + def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'") def read_config(self, offset:int, size:int): return int.from_bytes(self.cfg_fd.read(size, binary=True, offset=offset), byteorder='little') def write_config(self, offset:int, value:int, size:int): self.cfg_fd.write(value.to_bytes(size, byteorder='little'), binary=True, offset=offset) + + @functools.cache + def bar_fd(self, bar_idx:int) -> FileIOInterface: + return FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC) + @functools.cache + def bar_info(self, bar_idx:int) -> tuple[int, int]: + s, e, _ = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()[bar_idx].split() + return (int(s, 16), int(e, 16) - int(s, 16) + 1) def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: - fd, sz = self.bar_fds[bar], size or (self.bar_info[bar].size - off) + fd, sz = self.bar_fd(bar), size or (self.bar_info(bar)[1] - off) libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz, libc.MADV_DONTFORK) return MMIOInterface(loc, sz, fmt=fmt) - def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'") + def resize_bar(self, bar_idx:int): + rpath = f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}_resize" + try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1)) + except OSError as e: raise RuntimeError(f"Cannot resize BAR {bar_idx}: {e}. Ensure the resizable BAR option is enabled.") from e class USBPCIDevice(PCIDevice): - def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None): + def __init__(self, devpref:str, pcibus:str): self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock") self.usb = ASM24Controller() - self.pcibus, self.bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30)) + self.pcibus, self._bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30)) self.sram = BumpAllocator(size=0x80000, wrap=False) # asm24 controller sram - def read_config(self, offset:int, size:int): return self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, size=size) - def write_config(self, offset:int, value:int, size:int): self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, value=value, size=size) - def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'): - return USBMMIOInterface(self.usb, self.bar_info[bar].addr + off, size or self.bar_info[bar].size, fmt) + def dma_view(self, ctrl_addr, size): return USBMMIOInterface(self.usb, ctrl_addr, size, fmt='B', pcimem=False) def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]: return self.dma_view(0xf000 + (off:=self.sram.alloc(size)), size), [0x200000 + off] + def read_config(self, offset:int, size:int): return self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, size=size) + def write_config(self, offset:int, value:int, size:int): self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, value=value, size=size) + + def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._bar_info[bar_idx] # type: ignore[override] + def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'): + return USBMMIOInterface(self.usb, self.bar_info(bar)[0] + off, size or self.bar_info(bar)[1], fmt) + def resize_bar(self, bar_idx:int): pass # already resized + @dataclasses.dataclass class PCIAllocationMeta: mapping:VirtMapping; has_cpu_mapping:bool; hMemory:int=0 # noqa: E702 class PCIIfaceBase: def is_local(self) -> bool: return not isinstance(self.pci_dev, RemotePCIDevice) - def is_bar_small(self) -> bool: return self.pci_dev.bar_info[self.vram_bar].size == (256 << 20) + def is_bar_small(self) -> bool: return self.pci_dev.bar_info(self.vram_bar)[1] == (256 << 20) - 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]]], vram_bar, va_start, va_size, dev_impl_t, base_class:int|None=None): - self.pci_dev = System.pci_probe_device(dev.__class__.__name__[:2], dev_id, vendor, devices, bars, resize_bars=[vram_bar], base_class=base_class) + self.pci_dev = System.pci_probe_device(dev.__class__.__name__[:2], dev_id, vendor, devices, base_class=base_class) if self.is_local(): System.reserve_va(va_start, va_size) + with contextlib.suppress(Exception): self.pci_dev.resize_bar(vram_bar) self.dev_impl = dev_impl_t(self.pci_dev) self.dev, self.vram_bar = dev, vram_bar - self.p2p_base_addr = self.pci_dev.bar_info[vram_bar].addr def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer: should_use_sysmem = host or ((cpu_access if self.is_bar_small() else (uncached and cpu_access)) and not force_devmem) @@ -254,7 +259,7 @@ class PCIIfaceBase: if self.is_local() and b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size) def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]: - return [(p + self.p2p_base_addr, sz) for p, sz in paddrs], AddrSpace.SYS + return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS def map(self, b:HCQBuffer): if not self.is_local(): raise RuntimeError(f"P2P mapping not supported for remote devices: {b.owner} -> {self.dev}") @@ -295,11 +300,10 @@ class RemoteMMIOInterface(MMIOInterface): return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset) class RemotePCIDevice(PCIDevice): - def __init__(self, devpref:str, pcibus:str, bars:list[int], sock:socket.socket): + def __init__(self, devpref:str, pcibus:str, sock:socket.socket): self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock") self.pcibus, self.sock = pcibus, sock for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20) - self.bar_info = {b: PCIBarInfo(0, self._rpc(RemoteCmd.MAP_BAR, b)[0]) for b in bars} def _recvall(self, n:int) -> bytes: data = b'' @@ -328,11 +332,16 @@ class RemotePCIDevice(PCIDevice): # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping. paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2]))) return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)] + + def reset(self): self._rpc(RemoteCmd.RESET, 0, 0, 0) def read_config(self, offset:int, size:int): return self._rpc(RemoteCmd.CFG_READ, 0, offset, size)[0] def write_config(self, offset:int, value:int, size:int): self._rpc(RemoteCmd.CFG_WRITE, 0, offset, size, value) - def reset(self): self._rpc(RemoteCmd.RESET, 0, 0, 0) + + @functools.cache + def bar_info(self, bar_idx:int) -> tuple[int, int]: return (0, self._rpc(RemoteCmd.MAP_BAR, bar_idx)[0]) def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface: - return RemoteMMIOInterface(self, bar, size or self.bar_info[bar].size, fmt).view(off, size, fmt) + return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt) + def resize_bar(self, bar_idx:int): pass # TODO: resizing not supported for remote devices class APLRemotePCIDevice(RemotePCIDevice): APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU" @@ -345,7 +354,7 @@ class APLRemotePCIDevice(RemotePCIDevice): system(f"ditto -xk {fetch(f'https://github.com/nimlgen/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications") print(system(f"{cls.APP_PATH} install")) - def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None): + def __init__(self, devpref:str, pcibus:str): self.ensure_app() sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) for i in range(100): @@ -355,4 +364,4 @@ class APLRemotePCIDevice(RemotePCIDevice): if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) time.sleep(0.05) else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.") - super().__init__(devpref, pcibus, bars, sock) + super().__init__(devpref, pcibus, sock)