python3Packages.pytorch: migrate to cudaPackages

pkgs/development/libraries/mpich/default.nix

@@ -42,6 +42,11 @@ stdenv.mkDerivation  rec {
     sed -i 's:FC="gfortran":FC=${gfortran}/bin/gfortran:' $out/bin/mpifort
   '';
 
+  passthru = {
+    cudaSupport = ch4backend.cudaSupport or false;
+    cudaPackages = ch4backend.cudaPackages or { };
+  };
+
   meta = with lib; {
     description = "Implementation of the Message Passing Interface (MPI) standard";
 

pkgs/development/libraries/science/math/magma/default.nix

@@ -1,10 +1,6 @@
-{ lib, stdenv, fetchurl, cmake, gfortran, ninja, cudaPackages, libpthreadstubs, lapack, blas }:
+{ lib, stdenv, fetchurl, cmake, gfortran, ninja, cudaPackages, libpthreadstubs, lapack, blas, symlinkJoin }:
 
-let
-  inherit (cudaPackages) cudatoolkit;
-in
-
-assert let majorIs = lib.versions.major cudatoolkit.version;
+assert let majorIs = cudaPackages.cudaMajorVersion;
        in majorIs == "9" || majorIs == "10" || majorIs == "11";
 
 let
@@ -37,8 +33,22 @@ let
     ];
   };
 
+  inherit (cudaPackages) cudaMajorVersion;
+  inherit (cudaPackages.cudatoolkit) cc;
+
+  cuda_joined = symlinkJoin {
+    name = "cuda-redist-${cudaPackages.cudaVersion}";
+    paths = with cudaPackages; [
+      cuda_nvcc
+      cuda_cudart # cuda_runtime.h
+      libcublas
+      libcusparse
+      cuda_nvprof # <cuda_profiler_api.h>
+    ];
+  };
+
   capabilityString = lib.strings.concatStringsSep ","
-    cudaCapabilities."cuda${lib.versions.major cudatoolkit.version}";
+    cudaCapabilities."cuda${cudaMajorVersion}";
 
 in stdenv.mkDerivation {
   pname = "magma";
@@ -51,14 +61,14 @@ in stdenv.mkDerivation {
 
   nativeBuildInputs = [ gfortran cmake ninja ];
 
-  buildInputs = [ cudatoolkit libpthreadstubs lapack blas ];
+  buildInputs = [ libpthreadstubs lapack blas cuda_joined ];
 
   cmakeFlags = [ "-DGPU_TARGET=${capabilityString}" ];
 
   doCheck = false;
 
   preConfigure = ''
-    export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
+    export CC=${cc}/bin/gcc CXX=${cc}/bin/g++
   '';
 
   enableParallelBuilding=true;
@@ -72,5 +82,7 @@ in stdenv.mkDerivation {
     maintainers = with maintainers; [ tbenst ];
   };
 
-  passthru.cudatoolkit = cudatoolkit;
+  passthru = {
+    inherit cudaPackages;
+  };
 }

pkgs/development/libraries/ucx/default.nix

@@ -1,10 +1,13 @@
 { lib, stdenv, fetchFromGitHub, autoreconfHook, doxygen
 , numactl, rdma-core, libbfd, libiberty, perl, zlib, symlinkJoin
 , enableCuda ? false
-, cudatoolkit
+, cudaPackages
 }:
 
 let
+  # TODO: use the redistributable cuda packages instead
+  inherit (cudaPackages) cudatoolkit;
+
   # Needed for configure to find all libraries
   cudatoolkit' = symlinkJoin {
     inherit (cudatoolkit) name meta;
@@ -43,6 +46,11 @@ in stdenv.mkDerivation rec {
 
   enableParallelBuilding = true;
 
+  passthru = {
+    cudaSupport = enableCuda;
+    inherit cudaPackages;
+  };
+
   meta = with lib; {
     description = "Unified Communication X library";
     homepage = "http://www.openucx.org";

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

@@ -30,23 +30,60 @@
   isPy3k, pythonOlder }:
 
 let
-  inherit (cudaPackages) cudatoolkit cudnn nccl;
+  inherit (cudaPackages) cudnn nccl;
 in
 
 # assert that everything needed for cuda is present and that the correct cuda versions are used
-assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
+assert cudaSupport -> (let majorIs = cudaPackages.cudaMajorVersion;
                         in majorIs == "9" || majorIs == "10" || majorIs == "11");
 
-# confirm that cudatoolkits are sync'd across dependencies
-assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
-assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
+# We expect referential equality of all cudaPackages used to ensure consistency
+# You can make an overlay and pass the same cudaPackages to pytorch, mpi, and magma
+# TODO: `==` is an implementation detail; move comparison logic to cudaPackages
+assert (MPISupport && cudaSupport) -> mpi.cudaPackages == cudaPackages;
+assert cudaSupport -> (magma.cudaPackages == cudaPackages);
 
 let
   setBool = v: if v then "1" else "0";
+
+  # cudatoolkit.cc is a passthru attribute
+  # that points at a compatible gcc version,
+  # we're not really using cudatoolkit;
+  # not sure we really need to hard-code gcc,
+  # but this is out of scope of current PR
+  # (which is: migrating pytorch to redist cuda)
+  inherit (cudaPackages.cudatoolkit) cc;
+
   cudatoolkit_joined = symlinkJoin {
-    name = "${cudatoolkit.name}-unsplit";
+    name = "cuda-redist-${cudaPackages.cudaVersion}";
+
+    # The list of required cuda redist packages can be found e.g.
+    # at https://github.com/pytorch/pytorch/blob/b09769992f83f94150eaef2ab9d03c37b36da159/cmake/Summary.cmake#L85
+    # Not explicitly listed there are: nvml, nvtx, cccl
+    #
     # nccl is here purely for semantic grouping it could be moved to nativeBuildInputs
-    paths = [ cudatoolkit.out cudatoolkit.lib nccl.dev nccl.out ];
+    paths = with cudaPackages; [
+      cuda_nvcc
+      cuda_nvml_dev # <nvml.h>
+      cuda_nvtx # -llibNVToolsExt
+      cuda_cccl # <thrust/*>
+      cuda_nvprof # <cuda_profiler_api.h>
+      cuda_nvrtc
+      cuda_cudart # cuda_runtime.h
+      libcublas
+      libcufft
+      libcusolver
+      libcusparse
+      libcurand
+      cuda_cupti
+      (lib.getDev nccl)
+      nccl
+    ];
+
+    # ld is going to look for static archives (e.g. libcudart_static.a) in lib64
+    postBuild = ''
+      ln -s $out/lib $out/lib64
+    '';
   };
 
   # Give an explicit list of supported architectures for the build, See:
@@ -105,15 +142,15 @@ let
   final_cudaArchList =
     if !cudaSupport || cudaArchList != null
     then cudaArchList
-    else cudaCapabilities."cuda${lib.versions.major cudatoolkit.version}";
+    else cudaCapabilities."cuda${cudaPackages.cudaMajorVersion}";
 
   # Normally libcuda.so.1 is provided at runtime by nvidia-x11 via
   # LD_LIBRARY_PATH=/run/opengl-driver/lib.  We only use the stub
   # libcuda.so from cudatoolkit for running tests, so that we don’t have
   # to recompile pytorch on every update to nvidia-x11 or the kernel.
   cudaStub = linkFarm "cuda-stub" [{
     name = "libcuda.so.1";
-    path = "${cudatoolkit}/lib/stubs/libcuda.so";
+    path = "${cudaPackages.cuda_cudart}/lib/stubs/libcuda.so";
   }];
   cudaStubEnv = lib.optionalString cudaSupport
     "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH ";
@@ -154,9 +191,11 @@ in buildPythonPackage rec {
     ./pthreadpool-disable-gcd.diff
   ];
 
+  # CUDAHOSTCXX goes into nvcc's -ccbin argument
   preConfigure = lib.optionalString cudaSupport ''
     export TORCH_CUDA_ARCH_LIST="${lib.strings.concatStringsSep ";" final_cudaArchList}"
-    export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
+    export CUDAHOSTCXX=${cc}/bin/
+    export CC=${cc}/bin/gcc CXX=${cc}/bin/g++
   '' + lib.optionalString (cudaSupport && cudnn != null) ''
     export CUDNN_INCLUDE_DIR=${cudnn}/include
   '';