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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5   cython_0,
   6   fastrlock,
   7   numpy,
   8   wheel,
   9   pytestCheckHook,
  10   mock,
  11   setuptools,
  12   cudaPackages,
  13   addDriverRunpath,
  14   pythonOlder,
  15   symlinkJoin,
  16   fetchpatch
  17 }:
  18
  19 let
  20   inherit (cudaPackages) cudnn cutensor nccl;
  21   outpaths = with cudaPackages; [
  22       cuda_cccl # <nv/target>
  23       cuda_cudart
  24       cuda_nvcc # <crt/host_defines.h>
  25       cuda_nvprof
  26       cuda_nvrtc
  27       cuda_nvtx
  28       cuda_profiler_api
  29       libcublas
  30       libcufft
  31       libcurand
  32       libcusolver
  33       libcusparse
  34
  35       # Missing:
  36       # cusparselt
  37   ];
  38   cudatoolkit-joined = symlinkJoin {
  39     name = "cudatoolkit-joined-${cudaPackages.cudaVersion}";
  40     paths = outpaths ++ lib.concatMap (f: lib.map f outpaths) [lib.getLib lib.getDev (lib.getOutput "static") (lib.getOutput "stubs")];
  41   };
  42 in
  43 buildPythonPackage rec {
  44   pname = "cupy";
  45   version = "13.2.0";
  46   format = "setuptools";
  47
  48   disabled = pythonOlder "3.7";
  49
  50   src = fetchFromGitHub {
  51     owner = "cupy";
  52     repo = "cupy";
  53     rev = "v13.2.0";
  54     hash = "sha256-vZAtpIZztmsYeJeuq7yl7kgZse2azrIM3efHDmUswJI=";
  55     fetchSubmodules = true;
  56   };
  57
  58   patches = [
  59     (fetchpatch {
  60       url =
  61         "https://github.com/cfhammill/cupy/commit/67526c756e4a0a70f0420bf0e7f081b8a35a8ee5.patch";
  62       hash = "sha256-WZgexBdM9J0ep5s+9CGZriVq0ZidCRccox+g0iDDywQ=";
  63     })
  64   ];
  65
  66   # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both
  67   # CUPY_NUM_BUILD_JOBS and CUPY_NUM_NVCC_THREADS to NIX_BUILD_CORES results in
  68   # a small amount of thrashing but it turns out there are a large number of
  69   # very short builds and a few extremely long ones, so setting both ends up
  70   # working nicely in practice.
  71   preConfigure = ''
  72     export CUPY_NUM_BUILD_JOBS="$NIX_BUILD_CORES"
  73     export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES"
  74   '';
  75
  76   nativeBuildInputs = [
  77     setuptools
  78     wheel
  79     addDriverRunpath
  80     cython_0
  81     cudaPackages.cuda_nvcc
  82   ];
  83
  84   buildInputs = [
  85     cudatoolkit-joined
  86     cudnn
  87     cutensor
  88     nccl
  89   ];
  90
  91   NVCC = "${lib.getExe cudaPackages.cuda_nvcc}"; # FIXME: splicing/buildPackages
  92   CUDA_PATH = "${cudatoolkit-joined}";
  93
  94   propagatedBuildInputs = [
  95     fastrlock
  96     numpy
  97   ];
  98
  99   nativeCheckInputs = [
 100     pytestCheckHook
 101     mock
 102   ];
 103
 104   # Won't work with the GPU, whose drivers won't be accessible from the build
 105   # sandbox
 106   doCheck = false;
 107
 108   postFixup = ''
 109     find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
 110       addDriverRunpath "$lib"
 111     done
 112   '';
 113
 114   enableParallelBuilding = true;
 115
 116   meta = with lib; {
 117     description = "NumPy-compatible matrix library accelerated by CUDA";
 118     homepage = "https://cupy.chainer.org/";
 119     changelog = "https://github.com/cupy/cupy/releases/tag/v${version}";
 120     license = licenses.mit;
 121     platforms = [ "x86_64-linux" ];
 122     maintainers = with maintainers; [ hyphon81 ];
 123   };
 124 }