tjh.dev
1{ cudaMajorMinorVersion, lib }:
2let
3 inherit (lib) attrsets modules trivial;
4 redistName = "cuda";
5
6 # Manifest files for CUDA redistributables (aka redist). These can be found at
7 # https://developer.download.nvidia.com/compute/cuda/redist/
8 # Maps a cuda version to the specific version of the manifest.
9 cudaVersionMap = {
10 "11.4" = "11.4.4";
11 "11.5" = "11.5.2";
12 "11.6" = "11.6.2";
13 "11.7" = "11.7.1";
14 "11.8" = "11.8.0";
15 "12.0" = "12.0.1";
16 "12.1" = "12.1.1";
17 "12.2" = "12.2.2";
18 "12.3" = "12.3.2";
19 "12.4" = "12.4.1";
20 "12.5" = "12.5.1";
21 "12.6" = "12.6.3";
22 "12.8" = "12.8.1";
23 "12.9" = "12.9.1";
24 };
25
26 # Check if the current CUDA version is supported.
27 cudaVersionMappingExists = builtins.hasAttr cudaMajorMinorVersion cudaVersionMap;
28
29 # fullCudaVersion : String
30 fullCudaVersion = cudaVersionMap.${cudaMajorMinorVersion};
31
32 evaluatedModules = modules.evalModules {
33 modules = [
34 ../modules
35 # We need to nest the manifests in a config.cuda.manifests attribute so the
36 # module system can evaluate them.
37 {
38 cuda.manifests = {
39 redistrib = trivial.importJSON (./manifests + "/redistrib_${fullCudaVersion}.json");
40 feature = trivial.importJSON (./manifests + "/feature_${fullCudaVersion}.json");
41 };
42 }
43 ];
44 };
45
46 # Generally we prefer to do things involving getting attribute names with feature_manifest instead
47 # of redistrib_manifest because the feature manifest will have *only* the redist system
48 # names as the keys, whereas the redistrib manifest will also have things like version, name, license,
49 # and license_path.
50 featureManifest = evaluatedModules.config.cuda.manifests.feature;
51 redistribManifest = evaluatedModules.config.cuda.manifests.redistrib;
52
53 # Builder function which builds a single redist package for a given platform.
54 # buildRedistPackage : callPackage -> PackageName -> Derivation
55 buildRedistPackage =
56 callPackage: pname:
57 callPackage ../generic-builders/manifest.nix {
58 inherit pname redistName;
59 # We pass the whole release to the builder because it has logic to handle
60 # the case we're trying to build on an unsupported platform.
61 redistribRelease = redistribManifest.${pname};
62 featureRelease = featureManifest.${pname};
63 };
64
65 # Build all the redist packages given final and prev.
66 redistPackages =
67 final: _prev:
68 # Wrap the whole thing in an optionalAttrs so we can return an empty set if the CUDA version
69 # is not supported.
70 # NOTE: We cannot include the call to optionalAttrs *in* the pipe as we would strictly evaluate the
71 # attrNames before we check if the CUDA version is supported.
72 attrsets.optionalAttrs cudaVersionMappingExists (
73 trivial.pipe featureManifest [
74 # Get all the package names
75 builtins.attrNames
76 # Build the redist packages
77 (trivial.flip attrsets.genAttrs (buildRedistPackage final.callPackage))
78 ]
79 );
80in
81redistPackages