pyrox.dev
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
16let
17 inherit (cudaPackages) cudatoolkit cudnn cutensor nccl;
18in 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}