From f67ff2c2ff3e1c1fa1a31ef1a79670e4713291ae Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Tue, 23 Mar 2021 04:25:41 +0000
Subject: [PATCH] python3Packages.cupy: 8.40 -> 8.5.0

While we are at it, also switch to using CUDA 11, and properly use the
cuda toolkit stubs instead of improperly linking a specific version of
the Nvidia drivers.
---
 .../python-modules/cupy/default.nix           | 40 +++++++++++++------
 pkgs/top-level/python-packages.nix            |  8 ++--
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/pkgs/development/python-modules/cupy/default.nix b/pkgs/development/python-modules/cupy/default.nix
index d4b42ac83e83..f5d262f0b8fd 100644
--- a/pkgs/development/python-modules/cupy/default.nix
+++ b/pkgs/development/python-modules/cupy/default.nix
@@ -1,33 +1,35 @@
 { lib, buildPythonPackage
-, fetchPypi, isPy3k, linuxPackages
-, fastrlock, numpy, six, wheel, pytest, mock, setuptools
+, fetchPypi, isPy3k, cython
+, fastrlock, numpy, six, wheel, pytestCheckHook, mock, setuptools
 , cudatoolkit, cudnn, cutensor, nccl
+, addOpenGLRunpath
 }:
 
 buildPythonPackage rec {
   pname = "cupy";
-  version = "8.4.0";
+  version = "8.5.0";
   disabled = !isPy3k;
 
   src = fetchPypi {
     inherit pname version;
-    sha256 = "58d19af6b2e83388d4f0f6ca4226bae4b947920d2ca4951c2eddc8bc78abf66b";
+    sha256 = "fb3f8d3b3454beb249b9880502a45fe493c5a44efacc4c72914cbe1a5dbdf803";
   };
 
-  checkInputs = [
-    pytest
-    mock
+  preConfigure = ''
+    export CUDA_PATH=${cudatoolkit}
+  '';
+
+  nativeBuildInputs = [
+    addOpenGLRunpath
+    cython
   ];
 
-  preConfigure = ''
-      export CUDA_PATH=${cudatoolkit}
-  '';
+  LDFLAGS = "-L${cudatoolkit}/lib/stubs";
 
   propagatedBuildInputs = [
     cudatoolkit
     cudnn
     cutensor
-    linuxPackages.nvidia_x11
     nccl
     fastrlock
     numpy
@@ -36,8 +38,20 @@ buildPythonPackage rec {
     wheel
   ];
 
-  # In python3, test was failed...
-  doCheck = !isPy3k;
+  checkInputs = [
+    pytestCheckHook
+    mock
+  ];
+
+  # Won't work with the GPU, whose drivers won't be accessible from the build
+  # sandbox
+  doCheck = false;
+
+  postFixup = ''
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      addOpenGLRunpath "$lib"
+    done
+  '';
 
   enableParallelBuilding = true;
 
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 2b35c242edb3..42c373ebdf96 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -1602,10 +1602,10 @@ in {
   cufflinks = callPackage ../development/python-modules/cufflinks { };
 
   cupy = callPackage ../development/python-modules/cupy {
-    cudatoolkit = pkgs.cudatoolkit_10_0;
-    cudnn = pkgs.cudnn_cudatoolkit_10_0;
-    nccl = pkgs.nccl_cudatoolkit_10;
-    cutensor = pkgs.cutensor_cudatoolkit_10;
+    cudatoolkit = pkgs.cudatoolkit_11;
+    cudnn = pkgs.cudnn_cudatoolkit_11;
+    nccl = pkgs.nccl_cudatoolkit_11;
+    cutensor = pkgs.cutensor_cudatoolkit_11;
   };
 
   curio = callPackage ../development/python-modules/curio { };