tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1# This is what nvcc uses as a backend,
2# and it has to be an officially supported one (e.g. gcc14 for cuda12).
3#
4# It, however, propagates current stdenv's libstdc++ to avoid "GLIBCXX_* not found errors"
5# when linked with other C++ libraries.
6# E.g. for cudaPackages_12_9 we use gcc14 with gcc's libstdc++
7# Cf. https://github.com/NixOS/nixpkgs/pull/218265 for context
8{
9 _cuda,
10 config,
11 cudaMajorMinorVersion,
12 lib,
13 pkgs,
14 stdenv,
15 stdenvAdapters,
16}:
17let
18 inherit (builtins)
19 throw
20 toJSON
21 toString
22 ;
23 inherit (_cuda.db) allSortedCudaCapabilities cudaCapabilityToInfo nvccCompatibilities;
24 inherit (_cuda.lib)
25 _cudaCapabilityIsDefault
26 _cudaCapabilityIsSupported
27 _mkFailedAssertionsString
28 getRedistSystem
29 mkVersionedName
30 ;
31 inherit (lib)
32 assertMsg
33 filter
34 findFirst
35 flip
36 intersectLists
37 pipe
38 range
39 reverseList
40 subtractLists
41 toIntBase10
42 versionAtLeast
43 versionOlder
44 ;
45 inherit (lib.versions) major;
46
47 # NOTE: By virtue of processing a sorted list (allSortedCudaCapabilities), our groups will be sorted.
48
49 architectureSpecificCudaCapabilities = filter (
50 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isArchitectureSpecific
51 ) allSortedCudaCapabilities;
52
53 familySpecificCudaCapabilities = filter (
54 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isFamilySpecific
55 ) allSortedCudaCapabilities;
56
57 jetsonCudaCapabilities = filter (
58 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isJetson
59 ) allSortedCudaCapabilities;
60
61 passthruExtra = {
62 nvccHostCCMatchesStdenvCC = backendStdenv.cc == stdenv.cc;
63
64 # TODO(@connorbaker): Does it make sense to expose the `stdenv` we were called with and the `stdenv` selected
65 # prior to using `stdenvAdapters.useLibsFrom`?
66
67 # The Nix system of the host platform.
68 hostNixSystem = stdenv.hostPlatform.system;
69
70 # The Nix system of the host platform for the CUDA redistributable.
71 hostRedistSystem = getRedistSystem {
72 inherit (passthruExtra) cudaCapabilities;
73 inherit cudaMajorMinorVersion;
74 inherit (stdenv.hostPlatform) system;
75 };
76
77 # Sets whether packages should be built with forward compatibility.
78 # TODO(@connorbaker): If the requested CUDA capabilities are not supported by the current CUDA version,
79 # should we throw an evaluation warning and build with forward compatibility?
80 cudaForwardCompat = config.cudaForwardCompat or true;
81
82 # CUDA capabilities which are supported by the current CUDA version.
83 supportedCudaCapabilities = filter (
84 cudaCapability:
85 _cudaCapabilityIsSupported cudaMajorMinorVersion cudaCapabilityToInfo.${cudaCapability}
86 ) allSortedCudaCapabilities;
87
88 # Find the default set of capabilities for this CUDA version using the list of supported capabilities.
89 # Includes only baseline capabilities.
90 defaultCudaCapabilities = filter (
91 cudaCapability:
92 _cudaCapabilityIsDefault cudaMajorMinorVersion cudaCapabilityToInfo.${cudaCapability}
93 ) passthruExtra.supportedCudaCapabilities;
94
95 # The resolved requested or default CUDA capabilities.
96 cudaCapabilities =
97 if config.cudaCapabilities or [ ] != [ ] then
98 config.cudaCapabilities
99 else
100 passthruExtra.defaultCudaCapabilities;
101
102 # Requested architecture-specific CUDA capabilities.
103 requestedArchitectureSpecificCudaCapabilities = intersectLists architectureSpecificCudaCapabilities passthruExtra.cudaCapabilities;
104
105 # Whether the requested CUDA capabilities include architecture-specific CUDA capabilities.
106 hasArchitectureSpecificCudaCapability =
107 passthruExtra.requestedArchitectureSpecificCudaCapabilities != [ ];
108
109 # Requested family-specific CUDA capabilities.
110 requestedFamilySpecificCudaCapabilities = intersectLists familySpecificCudaCapabilities passthruExtra.cudaCapabilities;
111
112 # Whether the requested CUDA capabilities include family-specific CUDA capabilities.
113 hasFamilySpecificCudaCapability = passthruExtra.requestedFamilySpecificCudaCapabilities != [ ];
114
115 # Requested Jetson CUDA capabilities.
116 requestedJetsonCudaCapabilities = intersectLists jetsonCudaCapabilities passthruExtra.cudaCapabilities;
117
118 # Whether the requested CUDA capabilities include Jetson CUDA capabilities.
119 hasJetsonCudaCapability = passthruExtra.requestedJetsonCudaCapabilities != [ ];
120 };
121
122 assertions =
123 let
124 # Jetson devices (pre-Thor) cannot be targeted by the same binaries which target non-Jetson devices. While
125 # NVIDIA provides both `linux-aarch64` and `linux-sbsa` packages, which both target `aarch64`,
126 # they are built with different settings and cannot be mixed.
127 preThorJetsonCudaCapabilities = filter (flip versionOlder "10.1") passthruExtra.requestedJetsonCudaCapabilities;
128 postThorJetsonCudaCapabilities = filter (flip versionAtLeast "10.1") passthruExtra.requestedJetsonCudaCapabilities;
129
130 # Remove all known capabilities from the user's list to find unrecognized capabilities.
131 unrecognizedCudaCapabilities = subtractLists allSortedCudaCapabilities passthruExtra.cudaCapabilities;
132
133 # Capabilities which are too old for this CUDA version.
134 tooOldCudaCapabilities = filter (
135 cap:
136 let
137 # This can be null!
138 maybeMax = cudaCapabilityToInfo.${cap}.maxCudaMajorMinorVersion;
139 in
140 maybeMax != null && lib.versionOlder maybeMax cudaMajorMinorVersion
141 ) passthruExtra.cudaCapabilities;
142
143 # Capabilities which are too new for this CUDA version.
144 tooNewCudaCapabilities = filter (
145 cap: lib.versionOlder cudaMajorMinorVersion cudaCapabilityToInfo.${cap}.minCudaMajorMinorVersion
146 ) passthruExtra.cudaCapabilities;
147 in
148 [
149 {
150 message = "Requested unrecognized CUDA capabilities: ${toJSON unrecognizedCudaCapabilities}";
151 assertion = unrecognizedCudaCapabilities == [ ];
152 }
153 {
154 message = "Requested CUDA capabilities which are too old for CUDA ${cudaMajorMinorVersion}: ${toJSON tooOldCudaCapabilities}";
155 assertion = tooOldCudaCapabilities == [ ];
156 }
157 {
158 message = "Requested CUDA capabilities which are too new for CUDA ${cudaMajorMinorVersion}: ${toJSON tooNewCudaCapabilities}";
159 assertion = tooNewCudaCapabilities == [ ];
160 }
161 {
162 message =
163 "Requested Jetson CUDA capabilities (${toJSON passthruExtra.requestedJetsonCudaCapabilities}) require "
164 + "hostPlatform (${passthruExtra.hostNixSystem}) to be aarch64-linux";
165 assertion = passthruExtra.hasJetsonCudaCapability -> passthruExtra.hostNixSystem == "aarch64-linux";
166 }
167 {
168 message =
169 "Requested pre-Thor (10.1) Jetson CUDA capabilities (${toJSON preThorJetsonCudaCapabilities}) cannot be "
170 + "specified with other capabilities (${toJSON (subtractLists preThorJetsonCudaCapabilities passthruExtra.cudaCapabilities)})";
171 assertion =
172 # If there are preThorJetsonCudaCapabilities, they must be the only requested capabilities.
173 preThorJetsonCudaCapabilities != [ ]
174 -> preThorJetsonCudaCapabilities == passthruExtra.cudaCapabilities;
175 }
176 {
177 message =
178 "Requested pre-Thor (10.1) Jetson CUDA capabilities (${toJSON preThorJetsonCudaCapabilities}) require "
179 + "computed NVIDIA hostRedistSystem (${passthruExtra.hostRedistSystem}) to be linux-aarch64";
180 assertion =
181 preThorJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-aarch64";
182 }
183 {
184 message =
185 "Requested post-Thor (10.1) Jetson CUDA capabilities (${toJSON postThorJetsonCudaCapabilities}) require "
186 + "computed NVIDIA hostRedistSystem (${passthruExtra.hostRedistSystem}) to be linux-sbsa";
187 assertion = postThorJetsonCudaCapabilities != [ ] -> passthruExtra.hostRedistSystem == "linux-sbsa";
188 }
189 ];
190
191 failedAssertionsString = _mkFailedAssertionsString assertions;
192
193 # TODO(@connorbaker): Seems like `stdenvAdapters.useLibsFrom` breaks clangStdenv's ability to find header files.
194 # To reproduce: use `nix shell .#cudaPackages_12_6.backendClangStdenv.cc` since CUDA 12.6 supports at most Clang
195 # 18, but the current stdenv uses Clang 19, requiring this code path.
196 # With:
197 #
198 # ```cpp
199 # #include <cmath>
200 #
201 # int main() {
202 # double value = 0.5;
203 # double result = std::sin(value);
204 # return 0;
205 # }
206 # ```
207 #
208 # we get:
209 #
210 # ```console
211 # $ clang++ ./main.cpp
212 # ./main.cpp:1:10: fatal error: 'cmath' file not found
213 # 1 | #include <cmath>
214 # | ^~~~~~~
215 # 1 error generated.
216 # ```
217 # TODO(@connorbaker): Seems like even using unmodified `clangStdenv` causes issues -- saxpy fails to build CMake
218 # errors during CUDA compiler identification about invalid redefinitions of things like `realpath`.
219 backendStdenv =
220 let
221 hostCCName =
222 if stdenv.cc.isGNU then
223 "gcc"
224 else if stdenv.cc.isClang then
225 "clang"
226 else
227 throw "cudaPackages.backendStdenv: unsupported host compiler: ${stdenv.cc.name}";
228
229 versions = nvccCompatibilities.${cudaMajorMinorVersion}.${hostCCName};
230
231 stdenvIsSupportedVersion =
232 versionAtLeast (major stdenv.cc.version) versions.minMajorVersion
233 && versionAtLeast versions.maxMajorVersion (major stdenv.cc.version);
234
235 maybeGetVersionedCC =
236 if hostCCName == "gcc" then
237 version: pkgs."gcc${version}Stdenv" or null
238 else
239 version: pkgs."llvmPackages_${version}".stdenv or null;
240
241 maybeHostStdenv =
242 pipe (range (toIntBase10 versions.minMajorVersion) (toIntBase10 versions.maxMajorVersion))
243 [
244 # Convert integers to strings.
245 (map toString)
246 # Prefer the highest available version.
247 reverseList
248 # Map to the actual stdenvs or null if unavailable.
249 (map maybeGetVersionedCC)
250 # Get the first available version.
251 (findFirst (x: x != null) null)
252 ];
253 in
254 # If the current stdenv's compiler version is compatible, or we're on an unsupported host system, use stdenv
255 # directly.
256 # If we're on an unsupported host system (like darwin), there's not much else we can do, but we should not break
257 # evaluation on unsupported systems.
258 if stdenvIsSupportedVersion || passthruExtra.hostRedistSystem == "unsupported" then
259 stdenv
260 # Otherwise, try to find a compatible stdenv.
261 else
262 assert assertMsg (maybeHostStdenv != null)
263 "backendStdenv: no supported host compiler found (tried ${hostCCName} ${versions.minMajorVersion} to ${versions.maxMajorVersion})";
264 stdenvAdapters.useLibsFrom stdenv maybeHostStdenv;
265in
266# TODO: Consider testing whether we in fact use the newer libstdc++
267assert assertMsg (failedAssertionsString == "")
268 "${mkVersionedName "cudaPackages" cudaMajorMinorVersion}.backendStdenv has failed assertions:${failedAssertionsString}";
269backendStdenv.override (prevArgs: {
270 extraAttrs = prevArgs.extraAttrs or { } // passthruExtra;
271})