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

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