1 files changed, 84 insertions, 0 deletions
diff --git a/pkgs/development/cuda-modules/nccl-tests/default.nix b/pkgs/development/cuda-modules/nccl-tests/default.nix
new file mode 100644
index 0000000000000..5c2f29b7ed564
--- /dev/null
+++ b/pkgs/development/cuda-modules/nccl-tests/default.nix
@@ -0,0 +1,84 @@
+# 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: {
+
+    pname = "nccl-tests";
+    version = "2.13.8";
+
+    src = fetchFromGitHub {
+      owner = "NVIDIA";
+      repo = finalAttrs.pname;
+      rev = "v${finalAttrs.version}";
+      hash = "sha256-dxLoflsTHDBnZRTzoXdm30OyKpLlRa73b784YWALBHg=";
+    };
+
+    strictDeps = true;
+
+    nativeBuildInputs =
+      [which]
+      ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
+      ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [cuda_nvcc];
+
+    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 =
+      ["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;
+
+    installPhase = ''
+      mkdir -p $out/bin
+      cp -r build/* $out/bin/
+    '';
+
+    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;
+    };
+  }
+)