tjh.dev
1# Notes:
2#
3# Silvan (Tweag) covered some things on recursive attribute sets in the Nix Hour:
4# https://www.youtube.com/watch?v=BgnUFtd1Ivs
5#
6# I (@connorbaker) highly recommend watching it.
7#
8# Most helpful comment regarding recursive attribute sets:
9#
10# https://github.com/NixOS/nixpkgs/pull/256324#issuecomment-1749935979
11#
12# To summarize:
13#
14# - `prev` should only be used to access attributes which are going to be overridden.
15# - `final` should only be used to access `callPackage` to build new packages.
16# - Attribute names are evaluated eagerly ("NAMESET STRICTNESS").
17# - Extensions must not depend on `final` when computing names and count of new attributes.
18#
19# Silvan's recommendation then is to explicitly use `callPackage` to provide everything our
20# extensions need to compute the attribute names, without relying on `final`.
21#
22# I've (@connorbaker) attempted to do that, though I'm unsure of how this will interact with overrides.
23{
24 config,
25 _cuda,
26 cudaMajorMinorVersion,
27 lib,
28 pkgs,
29 stdenv,
30 runCommand,
31}:
32let
33 inherit (lib)
34 attrsets
35 customisation
36 fixedPoints
37 lists
38 strings
39 versions
40 ;
41
42 cudaLib = _cuda.lib;
43
44 # Since Jetson capabilities are never built by default, we can check if any of them were requested
45 # through final.config.cudaCapabilities and use that to determine if we should change some manifest versions.
46 # Copied from backendStdenv.
47 jetsonCudaCapabilities = lib.filter (
48 cudaCapability: _cuda.db.cudaCapabilityToInfo.${cudaCapability}.isJetson
49 ) _cuda.db.allSortedCudaCapabilities;
50 hasJetsonCudaCapability =
51 lib.intersectLists jetsonCudaCapabilities (config.cudaCapabilities or [ ]) != [ ];
52 redistSystem = _cuda.lib.getRedistSystem hasJetsonCudaCapability stdenv.hostPlatform.system;
53
54 # We must use an instance of Nixpkgs where the CUDA package set we're building is the default; if we do not, members
55 # of the versioned, non-default package sets may rely on (transitively) members of the default, unversioned CUDA
56 # package set.
57 # See `Using cudaPackages.pkgs` in doc/languages-frameworks/cuda.section.md for more information.
58 pkgs' =
59 let
60 cudaPackagesUnversionedName = "cudaPackages";
61 cudaPackagesMajorVersionName = cudaLib.mkVersionedName cudaPackagesUnversionedName (
62 versions.major cudaMajorMinorVersion
63 );
64 cudaPackagesMajorMinorVersionName = cudaLib.mkVersionedName cudaPackagesUnversionedName cudaMajorMinorVersion;
65 in
66 # If the CUDA version of pkgs matches our CUDA version, we are constructing the default package set and can use
67 # pkgs without modification.
68 if pkgs.cudaPackages.cudaMajorMinorVersion == cudaMajorMinorVersion then
69 pkgs
70 else
71 pkgs.extend (
72 final: _: {
73 recurseForDerivations = false;
74 # The CUDA package set will be available as cudaPackages_x_y, so we need only update the aliases for the
75 # minor-versioned and unversioned package sets.
76 # cudaPackages_x = cudaPackages_x_y
77 ${cudaPackagesMajorVersionName} = final.${cudaPackagesMajorMinorVersionName};
78 # cudaPackages = cudaPackages_x
79 ${cudaPackagesUnversionedName} = final.${cudaPackagesMajorVersionName};
80 }
81 );
82
83 passthruFunction = final: {
84 # NOTE:
85 # It is important that _cuda is not part of the package set fixed-point. As described by
86 # @SomeoneSerge:
87 # > The layering should be: configuration -> (identifies/is part of) cudaPackages -> (is built using) cudaLib.
88 # > No arrows should point in the reverse directions.
89 # That is to say that cudaLib should only know about package sets and configurations, because it implements
90 # functionality for interpreting configurations, resolving them against data, and constructing package sets.
91 # This decision is driven both by a separation of concerns and by "NAMESET STRICTNESS" (see above).
92 # Also see the comment in `pkgs/top-level/all-packages.nix` about the `_cuda` attribute.
93
94 inherit cudaMajorMinorVersion;
95
96 pkgs = pkgs';
97
98 cudaNamePrefix = "cuda${cudaMajorMinorVersion}";
99
100 cudaMajorVersion = versions.major cudaMajorMinorVersion;
101 cudaOlder = strings.versionOlder cudaMajorMinorVersion;
102 cudaAtLeast = strings.versionAtLeast cudaMajorMinorVersion;
103
104 flags =
105 cudaLib.formatCapabilities {
106 inherit (final.backendStdenv) cudaCapabilities cudaForwardCompat;
107 inherit (_cuda.db) cudaCapabilityToInfo;
108 }
109 # TODO(@connorbaker): Enable the corresponding warnings in `../development/cuda-modules/aliases.nix` after some
110 # time to allow users to migrate to cudaLib and backendStdenv.
111 // {
112 inherit (cudaLib) dropDots;
113 cudaComputeCapabilityToName =
114 cudaCapability: _cuda.db.cudaCapabilityToInfo.${cudaCapability}.archName;
115 dropDot = cudaLib.dropDots;
116 isJetsonBuild = final.backendStdenv.hasJetsonCudaCapability;
117 };
118
119 # Loose packages
120 # Barring packages which share a home (e.g., cudatoolkit and cudatoolkit-legacy-runfile), new packages
121 # should be added to ../development/cuda-modules/packages in "by-name" style, where they will be automatically
122 # discovered and added to the package set.
123
124 # TODO: Move to aliases.nix once all Nixpkgs has migrated to the splayed CUDA packages
125 cudatoolkit = final.callPackage ../development/cuda-modules/cudatoolkit/redist-wrapper.nix { };
126 cudatoolkit-legacy-runfile = final.callPackage ../development/cuda-modules/cudatoolkit { };
127
128 tests =
129 let
130 bools = [
131 true
132 false
133 ];
134 configs = {
135 openCVFirst = bools;
136 useOpenCVDefaultCuda = bools;
137 useTorchDefaultCuda = bools;
138 };
139 builder =
140 {
141 openCVFirst,
142 useOpenCVDefaultCuda,
143 useTorchDefaultCuda,
144 }@config:
145 {
146 name = strings.concatStringsSep "-" (
147 [
148 "test"
149 (if openCVFirst then "opencv" else "torch")
150 ]
151 ++ lists.optionals (if openCVFirst then useOpenCVDefaultCuda else useTorchDefaultCuda) [
152 "with-default-cuda"
153 ]
154 ++ [
155 "then"
156 (if openCVFirst then "torch" else "opencv")
157 ]
158 ++ lists.optionals (if openCVFirst then useTorchDefaultCuda else useOpenCVDefaultCuda) [
159 "with-default-cuda"
160 ]
161 );
162 value = final.callPackage ../development/cuda-modules/tests/opencv-and-torch config;
163 };
164 in
165 attrsets.listToAttrs (attrsets.mapCartesianProduct builder configs)
166 // {
167 flags = final.callPackage ../development/cuda-modules/tests/flags.nix { };
168 };
169 };
170
171 composedExtension = fixedPoints.composeManyExtensions (
172 [
173 (
174 final: _:
175 {
176 cuda_compat = runCommand "cuda_compat" { meta.platforms = [ ]; } "false"; # Prevent missing attribute errors
177 }
178 // lib.packagesFromDirectoryRecursive {
179 inherit (final) callPackage;
180 directory = ../development/cuda-modules/packages;
181 }
182 )
183 (import ../development/cuda-modules/cuda/extension.nix { inherit cudaMajorMinorVersion lib; })
184 (import ../development/cuda-modules/generic-builders/multiplex.nix {
185 inherit
186 cudaLib
187 cudaMajorMinorVersion
188 lib
189 redistSystem
190 stdenv
191 ;
192 pname = "cudnn";
193 redistName = "cudnn";
194 releasesModule = ../development/cuda-modules/cudnn/releases.nix;
195 shimsFn = ../development/cuda-modules/cudnn/shims.nix;
196 })
197 (import ../development/cuda-modules/cutensor/extension.nix {
198 inherit
199 cudaLib
200 cudaMajorMinorVersion
201 lib
202 redistSystem
203 ;
204 })
205 (import ../development/cuda-modules/cusparselt/extension.nix {
206 inherit
207 cudaLib
208 lib
209 redistSystem
210 ;
211 })
212 (import ../development/cuda-modules/generic-builders/multiplex.nix {
213 inherit
214 cudaLib
215 cudaMajorMinorVersion
216 lib
217 redistSystem
218 stdenv
219 ;
220 pname = "tensorrt";
221 redistName = "tensorrt";
222 releasesModule = ../development/cuda-modules/tensorrt/releases.nix;
223 shimsFn = ../development/cuda-modules/tensorrt/shims.nix;
224 })
225 (import ../development/cuda-modules/cuda-samples/extension.nix {
226 inherit cudaMajorMinorVersion lib stdenv;
227 })
228 (import ../development/cuda-modules/cuda-library-samples/extension.nix { inherit lib stdenv; })
229 ]
230 ++ lib.optionals config.allowAliases [
231 (import ../development/cuda-modules/aliases.nix { inherit lib; })
232 ]
233 ++ _cuda.extensions
234 );
235
236 cudaPackages = customisation.makeScope pkgs'.newScope (
237 fixedPoints.extends composedExtension passthruFunction
238 );
239in
240# We want to warn users about the upcoming deprecation of old CUDA
241# versions, without breaking Nixpkgs CI with evaluation warnings. This
242# gross hack ensures that the warning only triggers if aliases are
243# enabled, which is true by default, but not for ofborg.
244lib.warnIf (cudaPackages.cudaOlder "12.0" && config.allowAliases)
245 "CUDA versions older than 12.0 will be removed in Nixpkgs 25.05; see the 24.11 release notes for more information"
246 cudaPackages