from __future__ import annotations import ctypes, mmap, struct, sys if sys.platform != "win32": from tinygrad.runtime.autogen import libc class MockUSB: def __init__(self, mem): self.mem = mem def read(self, address, size): return bytes(self.mem[address:address+size]) def write(self, address, data): self.mem[address:address+len(data)] = data def pcie_mem_read(self, address, nbytes): return bytes(self.mem[address:address+nbytes]) def pcie_mem_write(self, address, data): self.mem[address:address+len(data)] = data # *** ASM24 Controller Mock *** _mock_usb_state: MockASM24State|None = None class MockASM24State: """Mock custom ASM24 controller: XRAM, DMA windows, PCI config space, and GPU BARs. Memory map (64KB XRAM): 0xA000-0xAFFF: DMA window -> sys 0x820000 0xB000-0xB1FF: DMA window -> sys 0x800000 0xB200-0xB7FF: controller PCI MMIO 0xF000-0xFFFF: DMA window -> sys 0x200000 (512KB) """ XRAM_SIZE = 0x10000 def __init__(self, gpu, driver, vram_size:int, doorbell_size:int, mmio_size:int): self.gpu, self.driver = gpu, driver self._xram = bytearray(self.XRAM_SIZE) self._doorbell_addr = libc.mmap(0, doorbell_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, gpu.doorbell_fd, 0) self._doorbell = (ctypes.c_ubyte * doorbell_size).from_address(self._doorbell_addr) # DMA windows: ctrl_addr -> (host_addr, size) self._dma_regions: dict[int, tuple[int, int]] = {} self._add_dma_window(0xF000, 0x200000, 0x80000) self._add_dma_window(0xA000, 0x820000, 0x1000) self._add_dma_window(0xB000, 0x800000, 0x200) # PCI config space: (bus,dev,fn) -> bytearray(4096) self._pci_cfg: dict[tuple[int,int,int], bytearray] = {} # GPU BAR definitions: reg_offset -> (size, type_bits, is_64bit) self._gpu_bars: dict[int, tuple[int, int, bool]] = { 0x10: (vram_size, 0x0C, True), # BAR0: VRAM, 64-bit prefetchable 0x18: (doorbell_size, 0x00, False), # BAR2: doorbell, 32-bit 0x1C: (0, 0x00, False), # BAR3: unused 0x20: (0, 0x00, False), # BAR4: unused 0x24: (mmio_size, 0x00, False), # BAR5: MMIO, 32-bit } self._bar_addrs: dict[int, tuple[int, int]] = {} # reg_offset -> (addr, size) # Initialize GPU config space (bus=4, dev=0, fn=0) with BAR type bits and REBAR capability gpu_cfg = self._get_cfg(4, 0, 0) for reg_off, (sz, type_bits, _) in self._gpu_bars.items(): if sz > 0: struct.pack_into(' bytearray: if (key:=(bus, dev, fn)) not in self._pci_cfg: self._pci_cfg[key] = bytearray(4096) return self._pci_cfg[key] def _add_dma_window(self, ctrl_addr:int, sys_addr:int, size:int): host_addr = libc.mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, -1, 0) self._dma_regions[ctrl_addr] = (host_addr, size) for off in range(0, size, 0x1000): self.gpu._sysmem_map[sys_addr + off] = host_addr + off # --- XRAM access --- def _xram_read(self, addr:int, length:int) -> bytes: for ctrl_addr, (host_addr, dma_size) in self._dma_regions.items(): if ctrl_addr <= addr < ctrl_addr + dma_size: return bytes((ctypes.c_ubyte * length).from_address(host_addr + (addr - ctrl_addr))) return bytes(self._xram[addr:addr+length]) def _xram_write_byte(self, addr:int, value:int): for ctrl_addr, (host_addr, dma_size) in self._dma_regions.items(): if ctrl_addr <= addr < ctrl_addr + dma_size: (ctypes.c_ubyte * 1).from_address(host_addr + (addr - ctrl_addr))[0] = value return self._xram[addr] = value def _cfg_write(self, bus:int, dev:int, fn:int, byte_addr:int, val:int, size:int): cfg = self._get_cfg(bus, dev, fn) # Handle BAR register writes for GPU device (bus=4, dev=0, fn=0) if (bus, dev, fn) == (4, 0, 0) and 0x10 <= byte_addr < 0x28 and size == 4: reg_off = byte_addr & ~0x3 if (bar_def:=self._gpu_bars.get(reg_off)) is not None: bar_size, type_bits, is_64 = bar_def if bar_size == 0: return # unused BAR if val == 0xFFFFFFFF: # size probe struct.pack_into('> (8 * i)) & 0xFF def _find_bar(self, address:int, size:int) -> tuple[int, int]: for reg_off, (bar_addr, bar_size) in self._bar_addrs.items(): if bar_addr <= address and address + size <= bar_addr + bar_size: return reg_off, address - bar_addr raise ValueError(f"PCIe range {address:#x}+{size:#x} not mapped to any BAR") def _pcie_read(self, address:int, size:int) -> bytes: reg_off, offset = self._find_bar(address, size) if reg_off == 0x10: return bytes(self.gpu.vram[offset:offset+size]) if reg_off == 0x18: return bytes(self._doorbell[offset:offset+size]) if reg_off == 0x24: return bytes((self.gpu.mmio[(offset+i)//4] >> (8*((offset+i)&3))) & 0xFF for i in range(size)) raise RuntimeError(f"unsupported BAR register {reg_off:#x}") def _pcie_write(self, address:int, data:bytes): reg_off, offset = self._find_bar(address, len(data)) if reg_off == 0x10: self.gpu.vram[offset:offset+len(data)] = list(data) elif reg_off == 0x18: self._doorbell[offset:offset+len(data)] = list(data) self.driver._emulate_execute() elif reg_off == 0x24: updates: dict[int, int] = {} for i, byte in enumerate(data): idx, shift = (offset+i)//4, 8*((offset+i)&3) updates[idx] = (updates.get(idx, self.gpu.mmio[idx]) & ~(0xFF << shift)) | (byte << shift) for idx, val in updates.items(): self.gpu.mmio[idx] = val else: raise RuntimeError(f"unsupported BAR register {reg_off:#x}") def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None: if value is None: return int.from_bytes(self._pcie_read(address, size), 'little') self._pcie_write(address, value.to_bytes(size, 'little')) return None class MockUSB3: @classmethod def list_devices(cls, vendor, dev): return [(0, "usb:mock")] def __init__(self, *args, **kwargs): self.product = "custom mock" self._bulk_read_op: tuple[str, int, int]|None = None self._bulk_write_op: tuple[str, int, int]|None = None self._f0_reply = bytes(8) @property def state(self) -> MockASM24State: assert _mock_usb_state is not None return _mock_usb_state def control_write(self, request:int, value:int=0, index:int=0, data:bytes=b'', timeout:int=1000): if request == 0xF3: self.state._xram[0xB450] = 0x78 if value else 0 elif request == 0xE5: self.state._xram_write_byte(value, index) elif request == 0xF2: op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000 + (index & 0xFF) * 0x4000, (value & 0x7FFF) * 512) if value & 0x8000: self._bulk_read_op = op else: self._bulk_write_op = op elif request == 0xF0: address_lo, address_hi, payload = struct.unpack('> 8 if index == 1: self._bulk_write_op = ("pcie_write", address, payload * 4) elif index == 2: self._bulk_read_op = ("pcie_read", address, payload * 4) else: assert index == 0 and byte_en offset = (byte_en & -byte_en).bit_length() - 1 size, is_write, is_cfg = byte_en.bit_count(), bool(fmt_type & 0x40), (fmt_type & 0xBE) == 0x04 if is_cfg: bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC if is_write: self.state._cfg_write(bus, dev, fn, byte_addr + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size) else: payload = int.from_bytes(self.state._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little') elif is_write: self.state._pcie_dispatch(address + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size) else: payload = (self.state._pcie_dispatch(address + offset, None, size) or 0) << (8 * offset) self._f0_reply = struct.pack(' memoryview: if request == 0xE4: data = self.state._xram_read(value, length) elif request == 0xF0: data = self._f0_reply else: raise ValueError(f"unsupported control IN request 0x{request:02X}") return memoryview(data[:length]) def bulk_write(self, data:bytes, timeout:int=1000): assert self._bulk_write_op is not None op, address, size = self._bulk_write_op assert len(data) == size if op == "sram_write": ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1]) ctypes.memmove(host_addr + (address - ctrl), data, min(len(data), region_size - (address - ctrl))) self.state.driver._emulate_execute() # landed data may un-stall a ring polling on it (e.g. copyin sentinels) elif op == "pcie_write": self.state._pcie_write(address, data) else: raise RuntimeError(f"cannot bulk write for {op}") self._bulk_write_op = None def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int: # the mock completes transfers synchronously self.bulk_write(bytes(payload), timeout) return 0 def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int: self.control_write(request, value, index, data, timeout) return 0 def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]: return 0, self.control_read(request, length, value, index, timeout) def bulk_wait(self, tag:int): pass def bulk_read(self, length:int, timeout:int=1000) -> memoryview: assert self._bulk_read_op is not None op, address, size = self._bulk_read_op assert length == size if op == "sram_read": ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1]) data = bytes((ctypes.c_ubyte * min(length, region_size - (address - ctrl))).from_address(host_addr + (address - ctrl))) elif op == "pcie_read": data = self.state._pcie_read(address, length) else: raise RuntimeError(f"cannot bulk read for {op}") self._bulk_read_op = None return memoryview(data)