IQ.Pilot Release Commit @ b6534c0
This commit is contained in:
IQ.Lvbs CI [bot]
127
artifacts/package_sources/tinygrad/test/mockgpu/am/amdriver.py
Normal file
127
artifacts/package_sources/tinygrad/test/mockgpu/am/amdriver.py
Normal 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
310
artifacts/package_sources/tinygrad/test/mockgpu/am/amgpu.py
Normal 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
38
artifacts/package_sources/tinygrad/test/mockgpu/amd/README
Normal 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
191
artifacts/package_sources/tinygrad/test/mockgpu/amd/amddriver.py
Normal 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
448
artifacts/package_sources/tinygrad/test/mockgpu/amd/amdgpu.py
Normal 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
2325
artifacts/package_sources/tinygrad/test/mockgpu/amd/emu.py
Normal file
File diff suppressed because it is too large
Load Diff
1362
artifacts/package_sources/tinygrad/test/mockgpu/amd/pcode.py
Normal file
1362
artifacts/package_sources/tinygrad/test/mockgpu/amd/pcode.py
Normal file
File diff suppressed because it is too large
Load Diff
174
artifacts/package_sources/tinygrad/test/mockgpu/cuda/cuda.py
Normal file
174
artifacts/package_sources/tinygrad/test/mockgpu/cuda/cuda.py
Normal 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
39
artifacts/package_sources/tinygrad/test/mockgpu/driver.py
Normal 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
4
artifacts/package_sources/tinygrad/test/mockgpu/gpu.py
Normal 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
20
artifacts/package_sources/tinygrad/test/mockgpu/helpers.py
Normal 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
110
artifacts/package_sources/tinygrad/test/mockgpu/mockgpu.py
Normal 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
294
artifacts/package_sources/tinygrad/test/mockgpu/nv/nvdriver.py
Normal 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
220
artifacts/package_sources/tinygrad/test/mockgpu/nv/nvgpu.py
Normal 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
212
artifacts/package_sources/tinygrad/test/mockgpu/usb.py
Normal 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)
|
||||
Reference in New Issue
Block a user