diff --git a/test/mockgpu/cuda/cuda.py b/test/mockgpu/cuda/cuda.py index 0889f4899e..466d269a9f 100644 --- a/test/mockgpu/cuda/cuda.py +++ b/test/mockgpu/cuda/cuda.py @@ -1,15 +1,14 @@ 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 test.mockgpu.helpers import _try_dlopen_gpuocelot from tinygrad.helpers import mv_address for attr in dir(orig_cuda): if not attr.startswith('__'): globals()[attr] = getattr(orig_cuda, attr) -gpuocelot_lib = _try_dlopen_gpuocelot() # Global state class CUDAState: @@ -128,7 +127,7 @@ def cuModuleUnload(hmod) -> int: 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: gpuocelot_lib.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) + 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 diff --git a/test/mockgpu/helpers.py b/test/mockgpu/helpers.py index 5aa8d085f2..1d3a8648d9 100644 --- a/test/mockgpu/helpers.py +++ b/test/mockgpu/helpers.py @@ -1,18 +1,11 @@ -import ctypes, ctypes.util +import ctypes +from tinygrad.runtime.support import c -def _try_dlopen_gpuocelot(): - GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else [] - GPUOCELOT_PATHS += ["libgpuocelot.so", "/usr/local/lib/libgpuocelot.so", - "libgpuocelot.dylib", "/usr/local/lib/libgpuocelot.dylib", "/opt/homebrew/lib/libgpuocelot.dylib"] - for path in GPUOCELOT_PATHS: - try: - gpuocelot_lib = ctypes.CDLL(path) - gpuocelot_lib.ptx_run.argtypes = [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] - except OSError: pass - else: return gpuocelot_lib - print("Could not find libgpuocelot.so") - return None +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.""" diff --git a/test/mockgpu/nv/nvgpu.py b/test/mockgpu/nv/nvgpu.py index 27b82d6bd7..eb6cf7f7ff 100644 --- a/test/mockgpu/nv/nvgpu.py +++ b/test/mockgpu/nv/nvgpu.py @@ -2,7 +2,7 @@ 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 _try_dlopen_gpuocelot +from test.mockgpu.helpers import ptx_run from tinygrad.helpers import to_mv from tinygrad.runtime.support.c import init_c_struct_t @@ -16,8 +16,6 @@ def make_qmd_struct_type(): return init_c_struct_t(0x40 * 4, tuple(fields)) qmd_struct_t = make_qmd_struct_type() -gpuocelot_lib = _try_dlopen_gpuocelot() - class SchedResult(Enum): CONT = auto(); YIELD = auto() # noqa: E702 class GPFIFO: @@ -97,9 +95,7 @@ class GPFIFO: 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: - gpuocelot_lib.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) + 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')