1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
inputs@{
autoconf,
automake,
config,
cudaPackages,
fetchFromGitHub,
lib,
libtool,
stdenv,
ucx,
# Configuration options
enableAvx ? stdenv.hostPlatform.avxSupport,
enableCuda ? config.cudaSupport,
enableSse41 ? stdenv.hostPlatform.sse4_1Support,
enableSse42 ? stdenv.hostPlatform.sse4_2Support,
}:
let
inherit (lib.attrsets) getOutput;
inherit (lib.lists) optionals;
inherit (lib.strings) concatStringsSep;
inherit (cudaPackages)
cccl
cuda_cudart
cuda_nvcc
cuda_nvml_dev
flags
nccl
;
stdenv = throw "Use effectiveStdenv instead";
effectiveStdenv = if enableCuda then cudaPackages.backendStdenv else inputs.stdenv;
in
effectiveStdenv.mkDerivation (finalAttrs: {
__structuredAttrs = true;
# TODO(@connorbaker):
# When strictDeps is enabled, `cuda_nvcc` is required as the argument to `--with-cuda` in `configureFlags` or else
# configurePhase fails with `checking for cuda_runtime.h... no`.
# This is odd, especially given `cuda_runtime.h` is provided by `cuda_cudart.dev`, which is already in `buildInputs`.
strictDeps = true;
pname = "ucc";
version = "1.8.0";
src = fetchFromGitHub {
owner = "openucx";
repo = "ucc";
tag = "v${finalAttrs.version}";
hash = "sha256-rH/bG0pYO4VKfrnkXLXy/67hjaz6i3R0YoYdsGBqPsU=";
};
outputs = [
"out"
"dev"
];
enableParallelBuilding = true;
# NOTE: We use --replace-quiet because not all Makefile.am files contain /bin/bash.
postPatch = ''
for comp in $(find src/components -name Makefile.am); do
substituteInPlace "$comp" \
--replace-quiet \
"/bin/bash" \
"${effectiveStdenv.shell}"
done
'';
nativeBuildInputs = [
autoconf
automake
libtool
]
++ optionals enableCuda [ cuda_nvcc ];
buildInputs = [
ucx
]
++ optionals enableCuda [
cccl
cuda_cudart
cuda_nvml_dev
nccl
];
# NOTE: With `__structuredAttrs` enabled, `LDFLAGS` must be set under `env` so it is assured to be a string;
# otherwise, we might have forgotten to convert it to a string and Nix would make LDFLAGS a shell variable
# referring to an array!
env.LDFLAGS = toString (
optionals enableCuda [
# Fake libcuda.so (the real one is deployed impurely)
"-L${getOutput "stubs" cuda_cudart}/lib/stubs"
# Fake libnvidia-ml.so (the real one is deployed impurely)
"-L${getOutput "stubs" cuda_nvml_dev}/lib/stubs"
]
);
preConfigure = ''
./autogen.sh
'';
configureFlags =
optionals enableSse41 [ "--with-sse41" ]
++ optionals enableSse42 [ "--with-sse42" ]
++ optionals enableAvx [ "--with-avx" ]
++ optionals enableCuda [
"--with-cuda=${cuda_nvcc}"
"--with-nvcc-gencode=${concatStringsSep " " flags.gencode}"
];
postInstall = ''
find "$out/lib/" -name "*.la" -exec rm -f \{} \;
moveToOutput bin/ucc_info "$dev"
'';
meta = {
description = "Collective communication operations API";
homepage = "https://openucx.github.io/ucc/";
mainProgram = "ucc_info";
license = lib.licenses.bsd3;
maintainers = [ lib.maintainers.markuskowa ];
platforms = lib.platforms.linux;
};
})
|
