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

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