pyrox.dev
1{ lib, buildPythonPackage
2, fetchPypi, isPy3k, cython
3, fastrlock, numpy, six, wheel, pytestCheckHook, mock, setuptools
4, cudatoolkit, cudnn, cutensor, nccl
5, addOpenGLRunpath
6}:
7
8buildPythonPackage rec {
9 pname = "cupy";
10 version = "9.6.0";
11 disabled = !isPy3k;
12
13 src = fetchPypi {
14 inherit pname version;
15 sha256 = "22469ea1ad51ffbb4af2b139ed0820ac5d0b78f1265b2a095ed5e5d5299aab91";
16 };
17
18 preConfigure = ''
19 export CUDA_PATH=${cudatoolkit}
20 '';
21
22 nativeBuildInputs = [
23 addOpenGLRunpath
24 cython
25 ];
26
27 LDFLAGS = "-L${cudatoolkit}/lib/stubs";
28
29 propagatedBuildInputs = [
30 cudatoolkit
31 cudnn
32 cutensor
33 nccl
34 fastrlock
35 numpy
36 six
37 setuptools
38 wheel
39 ];
40
41 checkInputs = [
42 pytestCheckHook
43 mock
44 ];
45
46 # Won't work with the GPU, whose drivers won't be accessible from the build
47 # sandbox
48 doCheck = false;
49
50 postFixup = ''
51 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
52 addOpenGLRunpath "$lib"
53 done
54 '';
55
56 enableParallelBuilding = true;
57
58 meta = with lib; {
59 description = "A NumPy-compatible matrix library accelerated by CUDA";
60 homepage = "https://cupy.chainer.org/";
61 license = licenses.mit;
62 platforms = [ "x86_64-linux" ];
63 maintainers = with maintainers; [ hyphon81 ];
64 };
65}