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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
|
{
_cuda,
addDriverRunpath,
backendStdenv,
cmake,
cuda_cudart,
cuda_nvcc,
cuda_nvrtc,
cudaNamePrefix,
cudnn,
fetchFromGitHub,
flags,
gtest,
lib,
libcublas,
libcurand,
ninja,
python3Packages,
# Options
pythonSupport ? true,
enableF16C ? false,
enableTools ? false,
# passthru.updateScript
gitUpdater,
}:
let
inherit (_cuda.lib) _mkMetaBadPlatforms;
inherit (lib) licenses maintainers teams;
inherit (lib.asserts) assertMsg;
inherit (lib.attrsets) getBin;
inherit (lib.lists) all optionals;
inherit (lib.strings)
cmakeBool
cmakeFeature
optionalString
versionAtLeast
;
inherit (lib.trivial) flip;
in
# TODO: Tests.
assert assertMsg (!enableTools) "enableTools is not yet implemented";
backendStdenv.mkDerivation (finalAttrs: {
__structuredAttrs = true;
strictDeps = true;
# NOTE: Depends on the CUDA package set, so use cudaNamePrefix.
name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
pname = "cutlass";
version = "4.7.1";
src = fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
tag = "v${finalAttrs.version}";
hash = "sha256-z/jASQGRj9fYwgyJKvzWgowG7j2zNKb3xeFUzYwwLQA=";
};
# TODO: As a header-only library, we should make sure we have an `include` directory or similar which is not a
# superset of the `out` (`bin`) or `dev` outputs (whih is what the multiple-outputs setup hook does by default).
outputs = [ "out" ] ++ optionals pythonSupport [ "dist" ];
nativeBuildInputs = [
cuda_nvcc
cmake
ninja
python3Packages.python # Python is always required
]
++ optionals pythonSupport (
with python3Packages;
[
build
pythonOutputDistHook
setuptools
]
);
postPatch =
# Prepend some commands to the CUDA.cmake file so it can find the CUDA libraries using CMake's FindCUDAToolkit
# module. These target names are used throughout the project; I (@connorbaker) did not choose them.
''
nixLog "patching CUDA.cmake to use FindCUDAToolkit"
mv ./CUDA.cmake ./_CUDA_Append.cmake
cat > ./_CUDA_Prepend.cmake <<'EOF'
find_package(CUDAToolkit REQUIRED)
foreach(_target cudart cuda_driver nvrtc)
if (NOT TARGET CUDA::''${_target})
message(FATAL_ERROR "''${_target} Not Found")
endif()
message(STATUS "''${_target} library: ''${CUDA_''${_target}_LIBRARY}")
add_library(''${_target} ALIAS CUDA::''${_target})
endforeach()
EOF
cat ./_CUDA_Prepend.cmake ./_CUDA_Append.cmake > ./CUDA.cmake
''
# Patch cutlass to use the provided NVCC.
# '_CUDA_INSTALL_PATH = os.getenv("CUDA_INSTALL_PATH", _cuda_install_path_from_nvcc())' \
# '_CUDA_INSTALL_PATH = "${getBin cuda_nvcc}"'
+ ''
nixLog "patching python bindings to make cuda_install_path fail"
substituteInPlace ./python/cutlass_cppgen/__init__.py \
--replace-fail \
'def cuda_install_path():' \
'
def cuda_install_path():
raise RuntimeException("not supported with Nixpkgs CUDA packaging")
'
''
# Patch the python bindings to use environment variables set by Nixpkgs.
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L83
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L84
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L320
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L322
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L347
# https://github.com/NVIDIA/cutlass/blob/2e602843e75100d0e03934efb386b3e1e35d7907/python/cutlass_cppgen/backend/compiler.py#L363
+ ''
nixLog "patching python bindings to use environment variables"
substituteInPlace ./python/cutlass_cppgen/backend/compiler.py \
--replace-fail \
'self.include_paths = include_paths' \
'self.include_paths = include_paths + [root + "/include" for root in os.getenv("CUDAToolkit_ROOT").split(";")]' \
--replace-fail \
'self.flags = flags' \
'self.flags = flags + ["-L" + root + "/lib" for root in os.getenv("CUDAToolkit_ROOT").split(";")]' \
--replace-fail \
"\''${cuda_install_path}/bin/nvcc" \
'${getBin cuda_nvcc}/bin/nvcc' \
--replace-fail \
'"cuda_install_path": cuda_install_path(),' \
"" \
--replace-fail \
'f"{cuda_install_path()}/bin/nvcc"' \
'"${getBin cuda_nvcc}/bin/nvcc"' \
--replace-fail \
'cuda_install_path() + "/include",' \
""
'';
enableParallelBuilding = true;
buildInputs = [
cuda_cudart
cuda_nvrtc
libcurand
]
++ optionals enableTools [
cudnn
libcublas
];
cmakeFlags = [
(cmakeFeature "CUTLASS_NVCC_ARCHS" flags.cmakeCudaArchitecturesString)
(cmakeBool "CUTLASS_ENABLE_EXAMPLES" false)
# Tests.
(cmakeBool "CUTLASS_ENABLE_TESTS" finalAttrs.finalPackage.doCheck)
(cmakeBool "CUTLASS_ENABLE_GTEST_UNIT_TESTS" finalAttrs.finalPackage.doCheck)
(cmakeBool "CUTLASS_USE_SYSTEM_GOOGLETEST" true)
# NOTE: Both CUDNN and CUBLAS can be used by the examples and the profiler. Since they are large dependencies, they
# are disabled by default.
(cmakeBool "CUTLASS_ENABLE_TOOLS" enableTools)
(cmakeBool "CUTLASS_ENABLE_CUBLAS" enableTools)
(cmakeBool "CUTLASS_ENABLE_CUDNN" enableTools)
# NOTE: Requires x86_64 and hardware support.
(cmakeBool "CUTLASS_ENABLE_F16C" enableF16C)
# TODO: Unity builds are supposed to reduce build time, but this seems to just reduce the number of tasks
# generated?
# NOTE: Good explanation of unity builds:
# https://www.methodpark.de/blog/how-to-speed-up-clang-tidy-with-unity-builds.
(cmakeBool "CUTLASS_UNITY_BUILD_ENABLED" false)
];
postBuild = lib.optionalString pythonSupport ''
pushd "$NIX_BUILD_TOP/$sourceRoot"
nixLog "building Python wheel"
pyproject-build \
--no-isolation \
--outdir "$NIX_BUILD_TOP/$sourceRoot/''${cmakeBuildDir:?}/dist/" \
--wheel
popd >/dev/null
'';
doCheck = false;
checkInputs = [ gtest ];
# NOTE: Because the test cases immediately create and try to run the binaries, we don't have an opportunity
# to patch them with autoAddDriverRunpath. To get around this, we add the driver runpath to the environment.
# TODO: This would break Jetson when using cuda_compat, as it must come first.
preCheck = optionalString finalAttrs.finalPackage.doCheck ''
export LD_LIBRARY_PATH="$(readlink -mnv "${addDriverRunpath.driverLink}/lib")"
'';
# This is *not* a derivation you want to build on a small machine.
requiredSystemFeatures = optionals finalAttrs.finalPackage.doCheck [
"big-parallel"
"cuda"
];
passthru = {
updateScript = gitUpdater {
inherit (finalAttrs) pname version;
rev-prefix = "v";
};
# TODO:
# tests.test = cutlass.overrideAttrs { doCheck = true; };
# Include required architectures in compatibility check.
# https://github.com/NVIDIA/cutlass/tree/main?tab=readme-ov-file#compatibility
platformAssertions = [
{
message = "all capabilities are >= 7.0 (${builtins.toJSON flags.cudaCapabilities})";
assertion = all (flip versionAtLeast "7.0") flags.cudaCapabilities;
}
];
};
meta = {
description = "CUDA Templates for Linear Algebra Subroutines";
homepage = "https://github.com/NVIDIA/cutlass";
license = licenses.bsd3;
platforms = [
"aarch64-linux"
"x86_64-linux"
];
badPlatforms = _mkMetaBadPlatforms finalAttrs;
maintainers = [ maintainers.connorbaker ];
teams = [ teams.cuda ];
};
})
|
