tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1# NOTE: Though NCCL tests 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 backendStdenv,
6 _cuda,
7 cuda_cccl,
8 cuda_cudart,
9 cuda_nvcc,
10 cudaNamePrefix,
11 fetchFromGitHub,
12 flags,
13 gitUpdater,
14 lib,
15 mpi,
16 mpiSupport ? false,
17 nccl,
18 which,
19}:
20let
21 inherit (_cuda.lib) _mkMetaBroken;
22 inherit (lib) licenses maintainers teams;
23 inherit (lib.attrsets) getBin getInclude getLib;
24 inherit (lib.lists) optionals;
25in
26backendStdenv.mkDerivation (finalAttrs: {
27 __structuredAttrs = true;
28 strictDeps = true;
29
30 # NOTE: Depends on the CUDA package set, so use cudaNamePrefix.
31 name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
32 pname = "nccl-tests";
33 version = "2.17.8";
34
35 src = fetchFromGitHub {
36 owner = "NVIDIA";
37 repo = "nccl-tests";
38 rev = "v${finalAttrs.version}";
39 hash = "sha256-vZ3PP1NbNBv3bvJaiKq5rkIzaT7bcZpHawu3nbSqK60=";
40 };
41
42 postPatch = ''
43 nixLog "patching $PWD/src/common.mk to remove NVIDIA's ccbin declaration"
44 substituteInPlace ./src/common.mk \
45 --replace-fail \
46 '-ccbin $(CXX)' \
47 ""
48 '';
49
50 nativeBuildInputs = [
51 which
52 cuda_nvcc
53 ];
54
55 buildInputs = [
56 cuda_cccl # <nv/target>
57 cuda_cudart
58 nccl
59 ]
60 ++ optionals mpiSupport [ mpi ];
61
62 # NOTE: CUDA_HOME is expected to have the bin directory
63 # TODO: This won't work with cross-compilation since cuda_nvcc will come from hostPackages by default (aka pkgs).
64 makeFlags = [
65 "CXXSTD=-std=c++17"
66 "CUDA_HOME=${getBin cuda_nvcc}"
67 "CUDA_INC=${getInclude cuda_cudart}/include"
68 "CUDA_LIB=${getLib cuda_cudart}/lib"
69 "NVCC_GENCODE=${flags.gencodeString}"
70 "PREFIX=$(out)"
71 ]
72 ++ optionals mpiSupport [ "MPI=1" ];
73
74 enableParallelBuilding = true;
75
76 installPhase = ''
77 runHook preInstall
78 mkdir -p "$out/bin"
79 install -Dm755 \
80 $(find build -type f -executable) \
81 "$out/bin"
82 runHook postInstall
83 '';
84
85 passthru = {
86 brokenAssertions = [
87 {
88 message = "mpi is non-null when mpiSupport is true";
89 assertion = mpiSupport -> mpi != null;
90 }
91 ];
92
93 updateScript = gitUpdater {
94 inherit (finalAttrs) pname version;
95 rev-prefix = "v";
96 };
97 };
98
99 meta = {
100 description = "Tests to check both the performance and the correctness of NVIDIA NCCL operations";
101 homepage = "https://github.com/NVIDIA/nccl-tests";
102 platforms = [
103 "aarch64-linux"
104 "x86_64-linux"
105 ];
106 license = licenses.bsd3;
107 broken = _mkMetaBroken finalAttrs;
108 maintainers = with maintainers; [ jmillerpdt ];
109 teams = [ teams.cuda ];
110 };
111})