superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ f2a861c

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-02 15:07:09 -05:00
parent b42569dbca
commit e8748fd704
5497 changed files with 316070 additions and 179848 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
artifacts/package_sources/tinygrad/test/mockgpu/am/__init__.py Normal file
View File

127
artifacts/package_sources/tinygrad/test/mockgpu/am/amdriver.py Normal file
View File

@@ -0,0 +1,127 @@
from __future__ import annotations
import mmap, functools
from tinygrad.runtime.autogen import libc
from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
from test.mockgpu.am.amgpu import MockAMGPU, VRAM_SIZE
DOORBELL_SIZE = 0x2000
MMIO_SIZE = 2 << 20
PCIBUS = "mock:am:0"
_empty_bar = "0x0000000000000000 0x0000000000000000 0x0000000000000000"
_resource_lines = [
f"0x0000000000000000 0x{VRAM_SIZE-1:016x} 0x0000000000000000", _empty_bar,
f"0x0000000000000000 0x{DOORBELL_SIZE-1:016x} 0x0000000000000000", _empty_bar, _empty_bar,
f"0x0000000000000000 0x{MMIO_SIZE-1:016x} 0x0000000000000000", _empty_bar,
]
class PagemapFileDesc(VirtFileDesc):
def __init__(self, fd, gpu):
super().__init__(fd)
self.gpu = gpu
def seek(self, offset): self.off = offset
def read_contents(self, size=None):
entries = bytearray()
for i in range((size or 8) // 8):
vaddr = ((self.off // 8) + i) * 0x1000
paddr = self.gpu._next_sysmem_paddr
self.gpu._next_sysmem_paddr += 0x1000
self.gpu._sysmem_map[paddr] = vaddr
entries += ((1 << 63) | (paddr // 0x1000)).to_bytes(8, 'little')
self.off += len(entries)
return bytes(entries)
class PCIBarFileDesc(VirtFileDesc):
def __init__(self, fd, memfd, driver=None):
super().__init__(fd)
self.memfd, self.driver = memfd, driver
def mmap(self, start, sz, prot, flags, fd, off):
addr = libc.mmap(start, sz, prot, flags, self.memfd, off)
if self.driver is not None:
self.driver.track_address(addr, addr + sz, lambda mv, idx: None, lambda mv, idx: self.driver._emulate_execute())
return addr
class PCIMMIOBarFileDesc(VirtFileDesc):
def __init__(self, fd, bar5_addr):
super().__init__(fd)
self.bar5_addr = bar5_addr
def mmap(self, start, sz, prot, flags, fd, off): return self.bar5_addr + off
class PCIConfigFileDesc(VirtFileDesc):
def __init__(self, fd):
super().__init__(fd)
self.data = bytearray(256)
def read_contents(self, size=None): return bytes(self.data[self.off:self.off + (size or len(self.data) - self.off)])
def write_contents(self, content): self.data[self.off:self.off + len(content)] = content
def seek(self, offset): self.off = offset
class PCIEnableFileDesc(VirtFileDesc):
def __init__(self, fd): super().__init__(fd)
def read_contents(self, size=None): return "1\n"
def write_contents(self, content): pass
class AMDriver(VirtDriver):
def __init__(self):
super().__init__()
self.gpus:dict[int, MockAMGPU] = {}
self._executing = False
self.gpu = MockAMGPU(0)
self.gpus[0] = self.gpu
self.next_fd = 1 << 30
self._bar5_addr = libc.mmap(0, MMIO_SIZE, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, -1, 0)
mmio = self.gpu.mmio
self.track_address(self._bar5_addr, self._bar5_addr + MMIO_SIZE,
lambda mv, idx: _bar5_sync_read(mv, idx, mmio), lambda mv, idx: _bar5_sync_write(mv, idx, mmio))
p = f"/sys/bus/pci/devices/{PCIBUS}"
self.tracked_files += [
VirtFile("/proc/sys/vm/compact_unevictable_allowed", functools.partial(TextFileDesc, text="0\n")),
VirtFile("/proc/self/pagemap", functools.partial(PagemapFileDesc, gpu=self.gpu)),
VirtFile("/sys/bus/pci/devices", functools.partial(DirFileDesc, child_names=[PCIBUS])),
VirtFile(f"{p}/vendor", functools.partial(TextFileDesc, text="0x1002\n")),
VirtFile(f"{p}/device", functools.partial(TextFileDesc, text="0x74a1\n")),
VirtFile(f"{p}/enable", PCIEnableFileDesc),
VirtFile(f"{p}/config", PCIConfigFileDesc),
VirtFile(f"{p}/resource", functools.partial(TextFileDesc, text="\n".join(_resource_lines) + "\n")),
VirtFile(f"{p}/resource0", functools.partial(PCIBarFileDesc, memfd=self.gpu.vram_fd)),
VirtFile(f"{p}/resource2", functools.partial(PCIBarFileDesc, memfd=self.gpu.doorbell_fd, driver=self)),
VirtFile(f"{p}/resource5", functools.partial(PCIMMIOBarFileDesc, bar5_addr=self._bar5_addr)),
]
def _alloc_fd(self):
fd = self.next_fd
self.next_fd += 1
return fd
def open(self, name, flags, mode, virtfile): return virtfile.fdcls(self._alloc_fd())
def _emulate_execute(self):
if self._executing: return
self._executing = True
try:
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.executing: any_progress |= q.execute() > 0
finally:
self._executing = False
def _bar5_sync_read(mv, idx, mmio):
if isinstance(idx, slice):
for i in range(idx.start or 0, idx.stop or len(mv), idx.step or 1): mv[i] = mmio[i]
else: mv[idx] = mmio[idx]
def _bar5_sync_write(mv, idx, mmio):
if isinstance(idx, slice):
for i in range(idx.start or 0, idx.stop or len(mv), idx.step or 1): mmio[i] = mv[i]
else: mmio[idx] = mv[idx]
class AMUSBDriver(AMDriver):
def __init__(self):
import test.mockgpu.usb as _musb
super().__init__()
self.state = _musb.MockASM24State(self.gpu, self, VRAM_SIZE, DOORBELL_SIZE, MMIO_SIZE)
_musb._mock_usb_state = self.state

310
artifacts/package_sources/tinygrad/test/mockgpu/am/amgpu.py Normal file
View File

@@ -0,0 +1,310 @@
# mypy: ignore-errors
from __future__ import annotations
import ctypes, struct, functools, os, mmap
from tinygrad.runtime.autogen.am import am
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.amd import AMDReg, import_asic_regs
from test.mockgpu.amd.amdgpu import AMDGPU
VRAM_SIZE = 512 << 20
IP_VERSIONS = {
am.GC_HWIP: (12, 0, 0), am.SDMA0_HWIP: (7, 0, 0), am.MMHUB_HWIP: (4, 1, 0), am.NBIO_HWIP: (6, 3, 1),
am.MP0_HWIP: (14, 0, 2), am.MP1_HWIP: (14, 0, 2), am.HDP_HWIP: (7, 0, 0), am.OSSSYS_HWIP: (7, 0, 0),
}
def _pad(t, n=10): return t + (0,) * (n - len(t))
IP_BASES = {
am.GC_HWIP: _pad((0x00001260, 0x0000A000, 0x0001C000, 0x02402C00)),
am.SDMA0_HWIP: _pad((0x00001260, 0x0000A000, 0x0001C000, 0x02402C00)),
am.MMHUB_HWIP: _pad((0x0001A000, 0x02408800)),
am.NBIO_HWIP: _pad((0x00000000, 0x00000014, 0x00000D20, 0x00010400, 0x0241B000, 0x04040000)),
am.MP0_HWIP: _pad((0x00016000, 0x00DC0000, 0x00E00000, 0x00E40000, 0x0243FC00)),
am.MP1_HWIP: _pad((0x00016000, 0x00DC0000, 0x00E00000, 0x00E40000, 0x0243FC00)),
am.HDP_HWIP: _pad((0x00000F20, 0x0240A400)),
am.OSSSYS_HWIP: _pad((0x000010A0, 0x0240A000)),
}
IP_HWIDS = {hwip: am.hw_id_map[hwip] for hwip in IP_VERSIONS}
GC_INFO = dict(gc_num_se=2, gc_num_cu_per_sh=8, gc_num_sh_per_se=2, gc_num_rb_per_se=4,
gc_num_tccs=8, gc_wave_size=32, gc_max_waves_per_simd=16, gc_max_scratch_slots_per_cu=32, gc_lds_size=64)
def _build_ip_regs(prefix, hwip) -> dict[str, AMDReg]:
try: return import_asic_regs(prefix, IP_VERSIONS[hwip], cls=functools.partial(AMDReg, bases={0: IP_BASES[hwip]}))
except Exception: return {}
class MockMMU:
def __init__(self, gpu:MockAMGPU):
self.gpu = gpu
self.tlb: dict[int, tuple[int, int, bool]] = {}
def invalidate(self, pt_base:int, va_base:int):
new_tlb: dict[int, tuple[int, int, bool]] = {}
self._walk(pt_base, 0, 0, new_tlb, va_base)
for va, (pa, sz, is_sys) in new_tlb.items():
old = self.tlb.get(va)
if not is_sys and (old is None or old[0] != pa): self.gpu.map_vram_at(va, pa, sz)
if old is None: self.gpu.map_range(va, sz)
self.tlb = new_tlb
def _walk(self, pt_paddr:int, level:int, va_acc:int, out:dict, va_base:int):
shift = [39, 30, 21, 12][level]
for i in range(512):
pte = struct.unpack_from('<Q', self.gpu.vram, pt_paddr + i * 8)[0]
if not (pte & am.AMDGPU_PTE_VALID): continue
va, pa = va_acc | (i << shift), pte & 0x0000FFFFFFFFF000
if level == 3 or (pte & am.AMDGPU_PDE_PTE_GFX12):
out[va_base + va] = (pa, 1 << shift, bool(pte & am.AMDGPU_PTE_SYSTEM))
else:
self._walk(pa, level + 1, va, out, va_base)
def paddr_to_host(self, paddr:int) -> int:
page, off = paddr & ~0xFFF, paddr & 0xFFF
if page in self.gpu._sysmem_map: return self.gpu._sysmem_map[page] + off
if paddr < VRAM_SIZE: return self.gpu.vram_addr + paddr
raise ValueError(f"paddr {paddr:#x} not found in sysmem_map or VRAM")
def addr_to_host(self, addr:int) -> int:
gmc = self.gpu.mmio.gmc
sys_lo = self.gpu.mmio.regs.get(gmc.reg('regMMMC_VM_SYSTEM_APERTURE_LOW_ADDR') or 0, 0) << 18
sys_hi = self.gpu.mmio.regs.get(gmc.reg('regMMMC_VM_SYSTEM_APERTURE_HIGH_ADDR') or 0, 0) << 18
if sys_lo <= addr < sys_hi: return self.paddr_to_host(addr - self.gpu.mc_base)
for tva, (pa, sz, is_sys) in self.tlb.items():
if tva <= addr < tva + sz:
paddr = pa + (addr - tva)
if not is_sys: return self.gpu.vram_addr + paddr
return self.paddr_to_host(paddr)
raise ValueError(f"addr {addr:#x} not mapped (sys_aperture=[{sys_lo:#x}, {sys_hi:#x}])")
class MockIPBlock:
def __init__(self, gpu:MockAMGPU, mmio:MockMMIOInterface, regs:dict[str, AMDReg]):
self.gpu, self.mmio, self._regs = gpu, mmio, regs
self._n2a = {n: r.addr[0] for n, r in regs.items()}
self._a2n = {a: n for n, a in self._n2a.items()}
self.addrs = set(self._n2a.values())
def reg(self, name) -> int|None: return self._n2a.get(name)
def decode(self, name) -> dict: return self._regs[name].decode(self.mmio.regs.get(self._n2a[name], 0))
def read(self, reg:int) -> int: return self.mmio.regs.get(reg, 0)
def write(self, reg:int, val:int): self.mmio.regs[reg] = val
def _read_pair(self, pair) -> int:
if pair[0] is None: return 0
return self.mmio.regs.get(pair[0], 0) | (self.mmio.regs.get(pair[1], 0) << 32)
class MockPSP(MockIPBlock):
def __init__(self, gpu, mmio):
super().__init__(gpu, mmio, _build_ip_regs('mp', am.MP0_HWIP))
self._sos_alive, self._ring_wptr = False, 0
pref = "regMPASP_SMN_C2PMSG" if IP_VERSIONS[am.MP0_HWIP] >= (14,0,0) else "regMP0_SMN_C2PMSG"
def r(n): return self.reg(f"{pref}_{n}")
self._c2pmsg_35, self._c2pmsg_64, self._c2pmsg_67 = r(35), r(64), r(67)
self._c2pmsg_69, self._c2pmsg_70, self._c2pmsg_81 = r(69), r(70), r(81)
def read(self, reg:int) -> int:
if reg == self._c2pmsg_35: return 0x80000000
if reg == self._c2pmsg_81: return 0x1 if self._sos_alive else 0x0
if reg == self._c2pmsg_64: return 0x80000000 if self._sos_alive else 0x0
if reg == self._c2pmsg_67: return self._ring_wptr
return super().read(reg)
def write(self, reg:int, val:int):
super().write(reg, val)
if reg == self._c2pmsg_35 and val == am.PSP_BL__LOAD_SOSDRV: self._sos_alive = True
if reg == self._c2pmsg_67: self._ring_submit(val)
def _ring_submit(self, new_wptr:int):
old_wptr = self._ring_wptr
self._ring_wptr = new_wptr
lo, hi = self._c2pmsg_69, self._c2pmsg_70
if lo is None or hi is None: return
ring_mc = self.mmio.regs.get(lo, 0) | (self.mmio.regs.get(hi, 0) << 32)
ring_paddr = ring_mc - self.gpu.mc_base
frame_off = ring_paddr + old_wptr * 4
frame = am.struct_psp_gfx_rb_frame.from_buffer_copy(bytes(self.gpu.vram[frame_off:frame_off + ctypes.sizeof(am.struct_psp_gfx_rb_frame)]))
fence_paddr = ((frame.fence_addr_hi << 32) | frame.fence_addr_lo) - self.gpu.mc_base
if 0 <= fence_paddr < len(self.gpu.vram):
struct.pack_into('<I', self.gpu.vram, fence_paddr, frame.fence_value)
cmd_paddr = ((frame.cmd_buf_addr_hi << 32) | frame.cmd_buf_addr_lo) - self.gpu.mc_base
if 0 <= cmd_paddr < len(self.gpu.vram):
struct.pack_into('<I', self.gpu.vram, cmd_paddr + 864, 0)
class MockSMU(MockIPBlock):
def __init__(self, gpu, mmio):
regs = import_asic_regs('mp', (11, 0, 0), cls=functools.partial(AMDReg, bases={0: IP_BASES[am.MP1_HWIP]}))
super().__init__(gpu, mmio, regs)
self._msg_pending = False
def r(n): return self.reg(f"mmMP1_SMN_C2PMSG_{n}")
self._c2pmsg_53, self._c2pmsg_54, self._c2pmsg_66 = r(53), r(54), r(66)
self._c2pmsg_75, self._c2pmsg_82, self._c2pmsg_90 = r(75), r(82), r(90)
def read(self, reg:int) -> int:
if reg == self._c2pmsg_90 or reg == self._c2pmsg_54: return 0x1 if self._msg_pending else super().read(reg)
if reg == self._c2pmsg_82: return self.mmio.regs.get(reg, 3)
return super().read(reg)
def write(self, reg:int, val:int):
super().write(reg, val)
if reg == self._c2pmsg_66 or reg == self._c2pmsg_75: self._msg_pending = True
if (reg == self._c2pmsg_90 or reg == self._c2pmsg_54) and val == 0: self._msg_pending = False
class MockSDMA(MockIPBlock):
def __init__(self, gpu, mmio):
all_gc = _build_ip_regs('gc', am.GC_HWIP)
super().__init__(gpu, mmio, {n: r for n, r in all_gc.items() if 'SDMA' in n})
def write(self, reg:int, val:int):
super().write(reg, val)
name = self._a2n.get(reg, '')
if name.endswith('_RB_CNTL') and self._regs[name].decode(val).get('rb_enable', 0):
self._activate_queue(name.rsplit('_RB_CNTL', 1)[0])
def _activate_queue(self, prefix:str):
ring_addr = self._read_pair((self.reg(f'{prefix}_RB_BASE'), self.reg(f'{prefix}_RB_BASE_HI'))) << 8
rptr_addr = self._read_pair((self.reg(f'{prefix}_RB_RPTR_ADDR_LO'), self.reg(f'{prefix}_RB_RPTR_ADDR_HI')))
wptr_addr = self._read_pair((self.reg(f'{prefix}_RB_WPTR_POLL_ADDR_LO'), self.reg(f'{prefix}_RB_WPTR_POLL_ADDR_HI')))
rb_size = self.decode(f'{prefix}_RB_CNTL')['rb_size']
self.gpu.add_sdma_queue(self.gpu.mmu.addr_to_host(ring_addr), 4 << rb_size,
self.gpu.mmu.addr_to_host(rptr_addr), self.gpu.mmu.addr_to_host(wptr_addr))
class MockGFX(MockIPBlock):
def __init__(self, gpu, mmio):
super().__init__(gpu, mmio, _build_ip_regs('gc', am.GC_HWIP))
self._pt_base = (self.reg('regGCVM_CONTEXT0_PAGE_TABLE_BASE_ADDR_LO32'), self.reg('regGCVM_CONTEXT0_PAGE_TABLE_BASE_ADDR_HI32'))
self._pt_start = (self.reg('regGCVM_CONTEXT0_PAGE_TABLE_START_ADDR_LO32'), self.reg('regGCVM_CONTEXT0_PAGE_TABLE_START_ADDR_HI32'))
self._gc_inv_ack = self.reg('regGCVM_INVALIDATE_ENG17_ACK')
self._gc_inv_req = self.reg('regGCVM_INVALIDATE_ENG17_REQ')
self._hqd_active = self.reg('regCP_HQD_ACTIVE')
def read(self, reg:int) -> int:
if reg == self.reg('regCP_STAT') or reg == self.reg('regRLC_SAFE_MODE'): return 0
if reg == self.reg('regRLC_RLCS_BOOTLOAD_STATUS'): return 0x2
if reg == self._gc_inv_ack: return 0x1
return super().read(reg)
def write(self, reg:int, val:int):
super().write(reg, val)
if reg == self.reg('regCP_HQD_DEQUEUE_REQUEST'):
if self._hqd_active is not None: self.mmio.regs[self._hqd_active] = 0
if reg == self._hqd_active and val == 1: self._activate_pm4_queue()
if reg == self._gc_inv_req: self.gpu.mmu.invalidate(self.get_pt_base(), self.get_va_base())
def _activate_pm4_queue(self):
ring_addr = self._read_pair((self.reg('regCP_HQD_PQ_BASE'), self.reg('regCP_HQD_PQ_BASE_HI'))) << 8
rptr_addr = self._read_pair((self.reg('regCP_HQD_PQ_RPTR_REPORT_ADDR'), self.reg('regCP_HQD_PQ_RPTR_REPORT_ADDR_HI')))
wptr_addr = self._read_pair((self.reg('regCP_HQD_PQ_WPTR_POLL_ADDR'), self.reg('regCP_HQD_PQ_WPTR_POLL_ADDR_HI')))
queue_size = self.decode('regCP_HQD_PQ_CONTROL')['queue_size']
self.gpu.add_pm4_queue(self.gpu.mmu.addr_to_host(ring_addr), 4 << (queue_size + 1),
self.gpu.mmu.addr_to_host(rptr_addr), self.gpu.mmu.addr_to_host(wptr_addr))
def get_pt_base(self) -> int: return self._read_pair(self._pt_base) & 0x0000FFFFFFFFF000
def get_va_base(self) -> int: return self._read_pair(self._pt_start) << 12
class MockGMC(MockIPBlock):
def __init__(self, gpu, mmio, gfx:MockGFX):
super().__init__(gpu, mmio, _build_ip_regs('mmhub', am.MMHUB_HWIP))
self._gfx = gfx
self._inv_ack = self.reg('regMMVM_INVALIDATE_ENG17_ACK')
self._inv_sem = self.reg('regMMVM_INVALIDATE_ENG17_SEM')
self._inv_req = self.reg('regMMVM_INVALIDATE_ENG17_REQ')
self._fb_loc_top = self.reg('regMMMC_VM_FB_LOCATION_TOP')
def read(self, reg:int) -> int:
if reg == self._inv_ack or reg == self._inv_sem: return 0x1
if reg == self._fb_loc_top: return VRAM_SIZE >> 24
return super().read(reg)
def write(self, reg:int, val:int):
super().write(reg, val)
if reg == self._inv_req: self.gpu.mmu.invalidate(self._gfx.get_pt_base(), self._gfx.get_va_base())
class MockNBIO(MockIPBlock):
def __init__(self, gpu, mmio):
regs = _build_ip_regs('nbif', am.NBIO_HWIP)
regs.update(_build_ip_regs('hdp', am.HDP_HWIP))
super().__init__(gpu, mmio, regs)
self._remap_hdp = self.reg('regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL')
self._hdp_flush = self.reg('regHDP_MEM_FLUSH_CNTL')
def read(self, reg:int) -> int:
if reg == self._remap_hdp and self._hdp_flush is not None: return self._hdp_flush * 4
return super().read(reg)
class MockMMIOInterface:
def __init__(self, gpu:MockAMGPU):
self.gpu = gpu
self.regs: dict[int, int] = {}
gfx = MockGFX(gpu, self)
self.gmc = MockGMC(gpu, self, gfx)
self.blocks = [MockPSP(gpu, self), MockSMU(gpu, self), MockSDMA(gpu, self), gfx, self.gmc, MockNBIO(gpu, self)]
self._addr_block: dict[int, MockIPBlock] = {}
for block in self.blocks:
for addr in block.addrs: self._addr_block.setdefault(addr, block)
def __getitem__(self, index:int|slice) -> int|list[int]:
if isinstance(index, slice): return [self[i] for i in range(index.start or 0, index.stop or 0, index.step or 1)] # type: ignore[misc]
if index == 0xde3: return VRAM_SIZE >> 20
if block := self._addr_block.get(index): return block.read(index)
return self.regs.get(index, 0)
def __setitem__(self, index:int|slice, val:int|list[int]|tuple[int, ...]):
if isinstance(index, slice):
vals = val if isinstance(val, (list, tuple)) else [val] * ((index.stop - index.start) // (index.step or 1)) # type: ignore[operator]
for i, v in zip(range(index.start or 0, index.stop or 0, index.step or 1), vals): self[i] = v
return
assert isinstance(val, int)
self.regs[index] = val
if block := self._addr_block.get(index): block.write(index, val)
def __len__(self): return 0x10000000
class MockAMGPU(AMDGPU):
def __init__(self, gpuid:int=0):
super().__init__(gpuid)
self.vram_fd = libc.memfd_create(b"vram", libc.MFD_CLOEXEC)
os.ftruncate(self.vram_fd, VRAM_SIZE)
self.vram_addr = libc.mmap(0, VRAM_SIZE, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, self.vram_fd, 0)
self.vram = (ctypes.c_ubyte * VRAM_SIZE).from_address(self.vram_addr)
self.doorbell_fd = libc.memfd_create(b"doorbell", libc.MFD_CLOEXEC)
os.ftruncate(self.doorbell_fd, 0x2000)
self.arch = "rdna4"
self._sysmem_map:dict[int,int] = {}
self._next_sysmem_paddr = 0x100000000
self.mmu = MockMMU(self)
self.mmio = MockMMIOInterface(self)
self._preboot()
def translate_addr(self, addr:int) -> int: return self.mmu.addr_to_host(addr)
def map_vram_at(self, va:int, paddr:int, size:int):
libc.mmap(va, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | 0x10, self.vram_fd, paddr)
def _preboot(self):
ip_data = bytearray()
for hwip, (major, minor, rev) in IP_VERSIONS.items():
ip = am.struct_ip_v4(hw_id=IP_HWIDS[hwip], num_base_address=len(IP_BASES[hwip]), major=major, minor=minor, revision=rev)
ip_data += bytes(ip) + b'\x00'
for b in IP_BASES[hwip]: ip_data += struct.pack('<I', b)
dhdr = am.struct_die_header(num_ips=len(IP_VERSIONS))
ihdr = am.struct_ip_discovery_header(signature=am.DISCOVERY_TABLE_SIGNATURE, version=4, num_dies=1)
ip_disc_off = ctypes.sizeof(am.struct_binary_header)
ihdr.die_info[0].die_offset = ip_disc_off + ctypes.sizeof(am.struct_ip_discovery_header)
gc = am.struct_gc_info_v2_1()
gc.header.table_id, gc.header.version_major, gc.header.version_minor = am.GC, 2, 1
gc.header.size = ctypes.sizeof(am.struct_gc_info_v2_1)
for field, val in GC_INFO.items(): setattr(gc, field, val)
gc_off = ip_disc_off + ctypes.sizeof(am.struct_ip_discovery_header) + ctypes.sizeof(am.struct_die_header) + len(ip_data)
bhdr = am.struct_binary_header(binary_signature=am.BINARY_SIGNATURE)
bhdr.table_list[am.IP_DISCOVERY].offset = ip_disc_off
bhdr.table_list[am.GC].offset = gc_off
tbl = bytes(bhdr) + bytes(ihdr) + bytes(dhdr) + ip_data + bytes(gc)
tbl_offset = VRAM_SIZE - (64 << 10)
self.vram[tbl_offset:tbl_offset + len(tbl)] = list(tbl)
@property
def mc_base(self) -> int:
fb_loc_base = self.mmio.gmc.reg('regMMMC_VM_FB_LOCATION_BASE') or 0
return (self.mmio.regs.get(fb_loc_base, 0) & 0xFFFFFF) << 24

38
artifacts/package_sources/tinygrad/test/mockgpu/amd/README Normal file
View File

@@ -0,0 +1,38 @@
An integrated environment for AMD GPU assembly and emulation
Test with `pytest -n12 test/amd/`
`DEV=AMD:LLVM pytest -n12 test/amd/`
* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `DEV=MOCK{KFD|KFD|USB}+AMD`
* generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
* sqtt.py -- SQTT parser
The code should be as readable and deduplicated as possible. emu (in test/mockgpu/amd/) shouldn't be required for dsl.
The autogen folder is autogenerated from the AMD PDFs with `python3 -m tinygrad.renderer.amd.pdf --arch all`
test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
`DEV=MOCK+AMD pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
`DEV=MOCK+AMD:LLVM pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
The ops tests also pass, but they are very slow, so you should run them one at a time.
`SKIP_SLOW_TEST=1 DEV=MOCK+AMD pytest -n=12 test/backend/test_ops.py`
`SKIP_SLOW_TEST=1 DEV=NOCK+AMD:LLVM pytest -n=12 test/backend/test_ops.py`
When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
If a test is failing with `DEV=MOCK+AMD` it's because an instruction is emulated incorrectly.
You can test with just `DEV=AMD` to test on real hardware, if it works on real hardware there's a bug in the emulator.
IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.
Get line count with `cloc --by-file tinygrad/renderer/amd/*.py`

191
artifacts/package_sources/tinygrad/test/mockgpu/amd/amddriver.py Normal file
View File

@@ -0,0 +1,191 @@
import pathlib, re, ctypes, mmap, collections, functools, copy, os
from tinygrad.runtime.autogen import kfd, amdgpu_drm, libc
import tinygrad.runtime.autogen.am.am as am
from tinygrad.helpers import from_mv
from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props, GFX_TARGET_VERSION, MOCKGPU_ARCH
def ioctls_from_header():
# hdrpy = (pathlib.Path(__file__).parent.parent.parent.parent / "tinygrad" / "runtime" / "autogen" / "kfd.py").read_text()
# pattern = r'# (AMDKFD_IOC_[A-Z0-9_]+)\s=\s_(IOW?R?).*\(( 0x[0-9a-fA-F]+) ,\s+struct\s([A-Za-z0-9_]+)\s+\)'
# matches = re.findall(pattern, hdrpy, re.MULTILINE)
hdr = (pathlib.Path(__file__).parent.parent.parent.parent / "extra" / "hip_gpu_driver" / "kfd_ioctl.h").read_text().replace("\\\n", "")
pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
matches = re.findall(pattern, hdr, re.MULTILINE)
return type("KFD_IOCTLS", (object, ), {name: int(nr, 0x10) for name, _, nr, _ in matches}), \
{int(nr, 0x10): getattr(kfd, "struct_"+sname, None) for name, idir, nr, sname in matches}
kfd_ioctls, kfd_headers = ioctls_from_header()
class KFDFileDesc(VirtFileDesc):
def __init__(self, fd, driver):
super().__init__(fd)
self.driver = driver
def ioctl(self, fd, request, argp): return self.driver.kfd_ioctl(request, argp)
def mmap(self, start, sz, prot, flags, fd, offset): return offset
class DRMFileDesc(VirtFileDesc):
def __init__(self, fd, driver, gpu):
super().__init__(fd)
self.driver, self.gpu = driver, gpu
def ioctl(self, fd, request, argp):
struct = amdgpu_drm.struct_drm_amdgpu_info.from_address(argp)
if struct.query == amdgpu_drm.AMDGPU_INFO_DEV_INFO:
dev_info = amdgpu_drm.struct_drm_amdgpu_info_device.from_address(struct.return_pointer)
# mock of gfx1100
for se in range(4):
for sa in range(4): dev_info.cu_bitmap[se][sa] = 0xff if (se * 4 + sa) < 12 else 0
return 0
raise NotImplementedError(f"unknown DRM ioctl query {struct.query}")
def mmap(self, start, sz, prot, flags, fd, offset): return libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
class AMDDriver(VirtDriver):
def __init__(self, gpus=6):
super().__init__()
# NOTE: gpu ids start from one (id 0 is skipped in KFDIface._is_usable_gpu)
self.tracked_files += [VirtFile('/dev/kfd', functools.partial(KFDFileDesc, driver=self))] + \
[VirtFile('/sys/devices/virtual/kfd/kfd/topology/nodes', functools.partial(DirFileDesc, child_names=[str(i+1) for i in range(gpus)]))]
self.gpus = {}
self.next_fd = (1 << 30)
self.next_handle = 1
self.next_event = 1
self.object_by_handle = {}
self.doorbells = {}
self.next_doorbell = collections.defaultdict(int)
self.mmu_event_ids = []
self._executing = False # re-entrancy guard for _emulate_execute
for i in range(gpus): self._prepare_gpu(i+1)
def _alloc_fd(self):
my_fd = self.next_fd
self.next_fd = self.next_fd + 1
return my_fd
def _alloc_handle(self):
handle = self.next_handle
self.next_handle += 1
return handle
def _alloc_next_event_slot(self):
ev = self.next_event
self.next_event += 1
return ev
def _alloc_doorbell(self, gpu_id):
x = ctypes.addressof(from_mv(self.doorbells[gpu_id])) + self.next_doorbell[gpu_id] * 8
self.next_doorbell[gpu_id] += 1
return x
def _prepare_gpu(self, gpu_id):
self.doorbells[gpu_id] = memoryview(bytearray(0x2000))
self.gpus[gpu_id] = AMDGPU(gpu_id)
ip_versions = {"rdna3": {"gc": (11, 0, 0), "sdma": (6, 0, 0), "nbif": (4, 3, 0)},
"rdna4": {"gc": (12, 0, 0), "sdma": (6, 0, 0), "nbif": (6, 3, 1)},
"cdna4": {"gc": (9, 5, 0), "sdma": (4, 4, 5), "nbif": (7, 9, 0)}}[MOCKGPU_ARCH]
def ip_discovery_files(hwid, ver, base_addr):
p = f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}/0'
return [VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}', functools.partial(DirFileDesc, child_names=['0'])),
VirtFile(f'{p}/major', functools.partial(TextFileDesc, text=str(ver[0]))),
VirtFile(f'{p}/minor', functools.partial(TextFileDesc, text=str(ver[1]))),
VirtFile(f'{p}/revision', functools.partial(TextFileDesc, text=str(ver[2]))),
VirtFile(f'{p}/base_addr', functools.partial(TextFileDesc, text=base_addr))]
self.tracked_files += [
VirtFile('/sys/module/amdgpu', functools.partial(TextFileDesc, text="1")),
VirtFile('/sys/module/amdgpu/parameters/ppfeaturemask', functools.partial(TextFileDesc, text="0xffff3fff")),
VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}', functools.partial(DirFileDesc, child_names=['gpu_id', 'properties'])),
VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/gpu_id', functools.partial(TextFileDesc, text=f"{gpu_id}")),
VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/properties',
functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id, gfx_target_version=GFX_TARGET_VERSION))),
VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/power_dpm_force_performance_level',
functools.partial(TextFileDesc, text='profile_standard\n')),
VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
*ip_discovery_files(am.GC_HWID, ip_versions["gc"], '0x00001260\n0x0000A000\n0x0001C000\n0x02402C00'),
*ip_discovery_files(am.SDMA0_HWID, ip_versions["sdma"], '0x00001260\n0x0000A000\n0x0001C000\n0x02402C00'),
*ip_discovery_files(am.NBIF_HWID, ip_versions["nbif"], '0x00000000\n0x00000014\n0x00000D20\n0x00010400\n0x0241B000\n0x04040000'),
VirtFile(f'/dev/dri/renderD{gpu_id}', functools.partial(DRMFileDesc, driver=self, gpu=f"{self.gpus[gpu_id]}")),
]
def open(self, name, flags, mode, virtfile): return virtfile.fdcls(self._alloc_fd())
def kfd_ioctl(self, req, argp):
nr = req & 0xFF
struct = kfd_headers[nr].from_address(argp)
if nr == kfd_ioctls.AMDKFD_IOC_ACQUIRE_VM: pass
elif nr == kfd_ioctls.AMDKFD_IOC_RUNTIME_ENABLE: pass
elif nr == kfd_ioctls.AMDKFD_IOC_GET_VERSION:
struct.major_version = 1
struct.minor_version = 14
elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
if struct.gpu_id not in self.gpus: return -1
struct.handle = self._alloc_handle()
self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
# Track signal memory (uncached + coherent) - progress queues when written to
if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
self.object_by_handle.pop(struct.handle)
elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
for i in range(struct.n_devices):
gpu = self.gpus[dev_ids[i]]
mem_obj = self.object_by_handle[struct.handle]
gpu.map_range(mem_obj.va_addr, mem_obj.size)
struct.n_success = i + 1
elif nr == kfd_ioctls.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU:
dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
for i in range(struct.n_devices):
gpu = self.gpus[dev_ids[i]]
mem_obj = self.object_by_handle[struct.handle]
gpu.unmap_range(mem_obj.va_addr, mem_obj.size)
struct.n_success = i + 1
elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
struct.event_slot_index = self._alloc_next_event_slot()
struct.event_id = struct.event_slot_index
if struct.event_type == kfd.KFD_IOC_EVENT_MEMORY: self.mmu_event_ids.append(struct.event_id)
elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
gpu = self.gpus[struct.gpu_id]
if struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
gpu.add_sdma_queue(struct.ring_base_address, struct.ring_size, struct.read_pointer_address, struct.write_pointer_address)
elif struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE:
gpu.add_pm4_queue(struct.ring_base_address, struct.ring_size, struct.read_pointer_address, struct.write_pointer_address)
else: raise RuntimeError("Unsuported, queue")
# Track writes to doorbell, calling callback
struct.doorbell_offset = self._alloc_doorbell(struct.gpu_id)
self.track_address(struct.doorbell_offset, struct.doorbell_offset + 8, lambda mv,off: None, lambda mv, off: self._emulate_execute())
elif nr == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
evs = (kfd.struct_kfd_event_data * struct.num_events).from_address(struct.events_ptr)
for ev in evs:
if ev.event_id in self.mmu_event_ids and "MOCKGPU_EMU_FAULTADDR" in os.environ:
ev.memory_exception_data.gpu_id = 1
ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
ev.memory_exception_data.failure.NotPresent = 1
else:
name = "unknown"
for k,v in kfd_ioctls.__dict__.items():
if nr == v: name = k
assert False, f"unknown kfd ioctl, {nr} {name}"
exit(1)
return 0
def _emulate_execute(self):
if self._executing: return # prevent re-entrancy
self._executing = True
try:
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.executing: any_progress |= q.execute() > 0
finally:
self._executing = False

448
artifacts/package_sources/tinygrad/test/mockgpu/amd/amdgpu.py Normal file
View File

@@ -0,0 +1,448 @@
import ctypes, time
from dataclasses import replace
from test.mockgpu.gpu import VirtGPU
from test.mockgpu.helpers import PythonRemu
from tinygrad.helpers import getbits, to_mv, getenv, DEV
from tinygrad.runtime.support import c
MOCKGPU_ARCH = "cdna4" if DEV.arch == "gfx950" else "rdna4" if DEV.arch.startswith("gfx12") else "rdna3"
assert (ma:=getenv("MOCKGPU_ARCH", "")) == "", "MOCKGPU_ARCH is deprecated, use DEV=" + \
str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) # type: ignore
GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH]
import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.sdma_6_0_0 as sdma
SDMA_MAX_COPY_SIZE = 0x400000
regCOMPUTE_PGM_LO = 0x1bac + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_PGM_RSRC2 = 0x1bb3 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_TMPRING_SIZE = 0x1bb8 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0
regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1
regSQ_THREAD_TRACE_BUF0_BASE = 0x39e8 + amd_gpu.GC_BASE__INST0_SEG1
regSQ_THREAD_TRACE_BUF0_SIZE = {"rdna3": 0x39e9, "rdna4": 0x39e6, "cdna4": 0x39e9}[MOCKGPU_ARCH] + amd_gpu.GC_BASE__INST0_SEG1
regSQ_THREAD_TRACE_WPTR = 0x39ef + amd_gpu.GC_BASE__INST0_SEG1
regSQ_THREAD_TRACE_STATUS = 0x39f4 + amd_gpu.GC_BASE__INST0_SEG1
regCP_PERFMON_CNTL = 0x3808 + amd_gpu.GC_BASE__INST0_SEG1
regCPG_PERFCOUNTER1_LO = 0x3000 + amd_gpu.GC_BASE__INST0_SEG1
regGUS_PERFCOUNTER_HI = 0x3643 + amd_gpu.GC_BASE__INST0_SEG1
# RDNA 4
regSQ_THREAD_TRACE_BUF0_BASE_LO = 0x39e7 + amd_gpu.GC_BASE__INST0_SEG1
regSQ_THREAD_TRACE_BUF0_BASE_HI = regSQ_THREAD_TRACE_BUF0_BASE
class SQTT_EVENTS:
THREAD_TRACE_FINISH = 0x00000037
CACHE_FLUSH_AND_INV_TS_EVENT = 0x14
WAIT_REG_MEM_FUNCTION_ALWAYS = 0
WAIT_REG_MEM_FUNCTION_EQ = 3 # ==
WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
remu = PythonRemu()
def create_sdma_packets():
# TODO: clean up this, if we want to keep it
structs = {}
for name,pkt in [(name,s) for name,s in amd_gpu.__dict__.items() if name.startswith("rocr_AMD_SDMA_PKT_") and name.endswith("_TAG")]:
names = set()
fields = []
for pkt_fields in pkt._real_fields_:
if not pkt_fields[0].endswith("_UNION"): fields.append(pkt_fields)
else:
for union_fields in pkt_fields[1]._real_fields_[:-1]:
fname = union_fields[0]
if fname in names: fname = pkt_fields[0]+fname
names.add(fname)
# merge together 64-bit fields, otherwise just append them
if fname.endswith("_63_32") and fields[-1][0].endswith("_31_0"): fields[-1] = (fname[:-6], ctypes.c_ulong, fields[-1][2], 64, 0)
else: fields.append((fname, union_fields[1], union_fields[2] + pkt_fields[2], *union_fields[3:]))
new_name = name[18:-4].lower()
structs[new_name] = c.init_c_struct_t(ctypes.sizeof(pkt), tuple(fields))
return type("SDMA_PKTS", (object, ), structs)
sdma_pkts = create_sdma_packets()
class AMDQueue:
def __init__(self, base, size, rptr, wptr):
self.queue, self.size = to_mv(base, size).cast("I"), size
self.rptr = to_mv(rptr, 8).cast("Q") if isinstance(rptr, int) else rptr
self.wptr = to_mv(wptr, 8).cast("Q") if isinstance(wptr, int) else wptr
@property
def executing(self): return self.rptr[0] < self.wptr[0]
class PM4Executor(AMDQueue):
def __init__(self, gpu, base, size, rptr, wptr):
self.gpu = gpu
self.ib_executor: PM4Executor|None = None
super().__init__(base, size, rptr, wptr)
def _next_dword(self):
x = self.queue[self.rptr[0] % (self.size // 4)]
self.rptr[0] += 1
return x
@property
def executing(self): return self.rptr[0] < self.wptr[0] or self.ib_executor is not None
def execute(self):
prev_rptr, executed_in_ib, cont = self.rptr[0], 0, True
while self.executing and cont:
if self.ib_executor is not None:
executed_in_ib += self.ib_executor.execute()
if self.ib_executor.executing: break
self.ib_executor = None
continue # this continue is needed if PACKET3_INDIRECT_BUFFER is the last packet and rptr == wptr
header = self._next_dword()
packet_type = header >> 30
op = (header >> 8) & 0xFF
n = (header >> 16) & 0x3FFF
assert packet_type == 3, "Can parse only packet3"
if op == amd_gpu.PACKET3_SET_SH_REG: self._exec_set_reg(n, pm4.PACKET3_SET_SH_REG_START)
elif op == amd_gpu.PACKET3_SET_UCONFIG_REG: self._exec_set_reg(n, pm4.PACKET3_SET_UCONFIG_REG_START)
elif op == amd_gpu.PACKET3_ACQUIRE_MEM: self._exec_acquire_mem(n)
elif op == amd_gpu.PACKET3_RELEASE_MEM: self._exec_release_mem(n)
elif op == amd_gpu.PACKET3_COPY_DATA: self._exec_copy_data(n)
elif op == amd_gpu.PACKET3_WAIT_REG_MEM: cont = self._exec_wait_reg_mem(n)
elif op == amd_gpu.PACKET3_DISPATCH_DIRECT: self._exec_dispatch_direct(n)
elif op == amd_gpu.PACKET3_INDIRECT_BUFFER: self._exec_indirect_buffer(n)
elif op == amd_gpu.PACKET3_EVENT_WRITE: self._exec_event_write(n)
else: raise RuntimeError(f"PM4: Unknown opcode: {op}")
return (self.rptr[0] - prev_rptr) + executed_in_ib
def _exec_acquire_mem(self, n):
assert n in (5, 6)
for _ in range(n + 1): self._next_dword() # TODO: implement
def _exec_release_mem(self, n):
assert n == 6
mem_event_type = (self._next_dword() >> 0) & 0xff
selectors = self._next_dword()
mem_data_sel = (selectors >> 29) & 0b111
# int_sel = (selectors >> 24) & 0b11
# mem_dst_sel = (selectors >> 16) & 0b1
addr_lo = self._next_dword()
addr_hi = self._next_dword()
val_lo = self._next_dword()
val_hi = self._next_dword()
val = val_lo + (val_hi << 32)
_ = self._next_dword() # ev
ptr = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
if mem_data_sel == 1: ptr.cast('I')[0] = val & 0xffffffff
elif mem_data_sel == 2: ptr.cast('Q')[0] = val
elif mem_data_sel == 3:
if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
elif mem_data_sel == 0: pass # no write
else: raise RuntimeError(f"Unknown {mem_data_sel=}")
def _exec_copy_data(self, n):
assert n == 4
copy_data_flags = self._next_dword()
src_addr_lo = self._next_dword()
_src_addr_hi = self._next_dword()
dst_addr_lo = self._next_dword()
dst_addr_hi = self._next_dword()
assert copy_data_flags in {0x100204, 0x000204}, hex(copy_data_flags) # better fail than silently do the wrong thing
to_mv(self.gpu.translate_addr(dst_addr_hi<<32|dst_addr_lo), 4).cast('I')[0] = self.gpu.regs[src_addr_lo]
def _exec_wait_reg_mem(self, n):
assert n == 5
info = self._next_dword()
addr_lo = self._next_dword()
addr_hi = self._next_dword()
val = self._next_dword()
mask = self._next_dword()
_ = self._next_dword() # timeout
mem_function = (info >> 0) & 0b111
mem_space = (info >> 4) & 0b1
mem_op = (info >> 6) & 0b1
_ = (info >> 8) & 0b1 # mem_engine
if mem_space == 0 and mem_op == 1: mval = val # hack for memory barrier, should properly handle (req_req, reg_done)
elif mem_space == 0: mval = self.gpu.regs[addr_hi<<32|addr_lo]
elif mem_space == 1: mval = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 4).cast('I')[0]
mval &= mask
if mem_function == WAIT_REG_MEM_FUNCTION_GEQ: can_cont = bool(mval >= val)
elif mem_function == WAIT_REG_MEM_FUNCTION_NEQ: can_cont = bool(mval != val)
elif mem_function == WAIT_REG_MEM_FUNCTION_EQ: can_cont = bool(mval == val)
else: raise RuntimeError(f"Do not support {mem_function=}")
if not can_cont: self.rptr[0] = self.rptr[0] - 7 # revert this packet, need to wait again
return can_cont
def _exec_set_reg(self, n, off):
reg = off + self._next_dword()
for i in range(n):
self.gpu.regs[reg] = self._next_dword()
reg += 1
def _exec_dispatch_direct(self, n):
assert n == 3
gl = [self._next_dword() for _ in range(3)]
_ = self._next_dword() # flags
prg_addr = (self.gpu.regs[regCOMPUTE_PGM_LO] + (self.gpu.regs[regCOMPUTE_PGM_LO + 1] << 32)) << 8
args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32)
lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)]
rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2]
# Read all user data registers (hardware loads these directly into s[0:N])
user_sgpr_count = (rsrc2 >> 1) & 0x1F # USER_SGPR_COUNT is bits 1:5
user_data = []
for i in range(user_sgpr_count):
try: user_data.append(self.gpu.regs[regCOMPUTE_USER_DATA_0 + i])
except KeyError: user_data.append(0)
prg_sz = 0
for st,sz in self.gpu.mapped_ranges:
if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
# Get scratch size from COMPUTE_TMPRING_SIZE register
# WAVESIZE = ceildiv(lanes * size_per_thread, mem_alignment_size)
# GFX11+: mem_alignment_size=256, so size_per_thread = WAVESIZE * 256 / 64 = WAVESIZE * 4
# GFX9: mem_alignment_size=1024, so size_per_thread = WAVESIZE * 1024 / 64 = WAVESIZE * 16
try: tmpring_size = self.gpu.regs[regCOMPUTE_TMPRING_SIZE]
except KeyError: tmpring_size = 0
wavesize = (tmpring_size >> 12) & 0x3FFF # WAVESIZE field is bits 12:25
scratch_size = wavesize * (16 if self.gpu.arch == "cdna" else 4) # per-thread scratch size in bytes
assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
# Pass valid memory ranges, rsrc2, scratch_size, arch, and user data registers to the emulator
remu.valid_mem_ranges = self.gpu.mapped_ranges
remu.rsrc2 = rsrc2
remu.scratch_size = scratch_size
remu.arch = self.gpu.arch
remu.user_data = user_data
err = remu.run_asm(prg_addr, prg_sz, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr)
if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
def _exec_indirect_buffer(self, n):
addr_lo = self._next_dword()
addr_hi = self._next_dword()
buf_sz = self._next_dword() & (0x7fffff)
rptr = memoryview(bytearray(8)).cast('Q')
wptr = memoryview(bytearray(8)).cast('Q')
rptr[0] = 0
wptr[0] = buf_sz
self.ib_executor = PM4Executor(self.gpu, self.gpu.translate_addr((addr_hi << 32) | addr_lo), buf_sz * 4, rptr, wptr)
def _exec_event_write(self, n):
assert n == 0
event_dw = self._next_dword()
match (event_dw & 0xFF): # event type
case SQTT_EVENTS.THREAD_TRACE_FINISH:
# Get the most recent trace from the emulator (if available)
from test.mockgpu.amd.emu import sqtt_traces
blob = sqtt_traces.pop(0) if sqtt_traces else b''
old_idx = self.gpu.regs.grbm_index
for se in range(self.gpu.regs.n_se):
self.gpu.regs.grbm_index = 0b011 << 29 | se << 16 # select se, broadcast sa and instance
self.gpu.regs[regSQ_THREAD_TRACE_STATUS] = 1 << 12 # FINISH_PENDING==0 FINISH_DONE==1 BUSY==0
if MOCKGPU_ARCH == "rdna3":
buf_addr = ((self.gpu.regs[regSQ_THREAD_TRACE_BUF0_SIZE]&0xf)<<32|self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE])<<12
else:
buf_addr = ((self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE_HI])<<32|self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE_LO])<<12
# Use real trace blob for SE 0 (which has itrace enabled), empty blob for other SEs
se_blob = blob if se == 0 else b''
# Write blob to trace buffer
if se_blob: ctypes.memmove(buf_addr, se_blob, len(se_blob))
# RDNA3 has absolute address for wptr, RDNA4 has relative
wptr_val = (((buf_addr if MOCKGPU_ARCH == "rdna3" else 0) + len(se_blob)) // 32) & 0x1FFFFFFF
self.gpu.regs[regSQ_THREAD_TRACE_WPTR] = wptr_val
self.gpu.regs.grbm_index = old_idx
case _: pass # NOTE: for now most events aren't emulated
class SDMAExecutor(AMDQueue):
def __init__(self, gpu, base, size, rptr, wptr):
self.gpu, self.base = gpu, base
super().__init__(base, size, rptr, wptr)
def execute(self):
prev_rptr, cont = self.rptr[0], True
while self.executing and cont:
header = self.queue[(self.rptr[0] // 4) % (self.size // 4)]
op = (header >> 0) & 0xff
if op == 0: self.rptr[0] += 4
elif op == amd_gpu.SDMA_OP_FENCE: self._execute_fence()
elif op == amd_gpu.SDMA_OP_TRAP: self._execute_trap()
elif op == amd_gpu.SDMA_OP_POLL_REGMEM: cont = self._execute_poll_regmem()
elif op == amd_gpu.SDMA_OP_GCR: self._execute_gcr()
elif op == amd_gpu.SDMA_OP_COPY: self._execute_copy()
elif op == sdma.SDMA_OP_WRITE: self._execute_write()
elif op == amd_gpu.SDMA_OP_TIMESTAMP: self._execute_timestamp()
elif op == 32: self.rptr[0] += 4 # SDMA_OP_DUMMY_TRAP: pipeline flush, no interrupt
else: raise RuntimeError(f"Unknown SDMA op {op}")
return self.rptr[0] - prev_rptr
def _execute_fence(self):
struct = sdma_pkts.fence.from_address(self.base + self.rptr[0] % self.size)
to_mv(self.gpu.translate_addr(struct.addr), 8).cast('Q')[0] = struct.data
self.rptr[0] += ctypes.sizeof(struct)
def _execute_trap(self):
struct = sdma_pkts.trap.from_address(self.base + self.rptr[0] % self.size)
self.rptr[0] += ctypes.sizeof(struct)
def _execute_write(self):
packet = to_mv(self.base + self.rptr[0] % self.size, 16).cast('I')
addr, count = packet[1] | packet[2] << 32, packet[3] + 1
ctypes.memmove(self.gpu.translate_addr(addr), self.base + self.rptr[0] % self.size + 16, count * 4)
self.rptr[0] += (4 + count) * 4
def _execute_poll_regmem(self):
struct = sdma_pkts.poll_regmem.from_address(self.base + self.rptr[0] % self.size)
if struct.mem_poll == 0: mval = struct.value & struct.mask
elif struct.mem_poll == 1: mval = to_mv(self.gpu.translate_addr(struct.addr), 4).cast('I')[0] & struct.mask
if struct.func == WAIT_REG_MEM_FUNCTION_GEQ: can_cont = bool(mval >= struct.value)
elif struct.func == WAIT_REG_MEM_FUNCTION_EQ: can_cont = bool(mval == struct.value)
elif struct.func == WAIT_REG_MEM_FUNCTION_ALWAYS: can_cont = True
else: raise RuntimeError(f"Do not support {struct.func=}")
if not can_cont: return False
self.rptr[0] += ctypes.sizeof(struct)
return True
def _execute_timestamp(self):
struct = sdma_pkts.timestamp.from_address(self.base + self.rptr[0] % self.size)
mem = to_mv(self.gpu.translate_addr(struct.addr), 8).cast('Q')
mem[0] = int(time.perf_counter() * 1e8)
self.rptr[0] += ctypes.sizeof(struct)
def _execute_gcr(self):
struct = sdma_pkts.gcr.from_address(self.base + self.rptr[0] % self.size)
self.rptr[0] += ctypes.sizeof(struct)
def _execute_copy(self):
struct = sdma_pkts.copy_linear.from_address(self.base + self.rptr[0] % self.size)
count_cnt = to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF
ctypes.memmove(self.gpu.translate_addr(struct.dst_addr), self.gpu.translate_addr(struct.src_addr), count_cnt + 1)
self.rptr[0] += ctypes.sizeof(struct)
class AMDGPURegisters:
def __init__(self, n_se:int=6):
self.n_se = n_se
self.grbm_index = 0b111 << 0x1d # all broadcast. NOTE: only per-se register emulation is currently supported
self.regs: dict[tuple[int, int], int] = {}
def __getitem__(self, addr:int) -> int:
if addr == regGRBM_GFX_INDEX: return self.grbm_index
if regCPG_PERFCOUNTER1_LO < addr < regGUS_PERFCOUNTER_HI:
assert self.regs[(regCP_PERFMON_CNTL, 0)] == 0x401, "read mode should be enabled"
return addr << 16 | self.grbm_index
return self.regs[(addr, getbits(self.grbm_index, 16, 23))]
def __setitem__(self, addr:int, val:int):
if addr == regGRBM_GFX_INDEX: self.grbm_index = val
if getbits(self.grbm_index, 31, 31):
for se in range(self.n_se): self.regs[(addr, se)] = val
else:
self.regs[(addr, getbits(self.grbm_index, 16, 23))] = val
class AMDGPU(VirtGPU):
def __init__(self, gpuid):
super().__init__(gpuid)
self.regs = AMDGPURegisters()
self.mapped_ranges = set()
self.queues = []
self.arch = "cdna" if MOCKGPU_ARCH == "cdna4" else MOCKGPU_ARCH
def translate_addr(self, addr:int) -> int: return addr
def map_range(self, vaddr, size): self.mapped_ranges.add((vaddr, size))
def unmap_range(self, vaddr, size): self.mapped_ranges.remove((vaddr, size))
def add_pm4_queue(self, base, size, rptr, wptr):
self.queues.append(PM4Executor(self, base, size, rptr, wptr))
return len(self.queues) - 1
def add_sdma_queue(self, base, size, rptr, wptr):
self.queues.append(SDMAExecutor(self, base, size, rptr, wptr))
return len(self.queues) - 1
_gpu_props_rdna = """cpu_cores_count 0
simd_count 192
mem_banks_count 1
caches_count 206
io_links_count 1
p2p_links_count 5
cpu_core_id_base 0
simd_id_base 2147488032
max_waves_per_simd 16
lds_size_in_kb 64
gds_size_in_kb 0
num_gws 64
wave_front_size 32
array_count 12
simd_arrays_per_engine 2
cu_per_simd_array 8
simd_per_cu 2
max_slots_scratch_cu 32
gfx_target_version {gfx_target_version}
vendor_id 4098
device_id 29772
location_id 34304
domain 0
drm_render_minor {drm_render_minor}
hive_id 0
num_sdma_engines 2
num_sdma_xgmi_engines 0
num_sdma_queues_per_engine 6
num_cp_queues 8
max_engine_clk_fcompute 2482
local_mem_size 0
fw_version 2140
capability 671588992
debug_prop 1495
sdma_fw_version 20
unique_id 11673270660693242239
num_xcc 1
max_engine_clk_ccompute 2400"""
_gpu_props_cdna = """cpu_cores_count 0
simd_count 304
mem_banks_count 1
caches_count 206
io_links_count 1
p2p_links_count 5
cpu_core_id_base 0
simd_id_base 2147488032
max_waves_per_simd 16
lds_size_in_kb 160
gds_size_in_kb 0
num_gws 64
wave_front_size 64
array_count 16
simd_arrays_per_engine 4
cu_per_simd_array 19
simd_per_cu 2
max_slots_scratch_cu 32
gfx_target_version {gfx_target_version}
vendor_id 4098
device_id 29772
location_id 34304
domain 0
drm_render_minor {drm_render_minor}
hive_id 0
num_sdma_engines 2
num_sdma_xgmi_engines 0
num_sdma_queues_per_engine 6
num_cp_queues 8
max_engine_clk_fcompute 2100
local_mem_size 0
fw_version 2140
capability 671588992
debug_prop 1495
sdma_fw_version 20
unique_id 11673270660693242239
num_xcc 1
max_engine_clk_ccompute 2100"""
gpu_props = _gpu_props_cdna if MOCKGPU_ARCH == "cdna4" else _gpu_props_rdna

2325
artifacts/package_sources/tinygrad/test/mockgpu/amd/emu.py Normal file
View File

File diff suppressed because it is too large Load Diff

1362
artifacts/package_sources/tinygrad/test/mockgpu/amd/pcode.py Normal file
View File

File diff suppressed because it is too large Load Diff

174
artifacts/package_sources/tinygrad/test/mockgpu/cuda/cuda.py Normal file
View File

@@ -0,0 +1,174 @@
from __future__ import annotations
from typing import Any
import ctypes, time
from test.mockgpu.helpers import ptx_run
from tinygrad.runtime.autogen import cuda as orig_cuda
from tinygrad.helpers import mv_address
for attr in dir(orig_cuda):
if not attr.startswith('__'):
globals()[attr] = getattr(orig_cuda, attr)
# Global state
class CUDAState:
def __init__(self):
self.memory: dict[int, memoryview] = {}
self.events: dict[int, float] = {} # Event ID -> timestamp
self.modules: dict[int, memoryview] = {} # Module ID -> code
self.current_context: int|None = None
self.contexts: dict[int, dict] = {} # Context ID -> context data
self.devices: dict[int, dict] = {} # Device ID -> device data
self.next_ptr = 1000 # For memory allocation
self.next_event_id = 1
self.next_module_id = 1
self.next_context_id = 1
cuda_state = CUDAState()
# Helper functions
def check_context():
if cuda_state.current_context is None:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
return orig_cuda.CUDA_SUCCESS
# CUDA API simulation
def cuInit(flags: int) -> int:
return orig_cuda.CUDA_SUCCESS
def cuDeviceGet(device, ordinal: int) -> int:
if ordinal < 0:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
device._obj.value = ordinal
cuda_state.devices[ordinal] = {"compute_capability": (3, 5)}
return orig_cuda.CUDA_SUCCESS
def cuCtxCreate_v2(pctx, flags: int, dev: int) -> int:
ctx_id = cuda_state.next_context_id
cuda_state.next_context_id += 1
cuda_state.contexts[ctx_id] = {"device": dev, "flags": flags}
pctx._obj.value = ctx_id
return orig_cuda.CUDA_SUCCESS
def cuCtxSetCurrent(context) -> int:
if context.value not in cuda_state.contexts:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
cuda_state.current_context = context.value
return orig_cuda.CUDA_SUCCESS
def cuMemAlloc_v2(dptr, bytesize: int) -> int:
x = memoryview(bytearray(bytesize))
dptr._obj.value = mv_address(x)
cuda_state.memory[dptr._obj.value] = x
return orig_cuda.CUDA_SUCCESS
def cuMemFree_v2(dptr) -> int:
if dptr.value in cuda_state.memory:
del cuda_state.memory[dptr.value]
return orig_cuda.CUDA_SUCCESS
return orig_cuda.CUDA_ERROR_INVALID_VALUE
def cuMemcpyHtoDAsync_v2(dst, src: ctypes.c_void_p, bytesize: int, stream: Any) -> int:
ctypes.memmove(dst if isinstance(dst, int) else dst.value, src, bytesize)
return orig_cuda.CUDA_SUCCESS
def cuMemcpyDtoH_v2(dst: ctypes.c_void_p, src, bytesize: int) -> int:
ctypes.memmove(dst, src if isinstance(src, int) else src.value, bytesize)
return orig_cuda.CUDA_SUCCESS
def cuEventCreate(phEvent, flags: int) -> int:
event_id = cuda_state.next_event_id
cuda_state.next_event_id += 1
cuda_state.events[event_id] = 0.0
phEvent._obj.value = event_id
return orig_cuda.CUDA_SUCCESS
def cuEventRecord(hEvent, hStream: Any) -> int:
if hEvent.value not in cuda_state.events:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
cuda_state.events[hEvent.value] = time.perf_counter_ns()
return orig_cuda.CUDA_SUCCESS
def cuEventSynchronize(hEvent) -> int:
if hEvent.value not in cuda_state.events:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
return orig_cuda.CUDA_SUCCESS
def cuEventElapsedTime(pMilliseconds, hStart, hEnd) -> int:
if hStart.value not in cuda_state.events or hEnd.value not in cuda_state.events:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
elapsed = (cuda_state.events[hEnd.value] - cuda_state.events[hStart.value]) * 1e-6
pMilliseconds._obj.value = elapsed
return orig_cuda.CUDA_SUCCESS
def cuEventDestroy_v2(hEvent) -> int:
if hEvent.value in cuda_state.events:
del cuda_state.events[hEvent.value]
return orig_cuda.CUDA_SUCCESS
def cuModuleLoadData(module, image: bytes) -> int:
module_id = cuda_state.next_module_id
cuda_state.next_module_id += 1
cuda_state.modules[module_id] = memoryview(bytearray(image))
module._obj.value = module_id
return orig_cuda.CUDA_SUCCESS
def cuModuleGetFunction(hfunc, hmod, name: bytes) -> int:
if hmod.value not in cuda_state.modules:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
hfunc._obj.value = mv_address(cuda_state.modules[hmod.value])
return orig_cuda.CUDA_SUCCESS
def cuModuleUnload(hmod) -> int:
if hmod.value in cuda_state.modules:
del cuda_state.modules[hmod.value]
return orig_cuda.CUDA_SUCCESS
def cuLaunchKernel(f, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, sharedMemBytes: int,
hStream: Any, kernelParams: Any, extra: Any) -> int:
cargs = [ctypes.cast(getattr(extra, field[0]), ctypes.c_void_p) for field in extra._real_fields_]
try: ptx_run(ctypes.cast(f.value, ctypes.c_char_p), len(cargs), (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
except Exception as e:
print("Error in cuLaunchKernel:", e)
return orig_cuda.CUDA_ERROR_LAUNCH_FAILED
return orig_cuda.CUDA_SUCCESS
def cuDeviceComputeCapability(major, minor, dev: int) -> int:
if dev not in cuda_state.devices:
return orig_cuda.CUDA_ERROR_INVALID_VALUE
major._obj.value = 3
minor._obj.value = 5
return orig_cuda.CUDA_SUCCESS
def cuDeviceCanAccessPeer(canAccessPeer, dev: int, peerDev: int) -> int:
canAccessPeer._obj.value = 1 # Always allow peer access in simulation
return orig_cuda.CUDA_SUCCESS
def cuCtxEnablePeerAccess(peerContext, flags: int) -> int:
return orig_cuda.CUDA_SUCCESS
def cuMemHostAlloc(pp, bytesize: int, flags: int) -> int:
return cuMemAlloc_v2(pp, bytesize)
def cuMemFreeHost(p: ctypes.c_void_p) -> int: return cuMemFree_v2(p)
def cuMemcpyDtoDAsync_v2(dst, src, bytesize: int, stream: Any) -> int:
ctypes.memmove(dst if isinstance(dst, int) else dst.value, src if isinstance(src, int) else src.value, bytesize)
return orig_cuda.CUDA_SUCCESS
def cuFuncSetAttribute(hfunc, attrib: int, value: int) -> int:
return orig_cuda.CUDA_SUCCESS
def cuStreamWaitEvent(stream: Any, event, flags: int) -> int: return orig_cuda.CUDA_SUCCESS
def cuCtxSynchronize() -> int: return orig_cuda.CUDA_SUCCESS
def cuGetErrorString(error: int, pStr) -> int:
error_str = orig_cuda.enum_cudaError_enum.get(error, "Unknown CUDA error").encode()
buf = ctypes.create_string_buffer(error_str)
# Set the pointer to point to our error string buffer
pStr._obj.value = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char))
return orig_cuda.CUDA_SUCCESS
def cuDeviceGetCount(count) -> int:
count._obj.value = 1
return orig_cuda.CUDA_SUCCESS

39
artifacts/package_sources/tinygrad/test/mockgpu/driver.py Normal file
View File

@@ -0,0 +1,39 @@
from typing import Any
from dataclasses import dataclass
class VirtFileDesc:
def __init__(self, fd): self.fd, self.off = fd, 0
def ioctl(self, fd, req, argp): raise NotImplementedError()
def mmap(self, st, sz, prot, flags, fd, off): raise NotImplementedError()
def close(self, fd): return 0
class TextFileDesc(VirtFileDesc):
def __init__(self, fd, text):
super().__init__(fd)
self.content = text
def ioctl(self, fd, req, argp): return 0
def read_contents(self, size=None):
ret = self.content[self.off:self.off+(size or len(self.content))]
self.off += (size or len(self.content))
return ret
def seek(self, offset): self.off += offset
class DirFileDesc(VirtFileDesc):
def __init__(self, fd, child_names):
super().__init__(fd)
self.child_names = child_names
def ioctl(self, fd, req, argp): return 0
def list_contents(self): return self.child_names
@dataclass(frozen=True)
class VirtFile:
path: str
fdcls: Any # TODO: fix this Union[VirtFileDesc, functools.partial[VirtFileDesc]]
class VirtDriver:
def __init__(self):
self.tracked_files = []
self.tracked_addresses = []
def track_address(self, staddr, enaddr, rcb, wcb): self.tracked_addresses.append((staddr, enaddr, rcb, wcb))
def open(self, name, flags, mode, fdcls): raise NotImplementedError()

4
artifacts/package_sources/tinygrad/test/mockgpu/gpu.py Normal file
View File

@@ -0,0 +1,4 @@
class VirtGPU:
def __init__(self, gpuid): self.gpuid = gpuid
def map_range(self, vaddr, size): raise NotImplementedError()
def unmap_range(self, vaddr, size): raise NotImplementedError()

20
artifacts/package_sources/tinygrad/test/mockgpu/helpers.py Normal file
View File

@@ -0,0 +1,20 @@
import ctypes
from tinygrad.runtime.support import c
gpuocelot_lib = c.DLL("ocelot", "gpuocelot")
@gpuocelot_lib.bind(None, ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int, ctypes.c_int,
ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int)
def ptx_run(source:bytes, n_args:int, args:c.POINTER[ctypes.c_void_p], blck_x:int, blck_y:int, blck_z:int,
grid_x:int, grid_y:int, grid_z:int, shared_mem_size:int): pass
class PythonRemu:
"""Python RDNA3/RDNA4 emulator wrapper used by mockgpu."""
valid_mem_ranges: set[tuple[int, int]] = set()
rsrc2: int = 0x19c # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
scratch_size: int = 0 # private_segment_fixed_size from kernel descriptor
arch: str = "rdna3" # Architecture: rdna3 or rdna4
user_data: list[int] = [] # All COMPUTE_USER_DATA registers (loaded into s[0:N])
def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
from test.mockgpu.amd.emu import run_asm
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size, self.arch, self.user_data)

110
artifacts/package_sources/tinygrad/test/mockgpu/mockgpu.py Normal file
View File

@@ -0,0 +1,110 @@
import ctypes, time, os, builtins, fcntl, typing
from tinygrad.helpers import DEV
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.autogen import libc
from test.mockgpu.nv.nvdriver import NVDriver
from test.mockgpu.amd.amddriver import AMDDriver
from test.mockgpu.am.amdriver import AMDriver, AMUSBDriver
start = time.perf_counter()
drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCK+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
"MOCK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
tracked_fds: dict[int, typing.Any] = {}
original_memoryview = builtins.memoryview
class TrackedMemoryView:
def __init__(self, data, rcb, wcb):
self.mv = original_memoryview(data)
self.rcb, self.wcb = rcb, wcb
def __getitem__(self, index):
self.rcb(self.mv, index)
return self.mv[index]
def __setitem__(self, index, value):
self.mv[index] = value
self.wcb(self.mv, index)
def cast(self, new_type, **kwargs):
self.mv = self.mv.cast('B').cast(new_type, **kwargs)
return self
@property
def nbytes(self): return self.mv.nbytes
def __len__(self): return len(self.mv)
def __repr__(self): return repr(self.mv)
def _memoryview(cls, mem):
if isinstance(mem, int) or isinstance(mem, ctypes.Array):
addr = ctypes.addressof(mem) if isinstance(mem, ctypes.Array) else mem
for d in drivers:
for st,en,rcb,wcb in d.tracked_addresses:
if st <= addr <= en: return TrackedMemoryView(mem, rcb, wcb)
return original_memoryview(mem)
class _MockMemoryviewMeta(type):
def __instancecheck__(cls, instance): return isinstance(instance, (original_memoryview, TrackedMemoryView))
builtins.memoryview = _MockMemoryviewMeta("memoryview", (), {'__new__': _memoryview}) # type: ignore
def _open(path, flags):
for d in drivers:
for x in d.tracked_files:
if path == x.path:
virtfd = d.open(path, flags, 0o777, x)
tracked_fds[virtfd.fd] = virtfd
return virtfd.fd
return os.open(path, flags, 0o777) if os.path.exists(path) else None
class MockFileIOInterface(FileIOInterface):
def __init__(self, path:str="", flags:int=os.O_RDONLY, fd:int|None=None):
self.path = path
self.fd = fd or _open(path, flags)
def __del__(self):
if self.fd in tracked_fds:
tracked_fds[self.fd].close(self.fd)
tracked_fds.pop(self.fd)
elif self.fd is not None: os.close(self.fd)
def ioctl(self, request, arg):
if self.fd in tracked_fds:
return tracked_fds[self.fd].ioctl(self.fd, request, ctypes.addressof(arg))
return fcntl.ioctl(self.fd, request, arg)
def mmap(self, start, sz, prot, flags, offset):
if self.fd in tracked_fds:
return tracked_fds[self.fd].mmap(start, sz, prot, flags, self.fd, offset)
return libc.mmap(start, sz, prot, flags, self.fd, offset)
def read(self, size=None, binary=False, offset=None):
if self.fd in tracked_fds:
if offset is not None: tracked_fds[self.fd].seek(offset)
return tracked_fds[self.fd].read_contents(size)
if binary: raise NotImplementedError()
with open(self.fd, "rb" if binary else "r", closefd=False) as file:
if file.tell() >= os.fstat(self.fd).st_size: file.seek(0)
return file.read(size)
def listdir(self):
if self.fd in tracked_fds:
return tracked_fds[self.fd].list_contents()
return os.listdir(self.path)
def write(self, content, binary=False, offset=None):
if self.fd in tracked_fds:
if offset is not None: tracked_fds[self.fd].seek(offset)
return tracked_fds[self.fd].write_contents(content)
raise NotImplementedError()
def seek(self, offset):
if self.fd in tracked_fds:
tracked_fds[self.fd].seek(offset)
else:
os.lseek(self.fd, offset, os.SEEK_CUR)
@staticmethod
def anon_mmap(start, sz, prot, flags, offset):
return FileIOInterface._mmap(start, sz, prot, flags & ~0x4a000, -1, offset) # strip MAP_LOCKED|MAP_POPULATE|MAP_HUGETLB
@staticmethod
def exists(path): return _open(path, os.O_RDONLY) is not None
@staticmethod
def readlink(path): raise NotImplementedError()
@staticmethod
def eventfd(initval, flags=None): NotImplementedError()

294
artifacts/package_sources/tinygrad/test/mockgpu/nv/nvdriver.py Normal file
View File

@@ -0,0 +1,294 @@
import ctypes, mmap, collections, functools, os
from tinygrad.runtime.autogen import nv_570 as nv_gpu, libc
from typing import cast, Any
from tinygrad.helpers import to_mv
from test.mockgpu.driver import VirtDriver, VirtFileDesc, VirtFile
from test.mockgpu.nv.nvgpu import NVGPU
NVSubDevice = collections.namedtuple('NVSubDevice', ['device'])
NVUserMode = collections.namedtuple('NVUserMode', ['subdevice'])
NVVASpace = collections.namedtuple('NVVASpace', ['device'])
NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size', 'is_signal'])
NVChannelGroup = collections.namedtuple('NVChannelGroup', ['device'])
NVContextShare = collections.namedtuple('NVContextShare', ['channel_group'])
NVGPFIFO = collections.namedtuple('NVGPFIFO', ['device', 'token'])
NVProfiler = collections.namedtuple('NVProfiler', ['subdevice'])
class NVCtlFileDesc(VirtFileDesc):
def __init__(self, fd, driver):
super().__init__(fd)
self.driver = driver
def ioctl(self, fd, request, argp): return self.driver.ctl_ioctl(request, argp)
def mmap(self, start, sz, prot, flags, fd, offset): return libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
class NVUVMFileDesc(VirtFileDesc):
def __init__(self, fd, driver):
super().__init__(fd)
self.driver = driver
def ioctl(self, fd, request, argp): return self.driver.uvm_ioctl(request, argp)
def mmap(self, start, sz, prot, flags, fd, offset): return libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
class NVDevFileDesc(VirtFileDesc):
def __init__(self, fd, driver, gpu):
super().__init__(fd)
self.driver, self.gpu = driver, gpu
self._mapping_userland = False
self._mapping_signal = False
def ioctl(self, fd, request, argp): return self.driver.dev_ioctl(self.gpu, request, argp)
def mmap(self, start, sz, prot, flags, fd, offset):
start = libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
if self._mapping_userland or self._mapping_signal:
self.driver.track_address(start, start+sz, lambda mv,off: None, lambda mv, off: self.driver._gpu_mmio_write(mv, off, self.gpu))
self._mapping_signal = False
return start
class NVDriver(VirtDriver):
def __init__(self, gpus=6):
super().__init__()
self.tracked_files += [VirtFile('/dev/nvidiactl', functools.partial(NVCtlFileDesc, driver=self)),
VirtFile('/dev/nvidia-uvm', functools.partial(NVUVMFileDesc, driver=self))]
self.root_handle = None
self.gpus = {}
self.next_fd = (1 << 29)
self.next_handle = 1
self.object_by_handle = {}
self.opened_fds = {}
self.next_doorbell = collections.defaultdict(int)
self._executing = False # re-entrancy guard for _gpu_mmio_write
for i in range(gpus): self._prepare_gpu(i)
def _alloc_fd(self):
my_fd = self.next_fd
self.next_fd = self.next_fd + 1
return my_fd
def _alloc_handle(self):
handle = self.next_handle
self.next_handle += 1
return handle
def _prepare_gpu(self, gpu_id):
self.gpus[gpu_id] = NVGPU(gpu_id)
self.tracked_files += [VirtFile(f'/dev/nvidia{gpu_id}', functools.partial(NVDevFileDesc, driver=self, gpu=self.gpus[gpu_id]))]
def open(self, name, flags, mode, virtfile):
cl = virtfile.fdcls(self._alloc_fd())
self.opened_fds[cl.fd] = cl
return cl
def rm_alloc(self, argp):
struct = nv_gpu.NVOS21_PARAMETERS.from_address(argp)
params_ptr = cast(int, struct.pAllocParms)
if struct.hClass == nv_gpu.NV01_ROOT_CLIENT: self.root_handle = struct.hObjectNew = self._alloc_handle()
elif struct.hClass == nv_gpu.NV01_DEVICE_0:
params:Any = nv_gpu.NV0080_ALLOC_PARAMETERS.from_address(params_ptr)
assert params.hClientShare == self.root_handle
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = self.gpus[params.deviceId]
elif struct.hClass == nv_gpu.NV20_SUBDEVICE_0:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVSubDevice(self.object_by_handle[struct.hObjectParent])
elif struct.hClass == nv_gpu.NV01_MEMORY_VIRTUAL:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
struct.hObjectNew = self._alloc_handle()
elif struct.hClass == nv_gpu.TURING_USERMODE_A:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVSubDevice)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVUserMode(self.object_by_handle[struct.hObjectParent])
elif struct.hClass == nv_gpu.FERMI_VASPACE_A:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVVASpace(self.object_by_handle[struct.hObjectParent])
elif struct.hClass == nv_gpu.NV1_MEMORY_SYSTEM or struct.hClass == nv_gpu.NV1_MEMORY_USER:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
params = nv_gpu.NV_MEMORY_ALLOCATION_PARAMS.from_address(params_ptr)
struct.hObjectNew = self._alloc_handle()
is_signal = struct.hClass == nv_gpu.NV1_MEMORY_SYSTEM # signal memory uses NV1_MEMORY_SYSTEM (uncached)
self.object_by_handle[struct.hObjectNew] = NVAllocation(self.object_by_handle[struct.hObjectParent], params.size, is_signal)
elif struct.hClass == nv_gpu.KEPLER_CHANNEL_GROUP_A:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVChannelGroup(self.object_by_handle[struct.hObjectParent])
elif struct.hClass == nv_gpu.FERMI_CONTEXT_SHARE_A:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVChannelGroup)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVContextShare(self.object_by_handle[struct.hObjectParent])
elif struct.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
parent = self.object_by_handle.get(struct.hObjectParent)
assert parent is not None and isinstance(parent, (NVChannelGroup, NVGPU))
struct.hObjectNew = self._alloc_handle()
params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS.from_address(params_ptr)
gpu = parent.device if isinstance(parent, NVChannelGroup) else parent
gpfifo_token = gpu.add_gpfifo(params.gpFifoOffset, params.gpFifoEntries)
self.object_by_handle[struct.hObjectNew] = NVGPFIFO(gpu, gpfifo_token)
elif struct.hClass in (nv_gpu.AMPERE_DMA_COPY_B, nv_gpu.ADA_COMPUTE_A, nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER):
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPFIFO)
struct.hObjectNew = self._alloc_handle()
gpfifo = self.object_by_handle[struct.hObjectParent]
gpfifo.device.queues[gpfifo.token].bound_engines.add(struct.hClass)
elif struct.hClass == nv_gpu.GT200_DEBUGGER:
struct.hObjectNew = self._alloc_handle()
elif struct.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVSubDevice)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVProfiler(self.object_by_handle[struct.hObjectParent])
else: raise RuntimeError(f"Unknown {struct.hClass} to rm_alloc")
return 0
def rm_control(self, argp):
struct = nv_gpu.NVOS54_PARAMETERS.from_address(argp)
params_ptr = cast(int, struct.params)
if struct.cmd == nv_gpu.NV0000_CTRL_CMD_GPU_GET_ID_INFO_V2:
params:Any = nv_gpu.NV0000_CTRL_GPU_GET_ID_INFO_V2_PARAMS.from_address(params_ptr)
params.deviceInstance = params.gpuId # emulate them to be the same
elif struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST_V2 or struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:
if struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:
params = nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS.from_address(params_ptr)
else:
params = nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_V2_PARAMS.from_address(params_ptr)
classes = [50021, 51607, 51648, 50543, 51125, 51125, 51125, 51125, 50529, 36967, 36909, 37105, 33868, 36978, 37095, 37094, 36980, 37014, 49270,
41068, 41088, 41280, 50025, 96, 112, 115, 125, 20608, 20640, 20539, 20540, 41089, 41092, 50034, 50810, 50811, 50814, 51056, 51057,
51059, 51069, 51071, 51632, 51639, 51639, 51706, 52019, 222, 50287, 50273, 50031, 50017] # from ada102
params.numClasses = len(classes)
if struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:
if params.classList and params.numClasses > 0:
clslist = to_mv(params.classList, params.numClasses * 4).cast('I')
for i,c in enumerate(classes): clslist[i] = c
else: params.numClasses = len(classes)
else:
for i,c in enumerate(classes): params.classList[i] = c
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GR_GET_INFO:
info = {nv_gpu.NV2080_CTRL_GR_INFO_INDEX_SM_VERSION: nv_gpu.NV2080_CTRL_GR_INFO_SM_VERSION_3_5,
nv_gpu.NV2080_CTRL_GR_INFO_INDEX_LITTER_NUM_GPCS: 1,
nv_gpu.NV2080_CTRL_GR_INFO_INDEX_LITTER_NUM_TPC_PER_GPC: 1,
nv_gpu.NV2080_CTRL_GR_INFO_INDEX_LITTER_NUM_SM_PER_TPC: 1,
nv_gpu.NV2080_CTRL_GR_INFO_INDEX_MAX_WARPS_PER_SM: 1,
}
params = nv_gpu.NV2080_CTRL_GR_GET_INFO_PARAMS.from_address(params_ptr)
reqlist = (nv_gpu.NV2080_CTRL_GR_INFO * params.grInfoListSize).from_address(params.grInfoList)
for i in range(params.grInfoListSize): reqlist[i].data = info[reqlist[i].index]
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GPU_GET_GID_INFO:
assert struct.hObject in self.object_by_handle and isinstance(self.object_by_handle[struct.hObject], NVSubDevice)
gpu = self.object_by_handle[struct.hObject].device
params = nv_gpu.NV2080_CTRL_GPU_GET_GID_INFO_PARAMS.from_address(params_ptr)
if params.flags != nv_gpu.NV2080_GPU_CMD_GPU_GET_GID_FLAGS_FORMAT_BINARY: raise RuntimeError("Unknown format")
bts = gpu.gpu_uuid(sz=params.length)
for i in range(params.length): params.data[i] = bts[i]
elif struct.cmd == nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN:
assert struct.hObject in self.object_by_handle and isinstance(self.object_by_handle[struct.hObject], NVGPFIFO)
params = nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS.from_address(params_ptr)
gpu_fifo = self.object_by_handle[struct.hObject]
params.workSubmitToken = gpu_fifo.token
elif struct.cmd in (nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06F_CTRL_CMD_BIND, nv_gpu.NVA06F_CTRL_CMD_GPFIFO_SCHEDULE): pass
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_PERF_BOOST: pass
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_FB_FLUSH_GPU_CACHE: pass
elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_ALL_SM_ERROR_STATES:
params = nv_gpu.NV83DE_CTRL_DEBUG_READ_ALL_SM_ERROR_STATES_PARAMS.from_address(params_ptr)
params.mmuFault.valid = bool("MOCKGPU_EMU_FAULTADDR" in os.environ)
elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_MMU_FAULT_INFO:
params = nv_gpu.struct_NV83DE_CTRL_DEBUG_READ_MMU_FAULT_INFO_PARAMS.from_address(params_ptr)
params.count = 1
params.mmuFaultInfoList[0].faultAddress = int(os.environ['MOCKGPU_EMU_FAULTADDR'], base=16)
params.mmuFaultInfoList[0].faultType = 1
params.mmuFaultInfoList[0].accessType = 1
elif struct.cmd == nv_gpu.NV0000_CTRL_CMD_SYSTEM_GET_BUILD_VERSION_V2:
params = nv_gpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_V2_PARAMS.from_address(params_ptr)
params.driverVersionBuffer = b"570.00.00\0"
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GR_GET_TPC_MASK:
params = nv_gpu.NV2080_CTRL_GR_GET_TPC_MASK_PARAMS.from_address(params_ptr)
params.tpcMask = 0x1 # one TPC
# Profiler commands - just pass through for mockgpu
elif struct.cmd in (nv_gpu.NVB0CC_CTRL_CMD_POWER_REQUEST_FEATURES, nv_gpu.NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM,
nv_gpu.NVB0CC_CTRL_CMD_RESERVE_HWPM_LEGACY, nv_gpu.NVB0CC_CTRL_CMD_RESERVE_PM_AREA_PC_SAMPLER,
nv_gpu.NVB0CC_CTRL_CMD_BIND_PM_RESOURCES, nv_gpu.NVB0CC_CTRL_CMD_SET_HS_CREDITS,
nv_gpu.NVB0CC_CTRL_CMD_EXEC_REG_OPS, nv_gpu.NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT): pass
else: raise RuntimeError(f"Unknown {struct.cmd} to rm_control")
return 0
def ctl_ioctl(self, req, argp):
nr = req & 0xff
if nr == nv_gpu.NV_ESC_RM_ALLOC: return self.rm_alloc(argp)
elif nr == nv_gpu.NV_ESC_RM_CONTROL: return self.rm_control(argp)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
st:Any = nv_gpu.nv_ioctl_nvos33_parameters_with_fd.from_address(argp)
obj = self.object_by_handle.get(st.params.hMemory)
file = self.opened_fds.get(st.fd)
if isinstance(obj, NVUserMode) and isinstance(file, NVDevFileDesc):
file._mapping_userland = True
elif isinstance(obj, NVAllocation) and obj.is_signal and isinstance(file, NVDevFileDesc):
file._mapping_signal = True
elif nr == nv_gpu.NV_ESC_RM_FREE:
st = nv_gpu.NVOS00_PARAMETERS.from_address(argp)
self.object_by_handle.pop(st.hObjectOld)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY_DMA:
pass # mappings are same as uvm
elif nr == nv_gpu.NV_ESC_CARD_INFO:
for i,gpu in enumerate(self.gpus.values()):
st = nv_gpu.nv_ioctl_card_info_t.from_address(argp + i * ctypes.sizeof(nv_gpu.nv_ioctl_card_info_t))
st.gpu_id = gpu.gpuid
st.pci_info.device_id = 0x2684
st.valid = True
else: raise RuntimeError(f"Unknown {nr} to nvidiactl")
return 0
def uvm_ioctl(self, nr, argp):
if nr == nv_gpu.UVM_INITIALIZE: pass
elif nr == nv_gpu.UVM_MM_INITIALIZE: pass
elif nr == nv_gpu.UVM_REGISTER_GPU:
st:Any = nv_gpu.UVM_REGISTER_GPU_PARAMS.from_address(argp)
assert any(all(st.gpu_uuid.uuid[i] == gpu.gpu_uuid()[i] for i in range(16)) for gpu in self.gpus.values())
elif nr == nv_gpu.UVM_REGISTER_GPU_VASPACE: pass
elif nr == nv_gpu.UVM_ENABLE_PEER_ACCESS: pass # uvm and shared spaced are setup already, no emulation for now
elif nr == nv_gpu.UVM_CREATE_EXTERNAL_RANGE:
st = nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS.from_address(argp)
libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
elif nr == nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION:
st = nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS.from_address(argp)
for gpu_attr_id in range(st.gpuAttributesCount):
gpu = None
for _gpu in self.gpus.values():
if all(st.perGpuAttributes[gpu_attr_id].gpuUuid.uuid[i] == _gpu.gpu_uuid()[i] for i in range(16)):
gpu = _gpu
break
if gpu is None: return -1
gpu.map_range(st.base, st.length)
elif nr == nv_gpu.UVM_REGISTER_CHANNEL: pass
elif nr == nv_gpu.UVM_FREE:
st = nv_gpu.UVM_FREE_PARAMS.from_address(argp)
libc.munmap(st.base, st.length)
else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
return 0
def dev_ioctl(self, dev, req, argp):
nr = req & 0xff
# Handle NV_ESC_RM_ALLOC_MEMORY for host/signal memory
if nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
# Track host memory (signal memory) - progress queues when written to
if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
return 0
def _gpu_mmio_write(self, mv, off, gpu):
if self._executing: return # prevent re-entrancy
self._executing = True
try:
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.ctrl.GPGet != q.ctrl.GPPut:
any_progress |= q.execute()
finally:
self._executing = False

220
artifacts/package_sources/tinygrad/test/mockgpu/nv/nvgpu.py Normal file
View File

@@ -0,0 +1,220 @@
import ctypes, time
from tinygrad.runtime.autogen import nv_570 as nv_gpu
from enum import Enum, auto
from test.mockgpu.gpu import VirtGPU
from test.mockgpu.helpers import ptx_run
from tinygrad.helpers import to_mv
from tinygrad.runtime.support.c import init_c_struct_t
def make_qmd_struct_type():
fields = []
bits = [(name,dt) for name,dt in nv_gpu.__dict__.items() if name.startswith("NVC6C0_QMDV03_00") and isinstance(dt, tuple)]
bits += [(name+f"_{i}",dt(i)) for name,dt in nv_gpu.__dict__.items() for i in range(8) if name.startswith("NVC6C0_QMDV03_00") and callable(dt)]
bits = sorted(bits, key=lambda x: x[1][1])
for i,(name, data) in enumerate(bits):
fields.append((name.replace("NVC6C0_QMDV03_00_", "").lower(), ctypes.c_uint32, data[1]//8, data[0]-data[1]+1, data[1]%8))
return init_c_struct_t(0x40 * 4, tuple(fields))
qmd_struct_t = make_qmd_struct_type()
class SchedResult(Enum): CONT = auto(); YIELD = auto() # noqa: E702
class GPFIFO:
def __init__(self, token, base, entries_cnt):
self.token, self.base, self.entries_cnt = token, base, entries_cnt
self.gpfifo = to_mv(self.base, self.entries_cnt * 8).cast("Q")
self.ctrl = nv_gpu.AmpereAControlGPFifo.from_address(self.base + self.entries_cnt * 8)
self.state = {}
self.bound_engines: set[int] = set()
# Buf exec state
self.buf = None
self.buf_sz = 0
self.buf_ptr = 0
def _next_dword(self):
assert self.buf is not None
x = self.buf[self.buf_ptr]
self.buf_ptr += 1
return x
def _next_header(self):
header = self._next_dword()
typ = (header >> 28) & 0b111
size = (header >> 16) & 0xFFF
subc = (header >> 13) & 0x7
mthd = (header & 0x1FFF) << 2
return typ, size, subc, mthd
def _state(self, reg): return self.state[reg]
def _state64(self, reg): return (self.state[reg] << 32) + self.state[reg + 4]
def _state64_le(self, reg): return (self.state[reg + 4] << 32) + self.state[reg]
def _reset_buf_state(self): self.buf, self.buf_ptr = None, 0
def _set_buf_state(self, gpfifo_entry):
ptr = ((gpfifo_entry >> 2) & 0x3fffffffff) << 2
sz = ((gpfifo_entry >> 42) & 0x1fffff) << 2
self.buf = to_mv(ptr, sz).cast("I")
self.buf_sz = sz // 4
def execute(self) -> bool:
initial_off = self.buf_ptr
while self.ctrl.GPGet != self.ctrl.GPPut:
self._set_buf_state(self.gpfifo[self.ctrl.GPGet])
if not self.execute_buf():
# Buffer isn't executed fully, check if any progress and report.
# Do not move GPGet in this case, will continue from the same state next time.
return self.buf_ptr != initial_off
self.ctrl.GPGet = (self.ctrl.GPGet + 1) % self.entries_cnt
self._reset_buf_state()
return True
def execute_buf(self) -> bool:
while self.buf_ptr < self.buf_sz:
init_off = self.buf_ptr
_, size, _, mthd = self._next_header()
cmd_end_off = self.buf_ptr + size
while self.buf_ptr < cmd_end_off:
res = self.execute_cmd(mthd)
if res == SchedResult.YIELD:
self.buf_ptr = init_off # just revert to the header
return False
mthd += 4
return True
def execute_qmd(self, qmd_addr):
qmd = qmd_struct_t.from_address(qmd_addr)
prg_addr = qmd.program_address_lower + (qmd.program_address_upper << 32)
const0 = to_mv(qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32), 0x160).cast('I')
args_cnt, vals_cnt = const0[80], const0[81]
args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160
args = to_mv(args_addr, args_cnt*8).cast('Q')
vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q')
cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2
try: ptx_run(ctypes.cast(prg_addr, ctypes.c_char_p), args_cnt+vals_cnt, (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
except Exception as e: print("failed to execute:", e)
if qmd.release0_enable:
rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)
rel0[1] = int(time.perf_counter() * 1e9)
if qmd.release1_enable:
rel1 = to_mv(qmd.release1_address_lower + (qmd.release1_address_upper << 32), 0x10).cast('Q')
rel1[0] = qmd.release1_payload_lower + (qmd.release1_payload_upper << 32)
rel1[1] = int(time.perf_counter() * 1e9)
if qmd.dependent_qmd0_enable:
if qmd.dependent_qmd0_action == 1: self.execute_qmd(qmd.dependent_qmd0_pointer << 8)
else: raise RuntimeError("unsupported dependent qmd action")
def execute_cmd(self, cmd) -> SchedResult:
if cmd == nv_gpu.NVC56F_SEM_EXECUTE: return self._exec_signal()
elif cmd == nv_gpu.NVC6C0_LAUNCH_DMA: return self._exec_nvc6c0_dma()
elif cmd == nv_gpu.NVC6B5_LAUNCH_DMA: # NOTE: NVC6B5_LAUNCH_DMA == NVC9B0_EXECUTE == 0x300
return self._exec_vid_decode() if self.bound_engines & {nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER} else self._exec_nvc6b5_dma()
elif cmd == nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B: return self._exec_pcas2()
elif cmd == 0x0320: return self._exec_load_inline_qmd() # NVC6C0_LOAD_INLINE_QMD_DATA
elif cmd == nv_gpu.NVC9B0_SEMAPHORE_D: return self._exec_vid_semaphore()
else: self.state[cmd] = self._next_dword() # just state update
return SchedResult.CONT
def _exec_signal(self) -> SchedResult:
signal = self._state64_le(nv_gpu.NVC56F_SEM_ADDR_LO)
val = self._state64_le(nv_gpu.NVC56F_SEM_PAYLOAD_LO)
flags = self._next_dword()
typ = (flags >> 0) & 0b111
timestamp = (flags & (1 << 25)) == (1 << 25)
if typ == 1:
to_mv(signal, 8).cast('Q')[0] = val
if timestamp: to_mv(signal + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
elif typ == 3:
mval = to_mv(signal, 8).cast('Q')[0]
return SchedResult.CONT if mval >= val else SchedResult.YIELD
elif typ == 4: # ACQ_AND: (mem & payload) != 0
mval = to_mv(signal, 4).cast('I')[0]
return SchedResult.CONT if (mval & (val & 0xffffffff)) != 0 else SchedResult.YIELD
elif typ == 5: # ACQ_NOR: ~(mem | payload) != 0
mval = to_mv(signal, 4).cast('I')[0]
return SchedResult.CONT if (~(mval | (val & 0xffffffff)) & 0xffffffff) != 0 else SchedResult.YIELD
else: raise RuntimeError(f"Unsupported type={typ} in exec wait/signal")
return SchedResult.CONT
def _exec_vid_decode(self) -> SchedResult:
self._next_dword() # consume execute flags
# validate that all required decode state was set up correctly
assert self._state(nv_gpu.NVC9B0_SET_APPLICATION_ID) == nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC
pic_desc_addr = self._state(nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET) << 8
pic = nv_gpu.nvdec_hevc_pic_s.from_address(pic_desc_addr)
assert pic.stream_len > 0 and pic.pic_width_in_luma_samples > 0 and pic.pic_height_in_luma_samples > 0
assert self._state(nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET) != 0
assert self._state(nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET) != 0
assert self._state(nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET) != 0
assert self._state(nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET) != 0
return SchedResult.CONT
def _exec_vid_semaphore(self) -> SchedResult:
signal = self._state64(nv_gpu.NVC9B0_SEMAPHORE_A)
val = self._state(nv_gpu.NVC9B0_SEMAPHORE_C)
self._next_dword() # flags
to_mv(signal, 8).cast('Q')[0] = val
to_mv(signal + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
return SchedResult.CONT
def _exec_load_inline_qmd(self):
qmd_addr = self._state64(nv_gpu.NVC6C0_SET_INLINE_QMD_ADDRESS_A) << 8
assert qmd_addr != 0x0, f"invalid qmd address {qmd_addr}"
qmd_data = [self._next_dword() for _ in range(0x40)]
cdata = (ctypes.c_uint32 * len(qmd_data))(*qmd_data)
ctypes.memmove(qmd_addr, cdata, 0x40 * 4)
self.execute_qmd(qmd_addr)
def _exec_nvc6c0_dma(self):
addr = self._state64(nv_gpu.NVC6C0_OFFSET_OUT_UPPER)
sz = self._state(nv_gpu.NVC6C0_LINE_LENGTH_IN)
lanes = self._state(nv_gpu.NVC6C0_LINE_COUNT)
assert lanes == 1, f"unsupported lanes > 1 in _exec_nvc6c0_dma: {lanes}"
flags = self._next_dword()
assert flags == 0x41, f"unsupported flags in _exec_nvc6c0_dma: {flags}"
typ, dsize, _, mthd = self._next_header()
assert typ == 6 and mthd == nv_gpu.NVC6C0_LOAD_INLINE_DATA, f"Expected inline data not found after nvc6c0_dma, {typ=} {mthd=}"
copy_data = [self._next_dword() for _ in range(dsize)]
assert len(copy_data) * 4 == sz, f"different copy sizes in _exec_nvc6c0_dma: {len(copy_data) * 4} != {sz}"
cdata = (ctypes.c_uint32 * len(copy_data))(*copy_data)
ctypes.memmove(addr, cdata, sz)
def _exec_nvc6b5_dma(self):
flags = self._next_dword()
if (flags & 0b11) != 0:
src = self._state64(nv_gpu.NVC6B5_OFFSET_IN_UPPER)
dst = self._state64(nv_gpu.NVC6B5_OFFSET_OUT_UPPER)
sz = self._state(nv_gpu.NVC6B5_LINE_LENGTH_IN)
assert flags == 0x182, f"unsupported flags in _exec_nvc6b5_dma: {flags}"
ctypes.memmove(dst, src, sz)
elif ((flags >> 3) & 0b11) != 0:
src = to_mv(self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 0x10).cast('Q')
val = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
src[0] = val
src[1] = int(time.perf_counter() * 1e9)
else: raise RuntimeError("unknown nvc6b5_dma flags")
def _exec_pcas2(self):
qmd_addr = self._state(nv_gpu.NVC6C0_SEND_PCAS_A) << 8
typ = self._next_dword()
if typ == 2 or typ == 9: # schedule
self.execute_qmd(qmd_addr)
class NVGPU(VirtGPU):
def __init__(self, gpuid):
super().__init__(gpuid)
self.regs = {}
self.mapped_ranges = set()
self.queues = []
def map_range(self, vaddr, size): self.mapped_ranges.add((vaddr, size))
def unmap_range(self, vaddr, size): self.mapped_ranges.remove((vaddr, size))
def add_gpfifo(self, base, entries_count):
self.queues.append(GPFIFO(token:=len(self.queues), base, entries_count))
return token
def gpu_uuid(self, sz=16): return self.gpuid.to_bytes(sz, byteorder='big', signed=False)

212
artifacts/package_sources/tinygrad/test/mockgpu/usb.py Normal file
View File

@@ -0,0 +1,212 @@
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('<I', gpu_cfg, reg_off, type_bits)
struct.pack_into('<I', gpu_cfg, 0x100, 0x15 | (1 << 16)) # REBAR cap header: id=0x15, version=1, next=0
struct.pack_into('<I', gpu_cfg, 0x104, sum(1 << (i + 4) for i in range(10))) # supported sizes up to 512MB
def _get_cfg(self, bus:int, dev:int, fn:int) -> 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('<I', cfg, reg_off, (~(bar_size - 1)) & 0xFFFFFFF0 | type_bits)
else:
struct.pack_into('<I', cfg, reg_off, val)
hi = struct.unpack_from('<I', cfg, reg_off + 4)[0] if is_64 else 0
self._bar_addrs[reg_off] = ((hi << 32) | (val & ~0xF), bar_size)
return
# Check if upper 32 bits of a 64-bit BAR
for breg, (bsz, _, b64) in self._gpu_bars.items():
if b64 and reg_off == breg + 4:
struct.pack_into('<I', cfg, reg_off, 0xFFFFFFFF if val == 0xFFFFFFFF else val)
if val != 0xFFFFFFFF:
self._bar_addrs[breg] = ((val << 32) | (struct.unpack_from('<I', cfg, breg)[0] & ~0xF), bsz)
return
# Generic config write
for i in range(size): cfg[byte_addr + i] = (val >> (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, (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('<III', data)
address, fmt_type, byte_en = address_lo | (address_hi << 32), value & 0xFF, value >> 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('<I', payload & 0xFFFFFFFF) + bytes(4)
else: raise ValueError(f"unsupported control OUT request 0x{request:02X}")
def control_read(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> 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":
host_addr, region_size = self.state._dma_regions[address]
ctypes.memmove(host_addr, data, min(len(data), region_size))
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_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":
host_addr, region_size = self.state._dma_regions[address]
data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
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)