diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index f9decbda3d..27c7e16830 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU from tinygrad.uop.ops import sint from tinygrad.device import BufferSpec, CPUProgram -from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, DEBUG, prod, OSX, to_mv, hi32, lo32 +from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32 from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import NVRenderer from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler @@ -24,28 +24,6 @@ def nv_iowr(fd:FileIOInterface, nr, args): ret = fd.ioctl((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF), args) if ret != 0: raise RuntimeError(f"ioctl returned {ret}") -def rm_alloc(fd, clss, root, parant, params): - made = nv_gpu.NVOS21_PARAMETERS(hRoot=root, hObjectParent=parant, hClass=clss, - pAllocParms=ctypes.cast(ctypes.byref(params), ctypes.c_void_p) if params is not None else None) - nv_iowr(fd, nv_gpu.NV_ESC_RM_ALLOC, made) - if made.status != 0: - if made.status == nv_gpu.NV_ERR_NO_MEMORY: raise MemoryError(f"rm_alloc returned {get_error_str(made.status)}") - raise RuntimeError(f"rm_alloc returned {get_error_str(made.status)}") - return made - -def rm_control(cmd, sttyp, fd, client, obj, **kwargs): - made = nv_gpu.NVOS54_PARAMETERS(hClient=client, hObject=obj, cmd=cmd, paramsSize=ctypes.sizeof(params:=sttyp(**kwargs)), - params=ctypes.cast(ctypes.byref(params), ctypes.c_void_p) if params is not None else None) - nv_iowr(fd, nv_gpu.NV_ESC_RM_CONTROL, made) - if made.status != 0: raise RuntimeError(f"rm_control returned {get_error_str(made.status)}") - return params - -def make_rmctrl_type(): - return type("NVRMCTRL", (object,), {name[name.find("_CTRL_CMD_")+10:].lower(): functools.partial(rm_control, dt, sttyp) - for name,dt in nv_gpu.__dict__.items() if name.find("_CTRL_CMD_")>=0 and (sttyp:=getattr(nv_gpu, name.replace("_CTRL_CMD_", "_CTRL_")+"_PARAMS", \ - getattr(nv_gpu, name+"_PARAMS", getattr(nv_gpu, name.replace("_CTRL_CMD_", "_CTRL_DEBUG_")+"_PARAMETERS", None))))}) -rmctrl = make_rmctrl_type() - def uvm_ioctl(cmd, sttyp, fd:FileIOInterface, **kwargs): ret = fd.ioctl(cmd, made:=sttyp(**kwargs)) if ret != 0: raise RuntimeError(f"ioctl(uvm) returned {ret}") @@ -61,7 +39,7 @@ class QMD: fields: dict[str, dict[str, tuple[int, int]]] = {} def __init__(self, dev:NVDevice, addr:int|None=None, **kwargs): - self.ver, self.sz = (5, 0x60) if dev.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else (3, 0x40) + self.ver, self.sz = (5, 0x60) if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else (3, 0x40) # Init fields from module if (pref:="NVCEC0_QMDV05_00" if self.ver == 5 else "NVC6C0_QMDV03_00") not in QMD.fields: @@ -81,7 +59,7 @@ class QMD: self.mv[lo//8:hi//8+1] = int((num & ~mask) | ((value << (lo % 8)) & mask)).to_bytes((hi//8 - lo//8 + 1), "little") def write(self, **kwargs): - for k,val in kwargs.items(): self._rw_bits(*QMD.fields[self.pref][k.upper()], value=val) + for k,val in kwargs.items(): self._rw_bits(*QMD.fields[self.pref][k.upper()], value=val) # type: ignore [misc] def read(self, k, val=0): return self._rw_bits(*QMD.fields[self.pref][k.upper()]) @@ -252,7 +230,7 @@ class NVProgram(HCQProgram): self.constbuffer_0 = [0] * (cbuf0_size // 4) - if dev.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A: + if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A: self.constbuffer_0[188:192], self.constbuffer_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0 qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'register_count':self.regs_usage, 'program_address_upper_shifted4':hi32(self.prog_addr>>4), 'program_address_lower_shifted4':lo32(self.prog_addr>>4), @@ -298,16 +276,15 @@ class NVProgram(HCQProgram): class NVAllocator(HCQAllocator['NVDevice']): def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer: - if options.host: return self.dev._gpu_alloc(size, host=True, tag="user host memory") - return self.dev._gpu_alloc(size, cpu_access=options.cpu_access, tag=f"user memory ({options})") + return self.dev.iface._gpu_alloc(size, cpu_access=options.cpu_access, host=options.host) def _free(self, opaque:HCQBuffer, options:BufferSpec): try: self.dev.synchronize() - self.dev._gpu_free(opaque) + self.dev.iface._gpu_free(opaque) except AttributeError: pass - def map(self, buf:HCQBuffer): self.dev._gpu_map(buf._base if buf._base is not None else buf) + def map(self, buf:HCQBuffer): self.dev.iface._gpu_map(buf._base if buf._base is not None else buf) @dataclass class GPFifo: @@ -318,11 +295,7 @@ class GPFifo: put_value: int = 0 MAP_FIXED, MAP_NORESERVE = 0x10, 0x400 -class NVDevice(HCQCompiled[NVSignal]): - devices: ClassVar[list[HCQCompiled]] = [] - signal_pages: ClassVar[list[HCQBuffer]] = [] - signal_pool: ClassVar[list[HCQBuffer]] = [] - +class NVKIface: root = None fd_ctl: FileIOInterface fd_uvm: FileIOInterface @@ -335,20 +308,84 @@ class NVDevice(HCQCompiled[NVSignal]): uvm_vaddr_allocator: BumpAllocator = BumpAllocator(size=(1 << 48) - 1, base=low_uvm_vaddr_allocator.base + low_uvm_vaddr_allocator.size, wrap=False) host_object_enumerator: int = 0x1000 + def __init__(self, dev, device_id): + if NVKIface.root is None: + NVKIface.fd_ctl = FileIOInterface("/dev/nvidiactl", os.O_RDWR | os.O_CLOEXEC) + NVKIface.fd_uvm = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC) + self.fd_uvm_2 = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC) + NVKIface.root = self.rm_alloc(0, nv_gpu.NV01_ROOT_CLIENT, None, root=0) + uvm.initialize(self.fd_uvm) + with contextlib.suppress(RuntimeError): uvm.mm_initialize(self.fd_uvm_2, uvmFd=self.fd_uvm.fd) # this error is okay, CUDA hits it too + + nv_iowr(NVKIface.fd_ctl, nv_gpu.NV_ESC_CARD_INFO, gpus_info:=(nv_gpu.nv_ioctl_card_info_t*64)()) + visible_devices = [int(x) for x in (getenv('VISIBLE_DEVICES', getenv('CUDA_VISIBLE_DEVICES', ''))).split(',') if x.strip()] + NVKIface.gpus_info = [gpus_info[x] for x in visible_devices] if visible_devices else gpus_info + + self.dev, self.device_id = dev, device_id + if self.device_id >= len(NVKIface.gpus_info) or not NVKIface.gpus_info[self.device_id].valid: + raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?") + + self.fd_dev = self._new_gpu_fd() + self.gpu_info = self.rm_control(self.root, nv_gpu.NV0000_CTRL_CMD_GPU_GET_ID_INFO_V2, + nv_gpu.NV0000_CTRL_GPU_GET_ID_INFO_V2_PARAMS(gpuId=NVKIface.gpus_info[self.device_id].gpu_id)) + self.gpu_minor = NVKIface.gpus_info[self.device_id].minor_number + self.gpu_instance = self.gpu_info.deviceInstance + + def rm_alloc(self, parent, clss, params=None, root=None) -> int: + nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_ALLOC, made:=nv_gpu.NVOS21_PARAMETERS(hRoot=root if root is not None else self.root, + hObjectParent=parent, hClass=clss, pAllocParms=ctypes.cast(ctypes.byref(params), ctypes.c_void_p) if params is not None else None)) + if made.status == nv_gpu.NV_ERR_NO_MEMORY: raise MemoryError(f"rm_alloc returned {get_error_str(made.status)}") + if made.status != 0: raise RuntimeError(f"rm_alloc returned {get_error_str(made.status)}") + return made.hObjectNew + + def rm_control(self, obj, cmd, params=None): + nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_CONTROL, made:=nv_gpu.NVOS54_PARAMETERS(hClient=self.root, hObject=obj, cmd=cmd, + paramsSize=ctypes.sizeof(params), params=ctypes.cast(ctypes.byref(params), ctypes.c_void_p) if params is not None else None)) + if made.status != 0: raise RuntimeError(f"rm_control returned {get_error_str(made.status)}") + return params + + def setup_usermode(self): + clsinfo = self.rm_control(self.dev.nvdevice, nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST, nv_gpu.NV0080_CTRL_GPU_GET_CLASSLIST_PARAMS(numClasses=100, + classList=mv_address(classlist:=memoryview(bytearray(100 * 4)).cast('I')))) + self.nvclasses = {classlist[i] for i in range(clsinfo.numClasses)} + self.usermode_class:int = next(c for c in [nv_gpu.HOPPER_USERMODE_A, nv_gpu.TURING_USERMODE_A] if c in self.nvclasses) + self.gpfifo_class:int = next(c for c in [nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A, nv_gpu.AMPERE_CHANNEL_GPFIFO_A] if c in self.nvclasses) + self.compute_class:int = next(c for c in [nv_gpu.BLACKWELL_COMPUTE_B, nv_gpu.ADA_COMPUTE_A, nv_gpu.AMPERE_COMPUTE_B] if c in self.nvclasses) + self.dma_class:int = next(c for c in [nv_gpu.BLACKWELL_DMA_COPY_B, nv_gpu.AMPERE_DMA_COPY_B] if c in self.nvclasses) + + usermode = self.rm_alloc(self.dev.subdevice, self.usermode_class) + return usermode, MMIOInterface(self._gpu_map_to_cpu(usermode, mmio_sz:=0x10000, flags=2), mmio_sz, fmt='I') + + def setup_vm(self, vaspace): + self.rm_control(self.dev.subdevice, nv_gpu.NV2080_CTRL_CMD_GPU_GET_GID_INFO, raw_uuid:=nv_gpu.NV2080_CTRL_GPU_GET_GID_INFO_PARAMS( + flags=nv_gpu.NV2080_GPU_CMD_GPU_GET_GID_FLAGS_FORMAT_BINARY, length=16)) + self.gpu_uuid = nv_gpu.struct_nv_uuid(uuid=(ctypes.c_ubyte*16)(*[raw_uuid.data[i] for i in range(16)])) + + uvm.register_gpu(self.fd_uvm, rmCtrlFd=-1, gpu_uuid=self.gpu_uuid) + uvm.register_gpu_vaspace(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace) + + for dev in cast(list[NVDevice], self.dev.devices): + try: uvm.enable_peer_access(self.fd_uvm, gpuUuidA=self.gpu_uuid, gpuUuidB=dev.iface.gpu_uuid) + except RuntimeError as e: raise RuntimeError(f"{e}. Make sure GPUs #{self.gpu_minor} & #{dev.iface.gpu_minor} have P2P enabled.") from e + + def setup_gpfifo_vm(self, gpfifo): + uvm.register_channel(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, + hChannel=gpfifo, base=self._alloc_gpu_vaddr(0x4000000, force_low=True), length=0x4000000) + def _new_gpu_fd(self): - fd_dev = FileIOInterface(f"/dev/nvidia{NVDevice.gpus_info[self.device_id].minor_number}", os.O_RDWR | os.O_CLOEXEC) + fd_dev = FileIOInterface(f"/dev/nvidia{NVKIface.gpus_info[self.device_id].minor_number}", os.O_RDWR | os.O_CLOEXEC) nv_iowr(fd_dev, nv_gpu.NV_ESC_REGISTER_FD, nv_gpu.nv_ioctl_register_fd_t(ctl_fd=self.fd_ctl.fd)) return fd_dev def _gpu_map_to_cpu(self, memory_handle, size, target=None, flags=0, system=False): fd_dev = self._new_gpu_fd() if not system else FileIOInterface("/dev/nvidiactl", os.O_RDWR | os.O_CLOEXEC) made = nv_gpu.nv_ioctl_nvos33_parameters_with_fd(fd=fd_dev.fd, - params=nv_gpu.NVOS33_PARAMETERS(hClient=self.root, hDevice=self.nvdevice, hMemory=memory_handle, length=size, flags=flags)) + params=nv_gpu.NVOS33_PARAMETERS(hClient=self.root, hDevice=self.dev.nvdevice, hMemory=memory_handle, length=size, flags=flags)) nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_MAP_MEMORY, made) if made.params.status != 0: raise RuntimeError(f"_gpu_map_to_cpu returned {get_error_str(made.params.status)}") return fd_dev.mmap(target, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if target is not None else 0), 0) - def _gpu_alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, tag="") -> HCQBuffer: + def _gpu_alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0) -> HCQBuffer: # Uncached memory is "system". Use huge pages only for gpu memory. page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10))) size = round_up(size, page_size) @@ -360,9 +397,9 @@ class NVDevice(HCQCompiled[NVSignal]): flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \ | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30) - NVDevice.host_object_enumerator += 1 - made = nv_gpu.nv_ioctl_nvos02_parameters_with_fd(params=nv_gpu.NVOS02_PARAMETERS(hRoot=self.root, hObjectParent=self.nvdevice, flags=flags, - hObjectNew=NVDevice.host_object_enumerator, hClass=nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR, pMemory=va_addr, limit=size-1), fd=-1) + NVKIface.host_object_enumerator += 1 + made = nv_gpu.nv_ioctl_nvos02_parameters_with_fd(params=nv_gpu.NVOS02_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, flags=flags, + hObjectNew=NVKIface.host_object_enumerator, hClass=nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR, pMemory=va_addr, limit=size-1), fd=-1) nv_iowr(self.fd_dev, nv_gpu.NV_ESC_RM_ALLOC_MEMORY, made) if made.params.status != 0: raise RuntimeError(f"host alloc returned {get_error_str(made.params.status)}") @@ -380,28 +417,26 @@ class NVDevice(HCQCompiled[NVSignal]): alloc_func = nv_gpu.NV1_MEMORY_SYSTEM if uncached else nv_gpu.NV1_MEMORY_USER alloc_params = nv_gpu.NV_MEMORY_ALLOCATION_PARAMS(owner=self.root, alignment=page_size, offset=0, limit=size-1, format=6, size=size, type=nv_gpu.NVOS32_TYPE_NOTIFIER if uncached else nv_gpu.NVOS32_TYPE_IMAGE, attr=attr, attr2=attr2, flags=fl) - mem_handle = rm_alloc(self.fd_ctl, alloc_func, self.root, self.nvdevice, alloc_params).hObjectNew + mem_handle = self.rm_alloc(self.dev.nvdevice, alloc_func, alloc_params) if cpu_access: va_addr = self._gpu_map_to_cpu(mem_handle, size, target=va_addr, flags=map_flags, system=uncached) - return self._gpu_uvm_map(va_addr, size, mem_handle, has_cpu_mapping=cpu_access or host, tag=tag) + return self._gpu_uvm_map(va_addr, size, mem_handle, has_cpu_mapping=cpu_access or host) def _gpu_free(self, mem:HCQBuffer): - if mem.meta.hMemory > NVDevice.host_object_enumerator: # not a host object, clear phys mem. - made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.nvdevice, hObjectOld=mem.meta.hMemory) + if mem.meta.hMemory > NVKIface.host_object_enumerator: # not a host object, clear phys mem. + made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, hObjectOld=mem.meta.hMemory) nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_FREE, made) if made.status != 0: raise RuntimeError(f"_gpu_free returned {get_error_str(made.status)}") - self._debug_mappings.pop((cast(int, mem.va_addr), mem.size)) uvm.free(self.fd_uvm, base=cast(int, mem.va_addr), length=mem.size) if mem.meta.has_cpu_mapping: FileIOInterface.munmap(cast(int, mem.va_addr), mem.size) - def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False, tag="") -> HCQBuffer: + def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False) -> HCQBuffer: if create_range: uvm.create_external_range(self.fd_uvm, base=va_base, length=size) attrs = (nv_gpu.struct_c__SA_UvmGpuMappingAttributes*256)(nv_gpu.struct_c__SA_UvmGpuMappingAttributes(gpuUuid=self.gpu_uuid, gpuMappingType=1)) # NOTE: va_addr is set to make rawbufs compatible with HCQBuffer protocol. - self._debug_mappings[(va_base, size)] = tag return HCQBuffer(va_base, size, meta=uvm.map_external_allocation(self.fd_uvm, base=va_base, length=size, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid], has_cpu_mapping=has_cpu_mapping), @@ -410,84 +445,49 @@ class NVDevice(HCQCompiled[NVSignal]): def _gpu_map(self, mem:HCQBuffer): if self.gpu_uuid in mem.meta.mapped_gpu_ids: return mem.meta.mapped_gpu_ids.append(self.gpu_uuid) - self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False, tag="p2p mem") + self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False) def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False): - return NVDevice.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVDevice.uvm_vaddr_allocator.alloc(size, alignment) + return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment) - def _setup_nvclasses(self): - classlist = memoryview(bytearray(100 * 4)).cast('I') - clsinfo = rmctrl.gpu_get_classlist(self.fd_ctl, self.root, self.nvdevice, numClasses=100, classList=mv_address(classlist)) - self.nvclasses = {classlist[i] for i in range(clsinfo.numClasses)} - self.usermode_class:int = next(c for c in [nv_gpu.HOPPER_USERMODE_A, nv_gpu.TURING_USERMODE_A] if c in self.nvclasses) - self.gpfifo_class:int = next(c for c in [nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A, nv_gpu.AMPERE_CHANNEL_GPFIFO_A] if c in self.nvclasses) - self.compute_class:int = next(c for c in [nv_gpu.BLACKWELL_COMPUTE_B, nv_gpu.ADA_COMPUTE_A, nv_gpu.AMPERE_COMPUTE_B] if c in self.nvclasses) - self.dma_class:int = next(c for c in [nv_gpu.BLACKWELL_DMA_COPY_B, nv_gpu.AMPERE_DMA_COPY_B] if c in self.nvclasses) +class NVDevice(HCQCompiled[NVSignal]): + devices: ClassVar[list[HCQCompiled]] = [] + signal_pages: ClassVar[list[HCQBuffer]] = [] + signal_pool: ClassVar[list[HCQBuffer]] = [] def __init__(self, device:str=""): - if NVDevice.root is None: - NVDevice.fd_ctl = FileIOInterface("/dev/nvidiactl", os.O_RDWR | os.O_CLOEXEC) - NVDevice.fd_uvm = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC) - self.fd_uvm_2 = FileIOInterface("/dev/nvidia-uvm", os.O_RDWR | os.O_CLOEXEC) - NVDevice.root = rm_alloc(self.fd_ctl, nv_gpu.NV01_ROOT_CLIENT, 0, 0, None).hObjectNew - uvm.initialize(self.fd_uvm) - with contextlib.suppress(RuntimeError): uvm.mm_initialize(self.fd_uvm_2, uvmFd=self.fd_uvm.fd) # this error is okay, CUDA hits it too - - nv_iowr(NVDevice.fd_ctl, nv_gpu.NV_ESC_CARD_INFO, gpus_info:=(nv_gpu.nv_ioctl_card_info_t*64)()) - visible_devices = [int(x) for x in (getenv('VISIBLE_DEVICES', getenv('CUDA_VISIBLE_DEVICES', ''))).split(',') if x.strip()] - NVDevice.gpus_info = [gpus_info[x] for x in visible_devices] if visible_devices else gpus_info - self.device_id = int(device.split(":")[1]) if ":" in device else 0 + self.iface = NVKIface(self, self.device_id) - if self.device_id >= len(NVDevice.gpus_info) or not NVDevice.gpus_info[self.device_id].valid: - raise RuntimeError(f"No device found for {device}. Requesting more devices than the system has?") - - self.fd_dev = self._new_gpu_fd() - self.gpu_info = rmctrl.gpu_get_id_info_v2(self.fd_ctl, self.root, self.root, gpuId=NVDevice.gpus_info[self.device_id].gpu_id) - self.gpu_minor = NVDevice.gpus_info[self.device_id].minor_number - - device_params = nv_gpu.NV0080_ALLOC_PARAMETERS(deviceId=self.gpu_info.deviceInstance, hClientShare=self.root, + device_params = nv_gpu.NV0080_ALLOC_PARAMETERS(deviceId=self.iface.gpu_instance, hClientShare=self.iface.root, vaMode=nv_gpu.NV_DEVICE_ALLOCATION_VAMODE_MULTIPLE_VASPACES) - self.nvdevice = rm_alloc(self.fd_ctl, nv_gpu.NV01_DEVICE_0, self.root, self.root, device_params).hObjectNew - self.subdevice = rm_alloc(self.fd_ctl, nv_gpu.NV20_SUBDEVICE_0, self.root, self.nvdevice, None).hObjectNew + self.nvdevice = self.iface.rm_alloc(self.iface.root, nv_gpu.NV01_DEVICE_0, device_params) + self.subdevice = self.iface.rm_alloc(self.nvdevice, nv_gpu.NV20_SUBDEVICE_0) + self.usermode, self.gpu_mmio = self.iface.setup_usermode() - self._setup_nvclasses() - self._debug_mappings: dict[tuple[int, int], str] = dict() - - self.usermode = rm_alloc(self.fd_ctl, self.usermode_class, self.root, self.subdevice, None).hObjectNew - self.gpu_mmio = MMIOInterface(self._gpu_map_to_cpu(self.usermode, mmio_sz:=0x10000, flags=2), mmio_sz, fmt='I') - - rmctrl.perf_boost(self.fd_ctl, self.root, self.subdevice, duration=0xffffffff, flags=((nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CUDA_YES << 4) | \ - (nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CUDA_PRIORITY_HIGH << 6) | (nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CMD_BOOST_TO_MAX << 0))) + self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_PERF_BOOST, nv_gpu.NV2080_CTRL_PERF_BOOST_PARAMS(duration=0xffffffff, + flags=((nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CUDA_YES << 4) | (nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CUDA_PRIORITY_HIGH << 6) | \ + (nv_gpu.NV2080_CTRL_PERF_BOOST_FLAGS_CMD_BOOST_TO_MAX)))) vaspace_params = nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS(vaBase=0x1000, vaSize=0x1fffffb000000, flags=nv_gpu.NV_VASPACE_ALLOCATION_FLAGS_ENABLE_PAGE_FAULTING | nv_gpu.NV_VASPACE_ALLOCATION_FLAGS_IS_EXTERNALLY_OWNED) - vaspace = rm_alloc(self.fd_ctl, nv_gpu.FERMI_VASPACE_A, self.root, self.nvdevice, vaspace_params).hObjectNew + vaspace = self.iface.rm_alloc(self.nvdevice, nv_gpu.FERMI_VASPACE_A, vaspace_params) - raw_uuid = rmctrl.gpu_get_gid_info(self.fd_ctl, self.root, self.subdevice, flags=nv_gpu.NV2080_GPU_CMD_GPU_GET_GID_FLAGS_FORMAT_BINARY, length=16) - self.gpu_uuid = nv_gpu.struct_nv_uuid(uuid=(ctypes.c_ubyte*16)(*[raw_uuid.data[i] for i in range(16)])) - - uvm.register_gpu(self.fd_uvm, rmCtrlFd=-1, gpu_uuid=self.gpu_uuid) - uvm.register_gpu_vaspace(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace) - - for dev in cast(list[NVDevice], self.devices): - try: uvm.enable_peer_access(self.fd_uvm, gpuUuidA=self.gpu_uuid, gpuUuidB=dev.gpu_uuid) - except RuntimeError as e: raise RuntimeError(str(e) + f". Make sure GPUs #{self.gpu_minor} & #{dev.gpu_minor} have P2P enabled between.") from e + self.iface.setup_vm(vaspace) channel_params = nv_gpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS(engineType=nv_gpu.NV2080_ENGINE_TYPE_GRAPHICS) - channel_group = rm_alloc(self.fd_ctl, nv_gpu.KEPLER_CHANNEL_GROUP_A, self.root, self.nvdevice, channel_params).hObjectNew + channel_group = self.iface.rm_alloc(self.nvdevice, nv_gpu.KEPLER_CHANNEL_GROUP_A, channel_params) - gpfifo_area = self._gpu_alloc(0x200000, contiguous=True, cpu_access=True, map_flags=0x10d0000, tag="gpfifo") + gpfifo_area = self.iface._gpu_alloc(0x200000, contiguous=True, cpu_access=True, map_flags=0x10d0000) ctxshare_params = nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC) - ctxshare = rm_alloc(self.fd_ctl, nv_gpu.FERMI_CONTEXT_SHARE_A, self.root, channel_group, ctxshare_params).hObjectNew + ctxshare = self.iface.rm_alloc(channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A, ctxshare_params) self.compute_gpfifo = self._new_gpu_fifo(gpfifo_area, ctxshare, channel_group, offset=0, entries=0x10000, enable_debug=True) self.dma_gpfifo = self._new_gpu_fifo(gpfifo_area, ctxshare, channel_group, offset=0x100000, entries=0x10000) + self.iface.rm_control(channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1)) - rmctrl.gpfifo_schedule(self.fd_ctl, self.root, channel_group, bEnable=1) - - self.cmdq_page:HCQBuffer = self._gpu_alloc(0x200000, cpu_access=True, tag="cmdq") + self.cmdq_page:HCQBuffer = self.iface._gpu_alloc(0x200000, cpu_access=True) self.cmdq_allocator = BumpAllocator(size=self.cmdq_page.size, base=cast(int, self.cmdq_page.va_addr), wrap=True) self.cmdq = MMIOInterface(cast(int, self.cmdq_page.va_addr), 0x200000, fmt='I') @@ -505,25 +505,22 @@ class NVDevice(HCQCompiled[NVSignal]): self._setup_gpfifos() def _new_gpu_fifo(self, gpfifo_area, ctxshare, channel_group, offset=0, entries=0x400, enable_debug=False) -> GPFifo: - notifier = self._gpu_alloc(48 << 20, uncached=True) + notifier = self.iface._gpu_alloc(48 << 20, uncached=True) params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(hObjectError=notifier.meta.hMemory, hObjectBuffer=gpfifo_area.meta.hMemory, gpFifoOffset=gpfifo_area.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare, hUserdMemory=(ctypes.c_uint32*8)(gpfifo_area.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset)) - gpfifo = rm_alloc(self.fd_ctl, self.gpfifo_class, self.root, channel_group, params).hObjectNew - comp = rm_alloc(self.fd_ctl, self.compute_class, self.root, gpfifo, None).hObjectNew - rm_alloc(self.fd_ctl, self.dma_class, self.root, gpfifo, None) + gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params) + comp = self.iface.rm_alloc(gpfifo, self.iface.compute_class) + self.iface.rm_alloc(gpfifo, self.iface.dma_class) if enable_debug: self.debug_compute_obj, self.debug_channel = comp, gpfifo - debugger_params = nv_gpu.NV83DE_ALLOC_PARAMETERS(hAppClient=self.root, hClass3dObject=self.debug_compute_obj) - self.debugger = rm_alloc(self.fd_ctl, nv_gpu.GT200_DEBUGGER, self.root, self.nvdevice, debugger_params).hObjectNew + debugger_params = nv_gpu.NV83DE_ALLOC_PARAMETERS(hAppClient=self.iface.root, hClass3dObject=self.debug_compute_obj) + self.debugger = self.iface.rm_alloc(self.nvdevice, nv_gpu.GT200_DEBUGGER, debugger_params) - ws_token_params = rmctrl.gpfifo_get_work_submit_token(self.fd_ctl, self.root, gpfifo, workSubmitToken=-1) - assert ws_token_params.workSubmitToken != -1 - - channel_base = self._alloc_gpu_vaddr(0x4000000, force_low=True) - uvm.register_channel(self.fd_uvm, gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, - hChannel=gpfifo, base=channel_base, length=0x4000000) + self.iface.rm_control(gpfifo, nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN, + ws_token_params:=nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS(workSubmitToken=-1)) + self.iface.setup_gpfifo_vm(gpfifo) return GPFifo(ring=MMIOInterface(gpfifo_area.va_addr + offset, entries*8, fmt='Q'), entries_count=entries, token=ws_token_params.workSubmitToken, controls=nv_gpu.AmpereAControlGPFifo.from_address(gpfifo_area.va_addr + offset + entries * 8)) @@ -531,21 +528,22 @@ class NVDevice(HCQCompiled[NVSignal]): def _query_gpu_info(self, *reqs): nvrs = [getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_'+r.upper(), getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_LITTER_'+r.upper(),None)) for r in reqs] infos = (nv_gpu.NV2080_CTRL_GR_INFO*len(nvrs))(*[nv_gpu.NV2080_CTRL_GR_INFO(index=nvr) for nvr in nvrs]) - rmctrl.gr_get_info(self.fd_ctl, self.root, self.subdevice, grInfoListSize=len(infos), grInfoList=ctypes.addressof(infos)) + self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_GR_GET_INFO, + nv_gpu.NV2080_CTRL_GR_GET_INFO_PARAMS(grInfoListSize=len(infos), grInfoList=ctypes.addressof(infos))) return [x.data for x in infos] def _setup_gpfifos(self): self.slm_per_thread, self.shader_local_mem = 0, None # Set windows addresses to not collide with other allocated buffers. - self.shared_mem_window = 0x729400000000 if self.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 0xfe000000 - self.local_mem_window = 0x729300000000 if self.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 0xff000000 + self.shared_mem_window = 0x729400000000 if self.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 0xfe000000 + self.local_mem_window = 0x729300000000 if self.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 0xff000000 - NVComputeQueue().setup(compute_class=self.compute_class, local_mem_window=self.local_mem_window, shared_mem_window=self.shared_mem_window) \ + NVComputeQueue().setup(compute_class=self.iface.compute_class, local_mem_window=self.local_mem_window, shared_mem_window=self.shared_mem_window) \ .signal(self.timeline_signal, self.timeline_value).submit(self) cast(NVCopyQueue, NVCopyQueue().wait(self.timeline_signal, self.timeline_value)) \ - .setup(copy_class=self.dma_class) \ + .setup(copy_class=self.iface.dma_class) \ .signal(self.timeline_signal, self.timeline_value + 1).submit(self) self.timeline_value += 2 @@ -565,24 +563,24 @@ class NVDevice(HCQCompiled[NVSignal]): .signal(self.timeline_signal, self.next_timeline()).submit(self) def invalidate_caches(self): - rmctrl.fb_flush_gpu_cache(self.fd_ctl, self.root, self.subdevice, + self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_FB_FLUSH_GPU_CACHE, nv_gpu.NV2080_CTRL_FB_FLUSH_GPU_CACHE_PARAMS( flags=((nv_gpu.NV2080_CTRL_FB_FLUSH_GPU_CACHE_FLAGS_WRITE_BACK_YES << 2) | (nv_gpu.NV2080_CTRL_FB_FLUSH_GPU_CACHE_FLAGS_INVALIDATE_YES << 3) | - (nv_gpu.NV2080_CTRL_FB_FLUSH_GPU_CACHE_FLAGS_FLUSH_MODE_FULL_CACHE << 4))) + (nv_gpu.NV2080_CTRL_FB_FLUSH_GPU_CACHE_FLAGS_FLUSH_MODE_FULL_CACHE << 4)))) def on_device_hang(self): # Prepare fault report. # TODO: Restore the GPU using NV83DE_CTRL_CMD_CLEAR_ALL_SM_ERROR_STATES if needed. report = [] - sm_errors = rmctrl.debug_read_all_sm_error_states(self.fd_ctl, self.root, self.debugger, hTargetChannel=self.debug_channel, numSMsToRead=100) + sm_errors = self.iface.rm_control(self.debugger, nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_ALL_SM_ERROR_STATES, + nv_gpu.NV83DE_CTRL_DEBUG_READ_ALL_SM_ERROR_STATES_PARAMS(hTargetChannel=self.debug_channel, numSMsToRead=100)) if sm_errors.mmuFault.valid: - mmu_info = rmctrl.debug_read_mmu_fault_info(self.fd_ctl, self.root, self.debugger) - for i in range(mmu_info.count): - pfinfo = mmu_info.mmuFaultInfoList[i] + mmu = self.iface.rm_control(self.debugger, nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_MMU_FAULT_INFO, + nv_gpu.NV83DE_CTRL_DEBUG_READ_MMU_FAULT_INFO_PARAMS()) + for i in range(mmu.count): + pfinfo = mmu.mmuFaultInfoList[i] report += [f"MMU fault: 0x{pfinfo.faultAddress:X} | {NV_PFAULT_FAULT_TYPE[pfinfo.faultType]} | {NV_PFAULT_ACCESS_TYPE[pfinfo.accessType]}"] - if DEBUG >= 5: - report += ["GPU mappings:\n"+"\n".join(f"\t0x{x:X} - 0x{x+y-1:X} | {self._debug_mappings[(x,y)]}" for x,y in sorted(self._debug_mappings))] else: for i, e in enumerate(sm_errors.smErrorStateArray): if e.hwwGlobalEsr or e.hwwWarpEsr: report += [f"SM {i} fault: esr={e.hwwGlobalEsr} warp_esr={e.hwwWarpEsr} warp_pc={e.hwwWarpEsrPc64}"]