tjh.dev
1# This is what nvcc uses as a backend,
2# and it has to be an officially supported one (e.g. gcc11 for cuda11).
3#
4# It, however, propagates current stdenv's libstdc++ to avoid "GLIBCXX_* not found errors"
5# when linked with other C++ libraries.
6# E.g. for cudaPackages_11_8 we use gcc11 with gcc12's libstdc++
7# Cf. https://github.com/NixOS/nixpkgs/pull/218265 for context
8{
9 config,
10 _cuda,
11 cudaMajorMinorVersion,
12 lib,
13 pkgs,
14 stdenv,
15 stdenvAdapters,
16}:
17let
18 inherit (builtins) toJSON;
19 inherit (_cuda.db) allSortedCudaCapabilities cudaCapabilityToInfo nvccCompatibilities;
20 inherit (_cuda.lib)
21 _cudaCapabilityIsDefault
22 _cudaCapabilityIsSupported
23 _evaluateAssertions
24 getRedistSystem
25 mkVersionedName
26 ;
27 inherit (lib) addErrorContext;
28 inherit (lib.customisation) extendDerivation;
29 inherit (lib.lists) filter intersectLists subtractLists;
30
31 # NOTE: By virtue of processing a sorted list (allSortedCudaCapabilities), our groups will be sorted.
32
33 architectureSpecificCudaCapabilities = filter (
34 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isArchitectureSpecific
35 ) allSortedCudaCapabilities;
36
37 familySpecificCudaCapabilities = filter (
38 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isFamilySpecific
39 ) allSortedCudaCapabilities;
40
41 jetsonCudaCapabilities = filter (
42 cudaCapability: cudaCapabilityToInfo.${cudaCapability}.isJetson
43 ) allSortedCudaCapabilities;
44
45 passthruExtra = {
46 nvccHostCCMatchesStdenvCC = backendStdenv.cc == stdenv.cc;
47
48 # The Nix system of the host platform.
49 hostNixSystem = stdenv.hostPlatform.system;
50
51 # The Nix system of the host platform for the CUDA redistributable.
52 hostRedistSystem = getRedistSystem passthruExtra.hasJetsonCudaCapability stdenv.hostPlatform.system;
53
54 # Sets whether packages should be built with forward compatibility.
55 # TODO(@connorbaker): If the requested CUDA capabilities are not supported by the current CUDA version,
56 # should we throw an evaluation warning and build with forward compatibility?
57 cudaForwardCompat = config.cudaForwardCompat or true;
58
59 # CUDA capabilities which are supported by the current CUDA version.
60 supportedCudaCapabilities = filter (
61 cudaCapability:
62 _cudaCapabilityIsSupported cudaMajorMinorVersion cudaCapabilityToInfo.${cudaCapability}
63 ) allSortedCudaCapabilities;
64
65 # Find the default set of capabilities for this CUDA version using the list of supported capabilities.
66 # Includes only baseline capabilities.
67 defaultCudaCapabilities = filter (
68 cudaCapability:
69 _cudaCapabilityIsDefault cudaMajorMinorVersion cudaCapabilityToInfo.${cudaCapability}
70 ) passthruExtra.supportedCudaCapabilities;
71
72 # The resolved requested or default CUDA capabilities.
73 cudaCapabilities =
74 if config.cudaCapabilities or [ ] != [ ] then
75 config.cudaCapabilities
76 else
77 passthruExtra.defaultCudaCapabilities;
78
79 # Requested architecture-specific CUDA capabilities.
80 requestedArchitectureSpecificCudaCapabilities = intersectLists architectureSpecificCudaCapabilities passthruExtra.cudaCapabilities;
81
82 # Whether the requested CUDA capabilities include architecture-specific CUDA capabilities.
83 hasArchitectureSpecificCudaCapability =
84 passthruExtra.requestedArchitectureSpecificCudaCapabilities != [ ];
85
86 # Requested family-specific CUDA capabilities.
87 requestedFamilySpecificCudaCapabilities = intersectLists familySpecificCudaCapabilities passthruExtra.cudaCapabilities;
88
89 # Whether the requested CUDA capabilities include family-specific CUDA capabilities.
90 hasFamilySpecificCudaCapability = passthruExtra.requestedFamilySpecificCudaCapabilities != [ ];
91
92 # Requested Jetson CUDA capabilities.
93 requestedJetsonCudaCapabilities = intersectLists jetsonCudaCapabilities passthruExtra.cudaCapabilities;
94
95 # Whether the requested CUDA capabilities include Jetson CUDA capabilities.
96 hasJetsonCudaCapability = passthruExtra.requestedJetsonCudaCapabilities != [ ];
97 };
98
99 assertions =
100 let
101 # Jetson devices cannot be targeted by the same binaries which target non-Jetson devices. While
102 # NVIDIA provides both `linux-aarch64` and `linux-sbsa` packages, which both target `aarch64`,
103 # they are built with different settings and cannot be mixed.
104 jetsonMesssagePrefix = "Jetson CUDA capabilities (${toJSON passthruExtra.requestedJetsonCudaCapabilities})";
105
106 # Remove all known capabilities from the user's list to find unrecognized capabilities.
107 unrecognizedCudaCapabilities = subtractLists allSortedCudaCapabilities passthruExtra.cudaCapabilities;
108
109 # Remove all supported capabilities from the user's list to find unsupported capabilities.
110 unsupportedCudaCapabilities = subtractLists passthruExtra.supportedCudaCapabilities passthruExtra.cudaCapabilities;
111 in
112 [
113 {
114 message = "Unrecognized CUDA capabilities: ${toJSON unrecognizedCudaCapabilities}";
115 assertion = unrecognizedCudaCapabilities == [ ];
116 }
117 {
118 message = "Unsupported CUDA capabilities: ${toJSON unsupportedCudaCapabilities}";
119 assertion = unsupportedCudaCapabilities == [ ];
120 }
121 {
122 message =
123 "${jetsonMesssagePrefix} require hostPlatform (currently ${passthruExtra.hostNixSystem}) "
124 + "to be aarch64-linux";
125 assertion = passthruExtra.hasJetsonCudaCapability -> passthruExtra.hostNixSystem == "aarch64-linux";
126 }
127 {
128 message =
129 let
130 # Find the capabilities which are not Jetson capabilities.
131 requestedNonJetsonCudaCapabilities = subtractLists (
132 passthruExtra.requestedJetsonCudaCapabilities
133 ++ passthruExtra.requestedArchitectureSpecificCudaCapabilities
134 ++ passthruExtra.requestedFamilySpecificCudaCapabilities
135 ) passthruExtra.cudaCapabilities;
136 in
137 "${jetsonMesssagePrefix} cannot be specified with non-Jetson capabilities "
138 + "(${toJSON requestedNonJetsonCudaCapabilities})";
139 assertion =
140 passthruExtra.hasJetsonCudaCapability
141 -> passthruExtra.requestedJetsonCudaCapabilities == passthruExtra.cudaCapabilities;
142 }
143 ];
144
145 assertCondition = addErrorContext "while evaluating ${mkVersionedName "cudaPackages" cudaMajorMinorVersion}.backendStdenv" (
146 _evaluateAssertions assertions
147 );
148
149 backendStdenv =
150 stdenvAdapters.useLibsFrom stdenv
151 pkgs."gcc${nvccCompatibilities.${cudaMajorMinorVersion}.gcc.maxMajorVersion}Stdenv";
152in
153# TODO: Consider testing whether we in fact use the newer libstdc++
154extendDerivation assertCondition passthruExtra backendStdenv