tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{ _cuda, lib }:
2{
3 /**
4 Returns a boolean indicating whether the provided redist system is supported by any of the provided redist systems.
5
6 NOTE: No guarantees are made about this function's stability. You may use it at your own risk.
7
8 # Type
9
10 ```
11 _redistSystemIsSupported
12 :: (redistSystem :: RedistSystem)
13 -> (redistSystems :: List RedistSystem)
14 -> Bool
15 ```
16
17 # Inputs
18
19 `redistSystem`
20
21 : The redist system to check
22
23 `redistSystems`
24
25 : The list of redist systems to check against
26
27 # Examples
28
29 :::{.example}
30 ## `cudaLib._redistSystemIsSupported` usage examples
31
32 ```nix
33 _redistSystemIsSupported "linux-x86_64" [ "linux-x86_64" ]
34 => true
35 ```
36
37 ```nix
38 _redistSystemIsSupported "linux-x86_64" [ "linux-aarch64" ]
39 => false
40 ```
41
42 ```nix
43 _redistSystemIsSupported "linux-x86_64" [ "linux-aarch64" "linux-x86_64" ]
44 => true
45 ```
46
47 ```nix
48 _redistSystemIsSupported "linux-x86_64" [ "linux-aarch64" "linux-all" ]
49 => true
50 ```
51 :::
52 */
53 _redistSystemIsSupported =
54 redistSystem: redistSystems:
55 lib.findFirst (
56 redistSystem':
57 redistSystem' == redistSystem || redistSystem' == "linux-all" || redistSystem' == "source"
58 ) null redistSystems != null;
59
60 /**
61 Maps a NVIDIA redistributable system to Nix systems.
62
63 NOTE: This function returns a list of systems because the redistributable systems `"linux-all"` and `"source"` can
64 be built on multiple systems.
65
66 NOTE: This function *will* be called by unsupported systems because `cudaPackages` is evaluated on all systems. As
67 such, we need to handle unsupported systems gracefully.
68
69 # Type
70
71 ```
72 getNixSystems :: (redistSystem :: RedistSystem) -> [String]
73 ```
74
75 # Inputs
76
77 `redistSystem`
78
79 : The NVIDIA redistributable system
80
81 # Examples
82
83 :::{.example}
84 ## `cudaLib.getNixSystems` usage examples
85
86 ```nix
87 getNixSystems "linux-sbsa"
88 => [ "aarch64-linux" ]
89 ```
90
91 ```nix
92 getNixSystems "linux-aarch64"
93 => [ "aarch64-linux" ]
94 ```
95 :::
96 */
97 getNixSystems =
98 redistSystem:
99 if redistSystem == "linux-x86_64" then
100 [ "x86_64-linux" ]
101 else if redistSystem == "linux-sbsa" || redistSystem == "linux-aarch64" then
102 [ "aarch64-linux" ]
103 else if redistSystem == "linux-all" || redistSystem == "source" then
104 [
105 "aarch64-linux"
106 "x86_64-linux"
107 ]
108 else
109 [ ];
110
111 /**
112 Maps a Nix system to a NVIDIA redistributable system.
113
114 NOTE: Certain Nix systems can map to multiple NVIDIA redistributable systems. In particular, ARM systems can map to
115 either `linux-sbsa` (for server-grade ARM chips) or `linux-aarch64` (for Jetson devices). Complicating matters
116 further, as of CUDA 13.0, Jetson Thor devices use `linux-sbsa` instead of `linux-aarch64`. (It is unknown whether
117 NVIDIA plans to make the Orin series use `linux-sbsa` as well for the CUDA 13.0 release.)
118
119 NOTE: This function *will* be called by unsupported systems because `cudaPackages` is evaluated on all systems. As
120 such, we need to handle unsupported systems gracefully.
121
122 NOTE: This function does not check whether the provided CUDA capabilities are valid for the given CUDA version.
123 The heavy validation work to ensure consistency of CUDA capabilities is performed by backendStdenv.
124
125 # Type
126
127 ```
128 getRedistSystem ::
129 { cudaCapabilities :: List String
130 , cudaMajorMinorVersion :: String
131 , system :: String
132 }
133 -> String
134 ```
135
136 # Inputs
137
138 `cudaCapabilities`
139
140 : The list of CUDA capabilities to build GPU code for
141
142 `cudaMajorMinorVersion`
143
144 : The major and minor version of CUDA (e.g. "12.6")
145
146 `system`
147
148 : The Nix system
149
150 # Examples
151
152 :::{.example}
153 ## `cudaLib.getRedistSystem` usage examples
154
155 ```nix
156 getRedistSystem {
157 cudaCapabilities = [ "8.7" ];
158 cudaMajorMinorVersion = "12.6";
159 system = "aarch64-linux";
160 }
161 => "linux-aarch64"
162 ```
163
164 ```nix
165 getRedistSystem {
166 cudaCapabilities = [ "11.0" ];
167 cudaMajorMinorVersion = "13.0";
168 system = "aarch64-linux";
169 }
170 => "linux-sbsa"
171 ```
172
173 ```nix
174 getRedistSystem {
175 cudaCapabilities = [ "8.0" "8.9" ];
176 cudaMajorMinorVersion = "12.6";
177 system = "aarch64-linux";
178 }
179 => "linux-sbsa"
180 ```
181 :::
182 */
183 getRedistSystem =
184 {
185 cudaCapabilities,
186 cudaMajorMinorVersion,
187 system,
188 }:
189 if system == "x86_64-linux" then
190 "linux-x86_64"
191 else if system == "aarch64-linux" then
192 # If all the Jetson devices are at least 10.1 (Thor, CUDA 12.9; CUDA 13.0 and later use 11.0 for Thor), then
193 # we've got SBSA.
194 if
195 lib.all (
196 cap: _cuda.db.cudaCapabilityToInfo.${cap}.isJetson -> lib.versionAtLeast cap "10.1"
197 ) cudaCapabilities
198 then
199 "linux-sbsa"
200 # Otherwise we've got some Jetson devices older than Thor and need to use linux-aarch64.
201 else
202 "linux-aarch64"
203 else
204 "unsupported";
205
206 /**
207 Function to generate a URL for something in the redistributable tree.
208
209 # Type
210
211 ```
212 mkRedistUrl :: (redistName :: RedistName) -> (relativePath :: NonEmptyStr) -> RedistUrl
213 ```
214
215 # Inputs
216
217 `redistName`
218
219 : The name of the redistributable
220
221 `relativePath`
222
223 : The relative path to a file in the redistributable tree
224 */
225 mkRedistUrl =
226 redistName: relativePath:
227 lib.concatStringsSep "/" (
228 [ _cuda.db.redistUrlPrefix ]
229 ++ (
230 if redistName != "tensorrt" then
231 [
232 redistName
233 "redist"
234 ]
235 else
236 [ "machine-learning" ]
237 )
238 ++ [ relativePath ]
239 );
240
241 /**
242 Function which accepts an attribute set mapping redistributable name to version and retrieves the corresponding
243 collection of manifests from `_cuda.manifests`. Additionally, the version provided is used to populate the
244 `release_label` field in the corresponding manifest if it is missing.
245
246 It is an error to provide a redistributable name and version for which there is no corresponding manifest.
247
248 # Type
249
250 ```
251 selectManifests :: (versions :: AttrSet RedistName Version) -> AttrSet RedistName Manifest
252 ```
253
254 # Inputs
255
256 `versions`
257
258 : An attribute set mapping redistributable name to manifest version
259 */
260 selectManifests = lib.mapAttrs (
261 name: version:
262 let
263 manifest = _cuda.manifests.${name}.${version};
264 in
265 manifest
266 // {
267 release_label = manifest.release_label or version;
268 }
269 );
270}