python3Packages.cupy: add aarch64-linux support (#378958)

pkgs/development/cuda-modules/cudnn/releases.nix

@@ -7,7 +7,7 @@
       {
         version = "8.9.5.30";
         minCudaVersion = "12.0";
-        maxCudaVersion = "12.2";
+        maxCudaVersion = "12.4";
         url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-aarch64/cudnn-linux-aarch64-8.9.5.30_cuda12-archive.tar.xz";
         hash = "sha256-BJH3sC9VwiB362eL8xTB+RdSS9UHz1tlgjm/mKRyM6E=";
       }
@@ -75,7 +75,7 @@
       {
         version = "8.9.7.29";
         minCudaVersion = "12.0";
-        maxCudaVersion = "12.2";
+        maxCudaVersion = "12.4";
         url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.9.7.29_cuda12-archive.tar.xz";
         hash = "sha256-6Yt8gAEHheXVygHuTOm1sMjHNYfqb4ZIvjTT+NHUe9E=";
       }
@@ -183,7 +183,7 @@
       {
         version = "8.9.7.29";
         minCudaVersion = "12.0";
-        maxCudaVersion = "12.2";
+        maxCudaVersion = "12.4";
         url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz";
         hash = "sha256-R1MzYlx+QqevPKCy91BqEG4wyTsaoAgc2cE++24h47s=";
       }

pkgs/development/python-modules/cupy/default.nix

@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
   cython_0,
@@ -13,31 +14,46 @@
   addDriverRunpath,
   pythonOlder,
   symlinkJoin,
-  fetchpatch
 }:
 
 let
-  inherit (cudaPackages) cudnn cutensor nccl;
-  outpaths = with cudaPackages; [
-      cuda_cccl # <nv/target>
-      cuda_cudart
-      cuda_nvcc # <crt/host_defines.h>
-      cuda_nvprof
-      cuda_nvrtc
-      cuda_nvtx
-      cuda_profiler_api
-      libcublas
-      libcufft
-      libcurand
-      libcusolver
-      libcusparse
+  inherit (cudaPackages) cudnn;
 
-      # Missing:
-      # cusparselt
+  shouldUsePkg =
+    pkg: if pkg != null && lib.meta.availableOn stdenv.hostPlatform pkg then pkg else null;
+
+  # some packages are not available on all platforms
+  cuda_nvprof = shouldUsePkg (cudaPackages.nvprof or null);
+  cutensor = shouldUsePkg (cudaPackages.cutensor or null);
+  nccl = shouldUsePkg (cudaPackages.nccl or null);
+
+  outpaths = with cudaPackages; [
+    cuda_cccl # <nv/target>
+    cuda_cudart
+    cuda_nvcc # <crt/host_defines.h>
+    cuda_nvprof
+    cuda_nvrtc
+    cuda_nvtx
+    cuda_profiler_api
+    libcublas
+    libcufft
+    libcurand
+    libcusolver
+    libcusparse
+
+    # Missing:
+    # cusparselt
   ];
   cudatoolkit-joined = symlinkJoin {
     name = "cudatoolkit-joined-${cudaPackages.cudaVersion}";
-    paths = outpaths ++ lib.concatMap (f: lib.map f outpaths) [lib.getLib lib.getDev (lib.getOutput "static") (lib.getOutput "stubs")];
+    paths =
+      outpaths
+      ++ lib.concatMap (f: lib.map f outpaths) [
+        lib.getLib
+        lib.getDev
+        (lib.getOutput "static")
+        (lib.getOutput "stubs")
+      ];
   };
 in
 buildPythonPackage rec {
@@ -47,6 +63,8 @@ buildPythonPackage rec {
 
   disabled = pythonOlder "3.7";
 
+  stdenv = cudaPackages.backendStdenv;
+
   src = fetchFromGitHub {
     owner = "cupy";
     repo = "cupy";
@@ -55,14 +73,6 @@ buildPythonPackage rec {
     fetchSubmodules = true;
   };
 
-  patches = [
-    (fetchpatch {
-      url =
-        "https://github.com/cfhammill/cupy/commit/67526c756e4a0a70f0420bf0e7f081b8a35a8ee5.patch";
-      hash = "sha256-WZgexBdM9J0ep5s+9CGZriVq0ZidCRccox+g0iDDywQ=";
-    })
-  ];
-
   # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both
   # CUPY_NUM_BUILD_JOBS and CUPY_NUM_NVCC_THREADS to NIX_BUILD_CORES results in
   # a small amount of thrashing but it turns out there are a large number of
@@ -118,7 +128,10 @@ buildPythonPackage rec {
     homepage = "https://cupy.chainer.org/";
     changelog = "https://github.com/cupy/cupy/releases/tag/v${version}";
     license = licenses.mit;
-    platforms = [ "x86_64-linux" ];
+    platforms = [
+      "aarch64-linux"
+      "x86_64-linux"
+    ];
     maintainers = with maintainers; [ hyphon81 ];
   };
 }

pkgs/top-level/python-packages.nix

@@ -2832,8 +2832,7 @@ self: super: with self; {
   cufflinks = callPackage ../development/python-modules/cufflinks { };
 
   cupy = callPackage ../development/python-modules/cupy {
-    # cupy 12.2.0 possibly incompatible with cutensor 2.0 that comes with cudaPackages_12
-    cudaPackages = pkgs.cudaPackages_11.overrideScope (cu-fi: _: {
+    cudaPackages = pkgs.cudaPackages.overrideScope (cu-fi: _: {
       # CuDNN 9 is not supported:
       # https://github.com/cupy/cupy/issues/8215
       cudnn = cu-fi.cudnn_8_9;