about summary refs log tree commit diff
path: root/pkgs/development/cuda-modules/cuda/extension.nix
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