pkgs/development/python-modules/gpuctypes/default.nix

   1 {
   2   lib,
   3   config,
   4   buildPythonPackage,
   5   fetchFromGitHub,
   6   substituteAll,
   7   addDriverRunpath,
   8   cudaSupport ? config.cudaSupport,
   9   rocmSupport ? config.rocmSupport,
  10   cudaPackages,
  11   setuptools,
  12   ocl-icd,
  13   rocmPackages,
  14   pytestCheckHook,
  15   gpuctypes,
  16   testCudaRuntime ? false,
  17   testOpenclRuntime ? false,
  18   testRocmRuntime ? false,
  19 }:
  20 assert testCudaRuntime -> cudaSupport;
  21 assert testRocmRuntime -> rocmSupport;
  22
  23 buildPythonPackage rec {
  24   pname = "gpuctypes";
  25   version = "0.3.0";
  26   pyproject = true;
  27
  28   src = fetchFromGitHub {
  29     repo = "gpuctypes";
  30     owner = "tinygrad";
  31     rev = "refs/tags/${version}";
  32     hash = "sha256-xUMvMBK1UhZaMZfik0Ia6+siyZGpCkBV+LTnQvzt/rw=";
  33   };
  34
  35   patches = [
  36     (substituteAll {
  37       src = ./0001-fix-dlopen-cuda.patch;
  38       inherit (addDriverRunpath) driverLink;
  39       libnvrtc =
  40         if cudaSupport then
  41           "${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so"
  42         else
  43           "Please import nixpkgs with `config.cudaSupport = true`";
  44     })
  45   ];
  46
  47   nativeBuildInputs = [ setuptools ];
  48
  49   postPatch =
  50     ''
  51       substituteInPlace gpuctypes/opencl.py \
  52         --replace "ctypes.util.find_library('OpenCL')" "'${ocl-icd}/lib/libOpenCL.so'"
  53     ''
  54     # hipGetDevicePropertiesR0600 is a symbol from rocm-6. We are currently at rocm-5.
  55     # We are not sure that this works. Remove when rocm gets updated to version 6.
  56     + lib.optionalString rocmSupport ''
  57       substituteInPlace gpuctypes/hip.py \
  58         --replace "/opt/rocm/lib/libamdhip64.so" "${rocmPackages.clr}/lib/libamdhip64.so" \
  59         --replace "/opt/rocm/lib/libhiprtc.so" "${rocmPackages.clr}/lib/libhiprtc.so" \
  60         --replace "hipGetDevicePropertiesR0600" "hipGetDeviceProperties"
  61
  62       substituteInPlace gpuctypes/comgr.py \
  63         --replace "/opt/rocm/lib/libamd_comgr.so" "${rocmPackages.rocm-comgr}/lib/libamd_comgr.so"
  64     '';
  65
  66   pythonImportsCheck = [ "gpuctypes" ];
  67
  68   nativeCheckInputs = [ pytestCheckHook ];
  69
  70   disabledTestPaths =
  71     lib.optionals (!testOpenclRuntime) [ "test/test_opencl.py" ]
  72     ++ lib.optionals (!rocmSupport) [ "test/test_hip.py" ]
  73     ++ lib.optionals (!cudaSupport) [ "test/test_cuda.py" ];
  74
  75   # Require GPU access to run (not available in the sandbox)
  76   pytestFlagsArray =
  77     lib.optionals (!testCudaRuntime) [
  78       "-k"
  79       "'not TestCUDADevice'"
  80     ]
  81     ++ lib.optionals (!testRocmRuntime) [
  82       "-k"
  83       "'not TestHIPDevice'"
  84     ]
  85     ++ lib.optionals (testCudaRuntime || testOpenclRuntime || testRocmRuntime) [ "-v" ];
  86
  87   # Running these tests requires special configuration on the builder.
  88   # e.g. https://github.com/NixOS/nixpkgs/pull/256230 implements a nix
  89   # pre-build hook which exposes the devices and the drivers in the sandbox
  90   # based on requiredSystemFeatures:
  91   requiredSystemFeatures =
  92     lib.optionals testCudaRuntime [ "cuda" ]
  93     ++ lib.optionals testOpenclRuntime [ "opencl" ]
  94     ++ lib.optionals testRocmRuntime [ "rocm" ];
  95
  96   passthru.gpuChecks = {
  97     cuda = gpuctypes.override {
  98       cudaSupport = true;
  99       testCudaRuntime = true;
 100     };
 101     opencl = gpuctypes.override { testOpenclRuntime = true; };
 102     rocm = gpuctypes.override {
 103       rocmSupport = true;
 104       testRocmRuntime = true;
 105     };
 106   };
 107
 108   preCheck = lib.optionalString (cudaSupport && !testCudaRuntime) ''
 109     addToSearchPath LD_LIBRARY_PATH ${lib.getLib cudaPackages.cuda_cudart}/lib/stubs
 110   '';
 111
 112   # If neither rocmSupport or cudaSupport is enabled, no tests are selected
 113   dontUsePytestCheck = !(rocmSupport || cudaSupport) && (!testOpenclRuntime);
 114
 115   meta = with lib; {
 116     description = "Ctypes wrappers for HIP, CUDA, and OpenCL";
 117     homepage = "https://github.com/tinygrad/gpuctypes";
 118     license = licenses.mit;
 119     maintainers = with maintainers; [
 120       GaetanLepage
 121       matthewcroughan
 122       wozeparrot
 123     ];
 124   };
 125 }