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