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