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

# 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.9";

  src = fetchFromGitHub {
    owner = "NVIDIA";
    repo = finalAttrs.pname;
    rev = "v${finalAttrs.version}";
    hash = "sha256-QYuMBPhvHHVo2ku14jD1CVINLPW0cyiXJkXxb77IxbE=";
  };

  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;
  };
})