blob: 5b87a3df0959a38fbe0441c71289f02e9aeb82d6 (
plain) (
blame)
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
|
{ cudaVersion, lib }:
let
inherit (lib) attrsets modules trivial;
redistName = "cuda";
# Manifest files for CUDA redistributables (aka redist). These can be found at
# https://developer.download.nvidia.com/compute/cuda/redist/
# Maps a cuda version to the specific version of the manifest.
cudaVersionMap = {
"11.4" = "11.4.4";
"11.5" = "11.5.2";
"11.6" = "11.6.2";
"11.7" = "11.7.1";
"11.8" = "11.8.0";
"12.0" = "12.0.1";
"12.1" = "12.1.1";
"12.2" = "12.2.2";
"12.3" = "12.3.0";
};
# Check if the current CUDA version is supported.
cudaVersionMappingExists = builtins.hasAttr cudaVersion cudaVersionMap;
# fullCudaVersion : String
fullCudaVersion = cudaVersionMap.${cudaVersion};
evaluatedModules = modules.evalModules {
modules = [
../modules
# We need to nest the manifests in a config.cuda.manifests attribute so the
# module system can evaluate them.
{
cuda.manifests = {
redistrib = trivial.importJSON (./manifests + "/redistrib_${fullCudaVersion}.json");
feature = trivial.importJSON (./manifests + "/feature_${fullCudaVersion}.json");
};
}
];
};
# Generally we prefer to do things involving getting attribute names with feature_manifest instead
# of redistrib_manifest because the feature manifest will have *only* the redist architecture
# names as the keys, whereas the redistrib manifest will also have things like version, name, license,
# and license_path.
featureManifest = evaluatedModules.config.cuda.manifests.feature;
redistribManifest = evaluatedModules.config.cuda.manifests.redistrib;
# Builder function which builds a single redist package for a given platform.
# buildRedistPackage : callPackage -> PackageName -> Derivation
buildRedistPackage =
callPackage: pname:
let
redistribRelease = redistribManifest.${pname};
featureRelease = featureManifest.${pname};
drv =
(callPackage ../generic-builders/manifest.nix {
# We pass the whole release to the builder because it has logic to handle
# the case we're trying to build on an unsupported platform.
inherit
pname
redistName
redistribRelease
featureRelease
;
}).overrideAttrs
(prevAttrs: {
# Add the package-specific license.
meta = prevAttrs.meta // {
license =
let
licensePath =
if redistribRelease.license_path != null then
redistribRelease.license_path
else
"${pname}/LICENSE.txt";
url = "https://developer.download.nvidia.com/compute/cuda/redist/${licensePath}";
in
lib.licenses.nvidiaCudaRedist // { inherit url; };
};
});
in
drv;
# Build all the redist packages given final and prev.
redistPackages =
final: _prev:
# Wrap the whole thing in an optionalAttrs so we can return an empty set if the CUDA version
# is not supported.
# NOTE: We cannot include the call to optionalAttrs *in* the pipe as we would strictly evaluate the
# attrNames before we check if the CUDA version is supported.
attrsets.optionalAttrs cudaVersionMappingExists (
trivial.pipe featureManifest [
# Get all the package names
builtins.attrNames
# Build the redist packages
(trivial.flip attrsets.genAttrs (buildRedistPackage final.callPackage))
]
);
in
redistPackages
|