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