diff --git a/pkgs/development/libraries/mpich/default.nix b/pkgs/development/libraries/mpich/default.nix index 60e543c462960..ad480e134cd20 100644 --- a/pkgs/development/libraries/mpich/default.nix +++ b/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"; diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix index 05d7d4fa18424..27070b655611e 100644 --- a/pkgs/development/libraries/science/math/magma/default.nix +++ b/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 # + ]; + }; + 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; + }; } diff --git a/pkgs/development/libraries/ucx/default.nix b/pkgs/development/libraries/ucx/default.nix index 0dc79d47dda7b..96dad0e4b3570 100644 --- a/pkgs/development/libraries/ucx/default.nix +++ b/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"; diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix index 810460a30adea..ff92c854a9a87 100644 --- a/pkgs/development/python-modules/pytorch/default.nix +++ b/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 # + cuda_nvtx # -llibNVToolsExt + cuda_cccl # + cuda_nvprof # + 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,7 +142,7 @@ 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 @@ -113,7 +150,7 @@ let # 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 '';