tjh.dev
1{ _cuda, lib }:
2let
3 cudaLib = _cuda.lib;
4in
5{
6 /**
7 Replaces dots in a string with underscores.
8
9 # Type
10
11 ```
12 dotsToUnderscores :: (str :: String) -> String
13 ```
14
15 # Inputs
16
17 `str`
18
19 : The string for which dots shall be replaced by underscores
20
21 # Examples
22
23 :::{.example}
24 ## `cudaLib.dotsToUnderscores` usage examples
25
26 ```nix
27 dotsToUnderscores "1.2.3"
28 => "1_2_3"
29 ```
30 :::
31 */
32 dotsToUnderscores = lib.replaceStrings [ "." ] [ "_" ];
33
34 /**
35 Removes the dots from a string.
36
37 # Type
38
39 ```
40 dropDots :: (str :: String) -> String
41 ```
42
43 # Inputs
44
45 `str`
46
47 : The string to remove dots from
48
49 # Examples
50
51 :::{.example}
52 ## `cudaLib.dropDots` usage examples
53
54 ```nix
55 dropDots "1.2.3"
56 => "123"
57 ```
58 :::
59 */
60 dropDots = lib.replaceStrings [ "." ] [ "" ];
61
62 /**
63 Produces an attribute set of useful data and functionality for packaging CUDA software within Nixpkgs.
64
65 # Type
66
67 ```
68 formatCapabilities
69 :: { cudaCapabilityToInfo :: AttrSet CudaCapability CudaCapabilityInfo
70 , cudaCapabilities :: List CudaCapability
71 , cudaForwardCompat :: Bool
72 }
73 -> { cudaCapabilities :: List CudaCapability
74 , cudaForwardCompat :: Bool
75 , gencode :: List String
76 , realArches :: List String
77 , virtualArches :: List String
78 , archNames :: List String
79 , arches :: List String
80 , gencodeString :: String
81 , cmakeCudaArchitecturesString :: String
82 }
83 ```
84
85 # Inputs
86
87 `cudaCapabilityToInfo`
88
89 : A mapping of CUDA capabilities to their information
90
91 `cudaCapabilities`
92
93 : A list of CUDA capabilities to use
94
95 `cudaForwardCompat`
96
97 : A boolean indicating whether to include the forward compatibility gencode (+PTX) to support future GPU
98 generations
99 */
100 formatCapabilities =
101 {
102 cudaCapabilityToInfo,
103 cudaCapabilities,
104 cudaForwardCompat,
105 }:
106 let
107 /**
108 The real architectures for the given CUDA capabilities.
109
110 # Type
111
112 ```
113 realArches :: List String
114 ```
115 */
116 realArches = lib.map cudaLib.mkRealArchitecture cudaCapabilities;
117
118 /**
119 The virtual architectures for the given CUDA capabilities.
120
121 These are typically used for forward compatibility, when trying to support an architecture newer than the CUDA
122 version allows.
123
124 # Type
125
126 ```
127 virtualArches :: List String
128 ```
129 */
130 virtualArches = lib.map cudaLib.mkVirtualArchitecture cudaCapabilities;
131
132 /**
133 The gencode flags for the given CUDA capabilities.
134
135 # Type
136
137 ```
138 gencode :: List String
139 ```
140 */
141 gencode =
142 let
143 base = lib.map (cudaLib.mkGencodeFlag "sm") cudaCapabilities;
144 forward = cudaLib.mkGencodeFlag "compute" (lib.last cudaCapabilities);
145 in
146 base ++ lib.optionals cudaForwardCompat [ forward ];
147 in
148 {
149 inherit
150 cudaCapabilities
151 cudaForwardCompat
152 gencode
153 realArches
154 virtualArches
155 ;
156
157 /**
158 The architecture names for the given CUDA capabilities.
159
160 # Type
161
162 ```
163 archNames :: List String
164 ```
165 */
166 # E.g. [ "Ampere" "Turing" ]
167 archNames = lib.pipe cudaCapabilities [
168 (lib.map (cudaCapability: cudaCapabilityToInfo.${cudaCapability}.archName))
169 lib.unique
170 lib.naturalSort
171 ];
172
173 /**
174 The architectures for the given CUDA capabilities, including both real and virtual architectures.
175
176 When `cudaForwardCompat` is enabled, the last architecture in the list is used as the forward compatibility architecture.
177
178 # Type
179
180 ```
181 arches :: List String
182 ```
183 */
184 # E.g. [ "sm_75" "sm_86" "compute_86" ]
185 arches = realArches ++ lib.optionals cudaForwardCompat [ (lib.last virtualArches) ];
186
187 /**
188 The CMake-compatible CUDA architectures string for the given CUDA capabilities.
189
190 # Type
191
192 ```
193 cmakeCudaArchitecturesString :: String
194 ```
195 */
196 cmakeCudaArchitecturesString = cudaLib.mkCmakeCudaArchitecturesString cudaCapabilities;
197
198 /**
199 The gencode string for the given CUDA capabilities.
200
201 # Type
202
203 ```
204 gencodeString :: String
205 ```
206 */
207 gencodeString = lib.concatStringsSep " " gencode;
208 };
209
210 /**
211 Produces a CMake-compatible CUDA architecture string from a list of CUDA capabilities.
212
213 # Type
214
215 ```
216 mkCmakeCudaArchitecturesString :: (cudaCapabilities :: List String) -> String
217 ```
218
219 # Inputs
220
221 `cudaCapabilities`
222
223 : The CUDA capabilities to convert
224
225 # Examples
226
227 :::{.example}
228 ## `cudaLib.mkCmakeCudaArchitecturesString` usage examples
229
230 ```nix
231 mkCmakeCudaArchitecturesString [ "8.9" "10.0a" ]
232 => "89;100a"
233 ```
234 :::
235 */
236 mkCmakeCudaArchitecturesString = lib.concatMapStringsSep ";" cudaLib.dropDots;
237
238 /**
239 Produces a gencode flag from a CUDA capability.
240
241 # Type
242
243 ```
244 mkGencodeFlag :: (archPrefix :: String) -> (cudaCapability :: String) -> String
245 ```
246
247 # Inputs
248
249 `archPrefix`
250
251 : The architecture prefix to use for the `code` field
252
253 `cudaCapability`
254
255 : The CUDA capability to convert
256
257 # Examples
258
259 :::{.example}
260 ## `cudaLib.mkGencodeFlag` usage examples
261
262 ```nix
263 mkGencodeFlag "sm" "8.9"
264 => "-gencode=arch=compute_89,code=sm_89"
265 ```
266
267 ```nix
268 mkGencodeFlag "compute" "10.0a"
269 => "-gencode=arch=compute_100a,code=compute_100a"
270 ```
271 :::
272 */
273 mkGencodeFlag =
274 archPrefix: cudaCapability:
275 let
276 cap = cudaLib.dropDots cudaCapability;
277 in
278 "-gencode=arch=compute_${cap},code=${archPrefix}_${cap}";
279
280 /**
281 Produces a real architecture string from a CUDA capability.
282
283 # Type
284
285 ```
286 mkRealArchitecture :: (cudaCapability :: String) -> String
287 ```
288
289 # Inputs
290
291 `cudaCapability`
292
293 : The CUDA capability to convert
294
295 # Examples
296
297 :::{.example}
298 ## `cudaLib.mkRealArchitecture` usage examples
299
300 ```nix
301 mkRealArchitecture "8.9"
302 => "sm_89"
303 ```
304
305 ```nix
306 mkRealArchitecture "10.0a"
307 => "sm_100a"
308 ```
309 :::
310 */
311 mkRealArchitecture = cudaCapability: "sm_" + cudaLib.dropDots cudaCapability;
312
313 /**
314 Create a versioned attribute name from a version by replacing dots with underscores.
315
316 # Type
317
318 ```
319 mkVersionedName :: (name :: String) -> (version :: Version) -> String
320 ```
321
322 # Inputs
323
324 `name`
325
326 : The name to use
327
328 `version`
329
330 : The version to use
331
332 # Examples
333
334 :::{.example}
335 ## `cudaLib.mkVersionedName` usage examples
336
337 ```nix
338 mkVersionedName "hello" "1.2.3"
339 => "hello_1_2_3"
340 ```
341
342 ```nix
343 mkVersionedName "cudaPackages" "12.8"
344 => "cudaPackages_12_8"
345 ```
346 :::
347 */
348 mkVersionedName = name: version: "${name}_${cudaLib.dotsToUnderscores version}";
349
350 /**
351 Produces a virtual architecture string from a CUDA capability.
352
353 # Type
354
355 ```
356 mkVirtualArchitecture :: (cudaCapability :: String) -> String
357 ```
358
359 # Inputs
360
361 `cudaCapability`
362
363 : The CUDA capability to convert
364
365 # Examples
366
367 :::{.example}
368 ## `cudaLib.mkVirtualArchitecture` usage examples
369
370 ```nix
371 mkVirtualArchitecture "8.9"
372 => "compute_89"
373 ```
374
375 ```nix
376 mkVirtualArchitecture "10.0a"
377 => "compute_100a"
378 ```
379 :::
380 */
381 mkVirtualArchitecture = cudaCapability: "compute_" + cudaLib.dropDots cudaCapability;
382}