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 }:
  15
  16 let
  17   inherit (cudaPackages) cudatoolkit cudnn cutensor nccl;
  18 in buildPythonPackage rec {
  19   pname = "cupy";
  20   version = "12.2.0";
  21
  22   disabled = pythonOlder "3.7";
  23
  24   src = fetchPypi {
  25     inherit pname version;
  26     hash = "sha256-+V/9Cv6sthewSP4Cjt4HuX3J6VrKFhCgIrHz0gqaAn4=";
  27   };
  28
  29   # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both
  30   # CUPY_NUM_BUILD_JOBS and CUPY_NUM_NVCC_THREADS to NIX_BUILD_CORES results in
  31   # a small amount of thrashing but it turns out there are a large number of
  32   # very short builds and a few extremely long ones, so setting both ends up
  33   # working nicely in practice.
  34   preConfigure = ''
  35     export CUDA_PATH=${cudatoolkit}
  36     export CUPY_NUM_BUILD_JOBS="$NIX_BUILD_CORES"
  37     export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES"
  38   '';
  39
  40   nativeBuildInputs = [
  41     addOpenGLRunpath
  42     cython
  43   ];
  44
  45   LDFLAGS = "-L${cudatoolkit}/lib/stubs";
  46
  47   propagatedBuildInputs = [
  48     cudatoolkit
  49     cudnn
  50     cutensor
  51     nccl
  52     fastrlock
  53     numpy
  54     setuptools
  55     wheel
  56   ];
  57
  58   nativeCheckInputs = [
  59     pytestCheckHook
  60     mock
  61   ];
  62
  63   # Won't work with the GPU, whose drivers won't be accessible from the build
  64   # sandbox
  65   doCheck = false;
  66
  67   postFixup = ''
  68     find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
  69       addOpenGLRunpath "$lib"
  70     done
  71   '';
  72
  73   enableParallelBuilding = true;
  74
  75   meta = with lib; {
  76     description = "A NumPy-compatible matrix library accelerated by CUDA";
  77     homepage = "https://cupy.chainer.org/";
  78     changelog = "https://github.com/cupy/cupy/releases/tag/v${version}";
  79     license = licenses.mit;
  80     platforms = [ "x86_64-linux" ];
  81     maintainers = with maintainers; [ hyphon81 ];
  82   };
  83 }