diff options
Diffstat (limited to 'pkgs/development/cuda-modules/nccl-tests/default.nix')
-rw-r--r-- | pkgs/development/cuda-modules/nccl-tests/default.nix | 84 |
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; + }; + } +) |