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