tjh.dev
cudaPackages.nccl: switch to cudaAtLeast, cudaOlder, and __structuredAttrs
+16
-14
pkgs/development/cuda-modules/nccl/default.nix
+16
-14
pkgs/development/cuda-modules/nccl/default.nix
···
17
17
cuda_cccl
18
18
cuda_cudart
19
19
cuda_nvcc
20
+
cudaAtLeast
20
21
cudaFlags
22
+
cudaOlder
21
23
cudatoolkit
22
-
cudaVersion
23
24
;
24
25
in
25
26
backendStdenv.mkDerivation (finalAttrs: {
···
33
34
hash = "sha256-ModIjD6RaRD/57a/PA1oTgYhZsAQPrrvhl5sNVXnO6c=";
34
35
};
35
36
37
+
__structuredAttrs = true;
36
38
strictDeps = true;
37
39
38
40
outputs = [
···
46
48
autoAddDriverRunpath
47
49
python3
48
50
]
49
-
++ lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
50
-
++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [ cuda_nvcc ];
51
+
++ lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
52
+
++ lib.optionals (cudaAtLeast "11.4") [ cuda_nvcc ];
51
53
52
54
buildInputs =
53
-
lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
54
-
++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
55
+
lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
56
+
++ lib.optionals (cudaAtLeast "11.4") [
55
57
cuda_nvcc.dev # crt/host_config.h
56
58
cuda_cudart
57
59
]
···
59
61
# against other version, like below, it's important that we use the same format. Otherwise,
60
62
# we'll get incorrect results.
61
63
# For example, lib.versionAtLeast "12.0" "12.0.0" == false.
62
-
++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [ cuda_cccl ];
64
+
++ lib.optionals (cudaAtLeast "12.0") [ cuda_cccl ];
63
65
64
66
env.NIX_CFLAGS_COMPILE = toString [ "-Wno-unused-function" ];
65
67
66
-
preConfigure = ''
68
+
postPatch = ''
67
69
patchShebangs ./src/device/generate.py
68
-
makeFlagsArray+=(
69
-
"NVCC_GENCODE=${lib.concatStringsSep " " cudaFlags.gencode}"
70
-
)
71
70
'';
72
71
73
-
makeFlags =
74
-
[ "PREFIX=$(out)" ]
75
-
++ lib.optionals (lib.versionOlder cudaVersion "11.4") [
72
+
makeFlagsArray =
73
+
[
74
+
"PREFIX=$(out)"
75
+
"NVCC_GENCODE=${cudaFlags.gencodeString}"
76
+
]
77
+
++ lib.optionals (cudaOlder "11.4") [
76
78
"CUDA_HOME=${cudatoolkit}"
77
79
"CUDA_LIB=${lib.getLib cudatoolkit}/lib"
78
80
"CUDA_INC=${lib.getDev cudatoolkit}/include"
79
81
]
80
-
++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
82
+
++ lib.optionals (cudaAtLeast "11.4") [
81
83
"CUDA_HOME=${cuda_nvcc}"
82
84
"CUDA_LIB=${lib.getLib cuda_cudart}/lib"
83
85
"CUDA_INC=${lib.getDev cuda_cudart}/include"