Merge pull request #169279 from samuela/samuela/fix-jaxlibwithcuda

Fix python3Packages.jaxlibWithCuda

pkgs/development/compilers/cudatoolkit/versions.toml

@@ -44,16 +44,16 @@ gcc = "gcc9"
 version = "11.4.2"
 url = "https://developer.download.nvidia.com/compute/cuda/11.4.2/local_installers/cuda_11.4.2_470.57.02_linux.run"
 sha256 = "sha256-u9h8oOkT+DdFSnljZ0c1E83e9VUILk2G7Zo4ZZzIHwo="
-gcc = "gcc10" # can bump to 11 along with stdenv.cc
+gcc = "gcc11"
 
 ["11.5"]
 version = "11.5.0"
 url = "https://developer.download.nvidia.com/compute/cuda/11.5.0/local_installers/cuda_11.5.0_495.29.05_linux.run"
 sha256 = "sha256-rgoWk9lJfPPYHmlIlD43lGNpANtxyY1Y7v2sr38aHkw="
-gcc = "gcc10" # can bump to 11 along with stdenv.cc
+gcc = "gcc11"
 
 ["11.6"]
 version = "11.6.1"
 url = "https://developer.download.nvidia.com/compute/cuda/11.6.1/local_installers/cuda_11.6.1_510.47.03_linux.run"
 sha256 = "sha256-qyGa/OALdCABEyaYZvv/derQN7z8I1UagzjCaEyYTX4="
-gcc = "gcc10" # can bump to 11 along with stdenv.cc
+gcc = "gcc11"

pkgs/top-level/all-packages.nix

@@ -3579,7 +3579,12 @@ with pkgs;
     inherit (darwin.apple_sdk.frameworks) Security;
   };
 
-  gpu-burn = callPackage ../applications/misc/gpu-burn { };
+  gpu-burn = callPackage ../applications/misc/gpu-burn {
+    # gpu-burn doesn't build on gcc11. CUDA 11.3 is the last version to use
+    # pre-gcc11, in particular gcc9.
+    cudatoolkit = cudaPackages_11_3.cudatoolkit;
+    stdenv = gcc9Stdenv;
+  };
 
   greg = callPackage ../applications/audio/greg {
     pythonPackages = python3Packages;
@@ -23097,6 +23102,10 @@ with pkgs;
 
   librealsenseWithCuda = callPackage ../development/libraries/librealsense {
     cudaSupport = true;
+    # librealsenseWithCuda doesn't build on gcc11. CUDA 11.3 is the last version
+    # to use pre-gcc11, in particular gcc9.
+    cudaPackages = cudaPackages_11_3;
+    stdenv = gcc9Stdenv;
   };
 
   librealsenseWithoutCuda = callPackage ../development/libraries/librealsense {

pkgs/top-level/python-packages.nix

@@ -1958,7 +1958,11 @@ in {
 
   cufflinks = callPackage ../development/python-modules/cufflinks { };
 
-  cupy = callPackage ../development/python-modules/cupy { };
+  cupy = callPackage ../development/python-modules/cupy {
+    # cupy doesn't build on gcc11. CUDA 11.3 is the last version to use
+    # pre-gcc11, in particular gcc9.
+    cudaPackages = pkgs.cudaPackages_11_3;
+  };
 
   curio = callPackage ../development/python-modules/curio { };
 
@@ -4243,20 +4247,23 @@ in {
 
   jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
     cudaSupport = pkgs.config.cudaSupport or false;
-    inherit (self.tensorflow) cudaPackages;
+    # At the time of writing (2022-04-18), `cudaPackages.nccl` is broken, so we
+    # pin to `cudaPackages_11_6` instead.
+    cudaPackages = pkgs.cudaPackages_11_6;
   };
 
   jaxlib-build = callPackage ../development/python-modules/jaxlib {
     # Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.
     cudaSupport = pkgs.config.cudaSupport or false;
-    inherit (self.tensorflow) cudaPackages;
+    # At the time of writing (2022-04-18), `cudaPackages.nccl` is broken, so we
+    # pin to `cudaPackages_11_6` instead.
+    cudaPackages = pkgs.cudaPackages_11_6;
   };
 
   jaxlib = self.jaxlib-build;
 
   jaxlibWithCuda = self.jaxlib-build.override {
     cudaSupport = true;
-
   };
 
   jaxlibWithoutCuda = self.jaxlib-build.override {