IQ.Lvbs/IQ.Pilot
8
7
Fork 3
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

IQ.Pilot Release Commit @ b6534c0

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-08-27 20:17:33 -05:00
commit 00f07cac48
4706 changed files with 1257146 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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