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