From ebef0c31074009f49a7d2e900deff099137ed0fd Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Thu, 30 Nov 2023 20:39:14 +0000
Subject: cudaPackages.nccl-tests: support building with CUDA < 11.4 with
 cudatoolkit

---
 .../libraries/science/math/nccl/tests.nix          | 131 ++++++++++++---------
 1 file changed, 73 insertions(+), 58 deletions(-)

(limited to 'pkgs')

diff --git a/pkgs/development/libraries/science/math/nccl/tests.nix b/pkgs/development/libraries/science/math/nccl/tests.nix
index 9d826b92f164a..5c2f29b7ed564 100644
--- a/pkgs/development/libraries/science/math/nccl/tests.nix
+++ b/pkgs/development/libraries/science/math/nccl/tests.nix
@@ -1,69 +1,84 @@
-{ backendStdenv
-, config
-, cuda_cccl
-, cuda_cudart
-, cuda_nvcc
-, cudaVersion
-, fetchFromGitHub
-, gitUpdater
-, lib
-, mpi
-, mpiSupport ? false
-, nccl
-, which
+# NOTE: Though NCCL tests is called within the cudaPackages package set, we avoid passing in
+# the names of dependencies from that package set directly to avoid evaluation errors
+# in the case redistributable packages are not available.
+{
+  config,
+  cudaPackages,
+  fetchFromGitHub,
+  gitUpdater,
+  lib,
+  mpi,
+  mpiSupport ? false,
+  which,
 }:
+let
+  inherit (cudaPackages)
+    backendStdenv
+    cuda_cccl
+    cuda_cudart
+    cuda_nvcc
+    cudatoolkit
+    cudaVersion
+    nccl
+    ;
+in
+backendStdenv.mkDerivation (
+  finalAttrs: {
 
-backendStdenv.mkDerivation (finalAttrs: {
+    pname = "nccl-tests";
+    version = "2.13.8";
 
-  pname = "nccl-tests";
-  version = "2.13.8";
+    src = fetchFromGitHub {
+      owner = "NVIDIA";
+      repo = finalAttrs.pname;
+      rev = "v${finalAttrs.version}";
+      hash = "sha256-dxLoflsTHDBnZRTzoXdm30OyKpLlRa73b784YWALBHg=";
+    };
 
-  src = fetchFromGitHub {
-    owner = "NVIDIA";
-    repo = finalAttrs.pname;
-    rev = "v${finalAttrs.version}";
-    hash = "sha256-dxLoflsTHDBnZRTzoXdm30OyKpLlRa73b784YWALBHg=";
-  };
+    strictDeps = true;
 
-  strictDeps = true;
+    nativeBuildInputs =
+      [which]
+      ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
+      ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [cuda_nvcc];
 
-  nativeBuildInputs = [
-    cuda_nvcc
-    which
-  ];
-
-  buildInputs = [
-    cuda_cudart
-    nccl
-  ] ++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [
-    cuda_cccl.dev # <nv/target>
-  ] ++ lib.optional mpiSupport mpi;
+    buildInputs =
+      [nccl]
+      ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
+      ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
+        cuda_nvcc.dev # crt/host_config.h
+        cuda_cudart
+      ]
+      ++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [
+        cuda_cccl.dev # <nv/target>
+      ]
+      ++ lib.optionals mpiSupport [mpi];
 
-  makeFlags = [
-    "CUDA_HOME=${cuda_nvcc}"
-    "NCCL_HOME=${nccl}"
-  ] ++ lib.optionals mpiSupport [
-    "MPI=1"
-  ];
+    makeFlags =
+      ["NCCL_HOME=${nccl}"]
+      ++ lib.optionals (lib.versionOlder cudaVersion "11.4") ["CUDA_HOME=${cudatoolkit}"]
+      ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") ["CUDA_HOME=${cuda_nvcc}"]
+      ++ lib.optionals mpiSupport ["MPI=1"];
 
-  enableParallelBuilding = true;
+    enableParallelBuilding = true;
 
-  installPhase = ''
-    mkdir -p $out/bin
-    cp -r build/* $out/bin/
-  '';
+    installPhase = ''
+      mkdir -p $out/bin
+      cp -r build/* $out/bin/
+    '';
 
-  passthru.updateScript = gitUpdater {
-    inherit (finalAttrs) pname version;
-    rev-prefix = "v";
-  };
+    passthru.updateScript = gitUpdater {
+      inherit (finalAttrs) pname version;
+      rev-prefix = "v";
+    };
 
-  meta = with lib; {
-    description = "Tests to check both the performance and the correctness of NVIDIA NCCL operations";
-    homepage = "https://github.com/NVIDIA/nccl-tests";
-    platforms = platforms.linux;
-    license = licenses.bsd3;
-    broken = !config.cudaSupport || (mpiSupport && mpi == null);
-    maintainers = with maintainers; [ jmillerpdt ] ++ teams.cuda.members;
-  };
-})
+    meta = with lib; {
+      description = "Tests to check both the performance and the correctness of NVIDIA NCCL operations";
+      homepage = "https://github.com/NVIDIA/nccl-tests";
+      platforms = platforms.linux;
+      license = licenses.bsd3;
+      broken = !config.cudaSupport || (mpiSupport && mpi == null);
+      maintainers = with maintainers; [jmillerpdt] ++ teams.cuda.members;
+    };
+  }
+)
-- 
cgit 1.4.1