Merge pull request #197157 from tpwrules/cutensor-fixes

python3Packages.cupy: upgrade and fix build

4 files changed, 11 insertions, 9 deletions

diff --git a/pkgs/development/libraries/science/math/cutensor/generic.nix b/pkgs/development/libraries/science/math/cutensor/generic.nix
index e4cbf0ab7bf72..c957fcdd99d4e 100644
--- a/pkgs/development/libraries/science/math/cutensor/generic.nix
+++ b/pkgs/development/libraries/science/math/cutensor/generic.nix
@@ -13,6 +13,7 @@
 let
   mostOfVersion = builtins.concatStringsSep "."
     (lib.take 3 (lib.versions.splitVersion version));
+  platform = "${stdenv.hostPlatform.parsed.kernel.name}-${stdenv.hostPlatform.parsed.cpu.name}";
 in
 
 stdenv.mkDerivation {
@@ -20,7 +21,9 @@ stdenv.mkDerivation {
   inherit version;
 
   src = fetchurl {
-    url = "https://developer.download.nvidia.com/compute/cutensor/${mostOfVersion}/local_installers/libcutensor-${stdenv.hostPlatform.parsed.kernel.name}-${stdenv.hostPlatform.parsed.cpu.name}-${version}.tar.gz";
+    url = if lib.versionOlder mostOfVersion "1.3.3"
+      then "https://developer.download.nvidia.com/compute/cutensor/${mostOfVersion}/local_installers/libcutensor-${platform}-${version}.tar.gz"
+      else "https://developer.download.nvidia.com/compute/cutensor/redist/libcutensor/${platform}/libcutensor-${platform}-${version}-archive.tar.xz";
     inherit hash;
   };
 
diff --git a/pkgs/development/python-modules/cupy/default.nix b/pkgs/development/python-modules/cupy/default.nix
index aba974e027045..9bc7efafff54c 100644
--- a/pkgs/development/python-modules/cupy/default.nix
+++ b/pkgs/development/python-modules/cupy/default.nix
@@ -9,12 +9,12 @@ let
   inherit (cudaPackages) cudatoolkit cudnn cutensor nccl;
 in buildPythonPackage rec {
   pname = "cupy";
-  version = "11.1.0";
+  version = "11.2.0";
   disabled = !isPy3k;
 
   src = fetchPypi {
     inherit pname version;
-    sha256 = "sha256-4TtvmQrd172HCQWvQp+tUQhNCFf4YA3TIGod9aRoTt0=";
+    sha256 = "sha256-wzNh8RejR6Y/aZbql0RtF/HAOPGh9TPlAkZCNQdpI+I=";
   };
 
   # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both
@@ -70,8 +70,5 @@ in buildPythonPackage rec {
     license = licenses.mit;
     platforms = [ "x86_64-linux" ];
     maintainers = with maintainers; [ hyphon81 ];
-
-    # See https://github.com/NixOS/nixpkgs/pull/179912#issuecomment-1206265922.
-    broken = true;
   };
 }
diff --git a/pkgs/test/cuda/cuda-library-samples/generic.nix b/pkgs/test/cuda/cuda-library-samples/generic.nix
index f0c30d801a1c9..e01664bab3191 100644
--- a/pkgs/test/cuda/cuda-library-samples/generic.nix
+++ b/pkgs/test/cuda/cuda-library-samples/generic.nix
@@ -54,6 +54,8 @@ in
 
     src = "${src}/cuTENSOR";
 
+    buildInputs = [ cutensor ];
+
     cmakeFlags = [
       "-DCUTENSOR_EXAMPLE_BINARY_INSTALL_DIR=${builtins.placeholder "out"}/bin"
     ];
diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix
index 70d57672a0f56..c57e486406e1f 100644
--- a/pkgs/top-level/cuda-packages.nix
+++ b/pkgs/top-level/cuda-packages.nix
@@ -28,15 +28,15 @@ let
       "1.2.2.5" = {
         hash = "sha256-lU7iK4DWuC/U3s1Ct/rq2Gr3w4F2U7RYYgpmF05bibY=";
       };
-      "1.3.1.3" = {
-        hash = "sha256-mNlVnabB2IC3HnYY0mb06RLqQzDxN9ePGVeBy3hkBC8=";
+      "1.5.0.3" = {
+        hash = "sha256-T96+lPC6OTOkIs/z3QWg73oYVSyidN0SVkBWmT9VRx0=";
       };
     };
 
     inherit (final) cudaMajorMinorVersion cudaMajorVersion;
 
     cutensor = buildCuTensorPackage rec {
-      version = if cudaMajorMinorVersion == "10.1" then "1.2.2.5" else "1.3.1.3";
+      version = if cudaMajorMinorVersion == "10.1" then "1.2.2.5" else "1.5.0.3";
       inherit (cuTensorVersions.${version}) hash;
       # This can go into generic.nix
       libPath = "lib/${if cudaMajorVersion == "10" then cudaMajorMinorVersion else cudaMajorVersion}";