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