pkgs/development/cuda-modules/nccl/default.nix

   1 # NOTE: Though NCCL is called within the cudaPackages package set, we avoid passing in
   2 # the names of dependencies from that package set directly to avoid evaluation errors
   3 # in the case redistributable packages are not available.
   4 {
   5   lib,
   6   fetchFromGitHub,
   7   python3,
   8   which,
   9   autoAddDriverRunpath,
  10   cudaPackages,
  11   # passthru.updateScript
  12   gitUpdater,
  13 }:
  14 let
  15   inherit (cudaPackages)
  16     backendStdenv
  17     cuda_cccl
  18     cuda_cudart
  19     cuda_nvcc
  20     cudaAtLeast
  21     cudaFlags
  22     cudaOlder
  23     cudatoolkit
  24     ;
  25 in
  26 backendStdenv.mkDerivation (finalAttrs: {
  27   pname = "nccl";
  28   version = "2.21.5-1";
  29
  30   src = fetchFromGitHub {
  31     owner = "NVIDIA";
  32     repo = "nccl";
  33     rev = "v${finalAttrs.version}";
  34     hash = "sha256-IF2tILwW8XnzSmfn7N1CO7jXL95gUp02guIW5n1eaig=";
  35   };
  36
  37   __structuredAttrs = true;
  38   strictDeps = true;
  39
  40   outputs = [
  41     "out"
  42     "dev"
  43   ];
  44
  45   nativeBuildInputs =
  46     [
  47       which
  48       autoAddDriverRunpath
  49       python3
  50     ]
  51     ++ lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
  52     ++ lib.optionals (cudaAtLeast "11.4") [ cuda_nvcc ];
  53
  54   buildInputs =
  55     lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
  56     ++ lib.optionals (cudaAtLeast "11.4") [
  57       cuda_nvcc # crt/host_config.h
  58       cuda_cudart
  59     ]
  60     # NOTE: CUDA versions in Nixpkgs only use a major and minor version. When we do comparisons
  61     # against other version, like below, it's important that we use the same format. Otherwise,
  62     # we'll get incorrect results.
  63     # For example, lib.versionAtLeast "12.0" "12.0.0" == false.
  64     ++ lib.optionals (cudaAtLeast "12.0") [ cuda_cccl ];
  65
  66   env.NIX_CFLAGS_COMPILE = toString [ "-Wno-unused-function" ];
  67
  68   postPatch = ''
  69     patchShebangs ./src/device/generate.py
  70   '';
  71
  72   makeFlags =
  73     [
  74       "PREFIX=$(out)"
  75       "NVCC_GENCODE=${cudaFlags.gencodeString}"
  76     ]
  77     ++ lib.optionals (cudaOlder "11.4") [
  78       "CUDA_HOME=${cudatoolkit}"
  79       "CUDA_LIB=${lib.getLib cudatoolkit}/lib"
  80       "CUDA_INC=${lib.getDev cudatoolkit}/include"
  81     ]
  82     ++ lib.optionals (cudaAtLeast "11.4") [
  83       "CUDA_HOME=${cuda_nvcc}"
  84       "CUDA_LIB=${lib.getLib cuda_cudart}/lib"
  85       "CUDA_INC=${lib.getDev cuda_cudart}/include"
  86     ];
  87
  88   enableParallelBuilding = true;
  89
  90   postFixup = ''
  91     moveToOutput lib/libnccl_static.a $dev
  92   '';
  93
  94   passthru.updateScript = gitUpdater {
  95     inherit (finalAttrs) pname version;
  96     rev-prefix = "v";
  97   };
  98
  99   meta = with lib; {
 100     description = "Multi-GPU and multi-node collective communication primitives for NVIDIA GPUs";
 101     homepage = "https://developer.nvidia.com/nccl";
 102     license = licenses.bsd3;
 103     platforms = platforms.linux;
 104     # NCCL is not supported on Jetson, because it does not use NVLink or PCI-e for inter-GPU communication.
 105     # https://forums.developer.nvidia.com/t/can-jetson-orin-support-nccl/232845/9
 106     badPlatforms = lib.optionals cudaFlags.isJetsonBuild [ "aarch64-linux" ];
 107     maintainers =
 108       with maintainers;
 109       [
 110         mdaiter
 111         orivej
 112       ]
 113       ++ teams.cuda.members;
 114   };
 115 })