diff options
author | Alyssa Ross <hi@alyssa.is> | 2024-03-24 11:04:41 +0100 |
---|---|---|
committer | Alyssa Ross <hi@alyssa.is> | 2024-03-24 11:04:41 +0100 |
commit | 5423cabbbf2b6dec5568f1ecabd288d5d9a642ec (patch) | |
tree | f316a6a921bfefd3a63bd4502c2eb50ff1644f67 /nixpkgs/pkgs/development/python-modules/gpuctypes | |
parent | 46a88117a05c3469af5d99433af140c3de8ca088 (diff) | |
parent | 8aa81f34981add12aecada6c702ddbbd0375ca36 (diff) | |
download | nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar.gz nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar.bz2 nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar.lz nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar.xz nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.tar.zst nixlib-5423cabbbf2b6dec5568f1ecabd288d5d9a642ec.zip |
Merge branch 'nixos-unstable-small' of https://github.com/NixOS/nixpkgs
Diffstat (limited to 'nixpkgs/pkgs/development/python-modules/gpuctypes')
3 files changed, 203 insertions, 0 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/gpuctypes/0001-fix-dlopen-cuda.patch b/nixpkgs/pkgs/development/python-modules/gpuctypes/0001-fix-dlopen-cuda.patch new file mode 100644 index 000000000000..bc9f6c7ec64b --- /dev/null +++ b/nixpkgs/pkgs/development/python-modules/gpuctypes/0001-fix-dlopen-cuda.patch @@ -0,0 +1,44 @@ +From d448321436e8314d3e2a6a09d4017c4bc10f612d Mon Sep 17 00:00:00 2001 +From: Gaetan Lepage <gaetan@glepage.com> +Date: Sat, 17 Feb 2024 17:37:22 +0100 +Subject: [PATCH] fix-dlopen-cuda + +--- + gpuctypes/cuda.py | 20 ++++++++++++++++++-- + 1 file changed, 18 insertions(+), 2 deletions(-) + +diff --git a/gpuctypes/cuda.py b/gpuctypes/cuda.py +index acba81c..091f7f7 100644 +--- a/gpuctypes/cuda.py ++++ b/gpuctypes/cuda.py +@@ -143,9 +143,25 @@ def char_pointer_cast(string, encoding='utf-8'): + + + ++NAME_TO_PATHS = { ++ "libcuda.so": ["@driverLink@/lib/libcuda.so"], ++ "libnvrtc.so": ["@libnvrtc@"], ++} ++def _try_dlopen(name): ++ try: ++ return ctypes.CDLL(name) ++ except OSError: ++ pass ++ for candidate in NAME_TO_PATHS.get(name, []): ++ try: ++ return ctypes.CDLL(candidate) ++ except OSError: ++ pass ++ raise RuntimeError(f"{name} not found") ++ + _libraries = {} +-_libraries['libcuda.so'] = ctypes.CDLL(ctypes.util.find_library('cuda')) +-_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc')) ++_libraries['libcuda.so'] = _try_dlopen('libcuda.so') ++_libraries['libnvrtc.so'] = _try_dlopen('libnvrtc.so') + + + cuuint32_t = ctypes.c_uint32 +-- +2.43.0 + diff --git a/nixpkgs/pkgs/development/python-modules/gpuctypes/default.nix b/nixpkgs/pkgs/development/python-modules/gpuctypes/default.nix new file mode 100644 index 000000000000..e5987f78ee07 --- /dev/null +++ b/nixpkgs/pkgs/development/python-modules/gpuctypes/default.nix @@ -0,0 +1,127 @@ +{ lib +, config +, buildPythonPackage +, fetchFromGitHub +, substituteAll +, addDriverRunpath +, cudaSupport ? config.cudaSupport +, rocmSupport ? config.rocmSupport +, cudaPackages +, setuptools +, ocl-icd +, rocmPackages +, pytestCheckHook +, gpuctypes +, testCudaRuntime ? false +, testOpenclRuntime ? false +, testRocmRuntime ? false +}: +assert testCudaRuntime -> cudaSupport; +assert testRocmRuntime -> rocmSupport; + +buildPythonPackage rec { + pname = "gpuctypes"; + version = "0.3.0"; + pyproject = true; + + src = fetchFromGitHub { + repo = "gpuctypes"; + owner = "tinygrad"; + rev = "refs/tags/${version}"; + hash = "sha256-xUMvMBK1UhZaMZfik0Ia6+siyZGpCkBV+LTnQvzt/rw="; + }; + + patches = [ + (substituteAll { + src = ./0001-fix-dlopen-cuda.patch; + inherit (addDriverRunpath) driverLink; + libnvrtc = + if cudaSupport + then "${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so" + else "Please import nixpkgs with `config.cudaSupport = true`"; + }) + ]; + + nativeBuildInputs = [ + setuptools + ]; + + postPatch = '' + substituteInPlace gpuctypes/opencl.py \ + --replace "ctypes.util.find_library('OpenCL')" "'${ocl-icd}/lib/libOpenCL.so'" + '' + # hipGetDevicePropertiesR0600 is a symbol from rocm-6. We are currently at rocm-5. + # We are not sure that this works. Remove when rocm gets updated to version 6. + + lib.optionalString rocmSupport '' + substituteInPlace gpuctypes/hip.py \ + --replace "/opt/rocm/lib/libamdhip64.so" "${rocmPackages.clr}/lib/libamdhip64.so" \ + --replace "/opt/rocm/lib/libhiprtc.so" "${rocmPackages.clr}/lib/libhiprtc.so" \ + --replace "hipGetDevicePropertiesR0600" "hipGetDeviceProperties" + + substituteInPlace gpuctypes/comgr.py \ + --replace "/opt/rocm/lib/libamd_comgr.so" "${rocmPackages.rocm-comgr}/lib/libamd_comgr.so" + ''; + + pythonImportsCheck = [ "gpuctypes" ]; + + nativeCheckInputs = [ + pytestCheckHook + ]; + + disabledTestPaths = lib.optionals (!testOpenclRuntime) [ + "test/test_opencl.py" + ] ++ lib.optionals (!rocmSupport) [ + "test/test_hip.py" + ] ++ lib.optionals (!cudaSupport) [ + "test/test_cuda.py" + ]; + + # Require GPU access to run (not available in the sandbox) + pytestFlagsArray = lib.optionals (!testCudaRuntime) [ + "-k" "'not TestCUDADevice'" + ] ++ lib.optionals (!testRocmRuntime) [ + "-k" "'not TestHIPDevice'" + ] ++ lib.optionals (testCudaRuntime || testOpenclRuntime || testRocmRuntime) [ + "-v" + ]; + + # Running these tests requires special configuration on the builder. + # e.g. https://github.com/NixOS/nixpkgs/pull/256230 implements a nix + # pre-build hook which exposes the devices and the drivers in the sandbox + # based on requiredSystemFeatures: + requiredSystemFeatures = lib.optionals testCudaRuntime [ + "cuda" + ] ++ lib.optionals testOpenclRuntime [ + "opencl" + ] ++ lib.optionals testRocmRuntime [ + "rocm" + ]; + + passthru.gpuChecks = { + cuda = gpuctypes.override { + cudaSupport = true; + testCudaRuntime = true; + }; + opencl = gpuctypes.override { + testOpenclRuntime = true; + }; + rocm = gpuctypes.override { + rocmSupport = true; + testRocmRuntime = true; + }; + }; + + preCheck = lib.optionalString (cudaSupport && !testCudaRuntime) '' + addToSearchPath LD_LIBRARY_PATH ${lib.getLib cudaPackages.cuda_cudart}/lib/stubs + ''; + + # If neither rocmSupport or cudaSupport is enabled, no tests are selected + dontUsePytestCheck = !(rocmSupport || cudaSupport) && (!testOpenclRuntime); + + meta = with lib; { + description = "Ctypes wrappers for HIP, CUDA, and OpenCL"; + homepage = "https://github.com/tinygrad/gpuctypes"; + license = licenses.mit; + maintainers = with maintainers; [ GaetanLepage matthewcroughan wozeparrot ]; + }; +} diff --git a/nixpkgs/pkgs/development/python-modules/gpuctypes/fix-dlopen-cuda.patch b/nixpkgs/pkgs/development/python-modules/gpuctypes/fix-dlopen-cuda.patch new file mode 100644 index 000000000000..8d3b69e35e11 --- /dev/null +++ b/nixpkgs/pkgs/development/python-modules/gpuctypes/fix-dlopen-cuda.patch @@ -0,0 +1,32 @@ +diff --git a/gpuctypes/cuda.py b/gpuctypes/cuda.py +index acba81c..aac5fc7 100644 +--- a/gpuctypes/cuda.py ++++ b/gpuctypes/cuda.py +@@ -143,9 +143,25 @@ def char_pointer_cast(string, encoding='utf-8'): + + + ++NAME_TO_PATHS = { ++ "libcuda.so": ["@driverLink@/lib/libcuda.so"], ++ "libnvrtc.so": ["@libnvrtc@"], ++} ++def _try_dlopen(name): ++ try: ++ return ctypes.CDLL(name) ++ except OSError: ++ pass ++ for candidate in NAME_TO_PATHS.get(name, []): ++ try: ++ return ctypes.CDLL(candidate) ++ except OSError: ++ pass ++ raise RuntimeError(f"{name} not found") ++ + _libraries = {} +-_libraries['libcuda.so'] = ctypes.CDLL(ctypes.util.find_library('cuda')) +-_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc')) ++_libraries['libcuda.so'] = _try_dlopen('libcuda.so') ++_libraries['libnvrtc.so'] = _try_dlopen('libnvrtc.so') + + + cuuint32_t = ctypes.c_uint32 |