pkgs/top-level/cuda-packages.nix

   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 overriden.
  15 # - `final` should only be used to access `callPackage` to build new packages.
  16 # - Attribute names should be computable without relying on `final`.
  17 #   - Extensions should take arguments to build attribute names before relying on `final`.
  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   callPackage,
  25   cudaVersion,
  26   lib,
  27   newScope,
  28   pkgs,
  29   config,
  30 }:
  31 let
  32   inherit (lib)
  33     attrsets
  34     customisation
  35     fixedPoints
  36     lists
  37     strings
  38     trivial
  39     versions
  40     ;
  41   # Backbone
  42   gpus = builtins.import ../development/cuda-modules/gpus.nix;
  43   nvccCompatibilities = builtins.import ../development/cuda-modules/nvcc-compatibilities.nix;
  44   flags = callPackage ../development/cuda-modules/flags.nix { inherit cudaVersion gpus; };
  45   passthruFunction = final: ({
  46     inherit cudaVersion lib pkgs;
  47     inherit gpus nvccCompatibilities flags;
  48     cudaMajorVersion = versions.major cudaVersion;
  49     cudaMajorMinorVersion = versions.majorMinor cudaVersion;
  50     cudaOlder = strings.versionOlder cudaVersion;
  51     cudaAtLeast = strings.versionAtLeast cudaVersion;
  52
  53     # Maintain a reference to the final cudaPackages.
  54     # Without this, if we use `final.callPackage` and a package accepts `cudaPackages` as an
  55     # argument, it's provided with `cudaPackages` from the top-level scope, which is not what we
  56     # want. We want to provide the `cudaPackages` from the final scope -- that is, the *current*
  57     # scope. However, we also want to prevent `pkgs/top-level/release-attrpaths-superset.nix` from
  58     # recursing more than one level here.
  59     cudaPackages = final // {
  60       __attrsFailEvaluation = true;
  61     };
  62
  63     # TODO(@connorbaker): `cudaFlags` is an alias for `flags` which should be removed in the future.
  64     cudaFlags = flags;
  65
  66     # Exposed as cudaPackages.backendStdenv.
  67     # This is what nvcc uses as a backend,
  68     # and it has to be an officially supported one (e.g. gcc11 for cuda11).
  69     #
  70     # It, however, propagates current stdenv's libstdc++ to avoid "GLIBCXX_* not found errors"
  71     # when linked with other C++ libraries.
  72     # E.g. for cudaPackages_11_8 we use gcc11 with gcc12's libstdc++
  73     # Cf. https://github.com/NixOS/nixpkgs/pull/218265 for context
  74     backendStdenv = final.callPackage ../development/cuda-modules/backend-stdenv.nix { };
  75
  76     # Loose packages
  77
  78     # TODO: Move to aliases.nix once all Nixpkgs has migrated to the splayed CUDA packages
  79     cudatoolkit = final.callPackage ../development/cuda-modules/cudatoolkit/redist-wrapper.nix { };
  80     cudatoolkit-legacy-runfile = final.callPackage ../development/cuda-modules/cudatoolkit { };
  81
  82     saxpy = final.callPackage ../development/cuda-modules/saxpy { };
  83     nccl = final.callPackage ../development/cuda-modules/nccl { };
  84     nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { };
  85
  86     tests =
  87       let
  88         bools = [
  89           true
  90           false
  91         ];
  92         configs = {
  93           openCVFirst = bools;
  94           useOpenCVDefaultCuda = bools;
  95           useTorchDefaultCuda = bools;
  96         };
  97         builder =
  98           {
  99             openCVFirst,
 100             useOpenCVDefaultCuda,
 101             useTorchDefaultCuda,
 102           }@config:
 103           {
 104             name = strings.concatStringsSep "-" (
 105               [
 106                 "test"
 107                 (if openCVFirst then "opencv" else "torch")
 108               ]
 109               ++ lists.optionals (if openCVFirst then useOpenCVDefaultCuda else useTorchDefaultCuda) [
 110                 "with-default-cuda"
 111               ]
 112               ++ [
 113                 "then"
 114                 (if openCVFirst then "torch" else "opencv")
 115               ]
 116               ++ lists.optionals (if openCVFirst then useTorchDefaultCuda else useOpenCVDefaultCuda) [
 117                 "with-default-cuda"
 118               ]
 119             );
 120             value = final.callPackage ../development/cuda-modules/tests/opencv-and-torch config;
 121           };
 122       in
 123       attrsets.listToAttrs (attrsets.mapCartesianProduct builder configs);
 124
 125     writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-test-python.nix { };
 126   });
 127
 128   mkVersionedPackageName =
 129     name: version:
 130     strings.concatStringsSep "_" [
 131       name
 132       (strings.replaceStrings [ "." ] [ "_" ] (versions.majorMinor version))
 133     ];
 134
 135   composedExtension = fixedPoints.composeManyExtensions (
 136     [
 137       (import ../development/cuda-modules/setup-hooks/extension.nix)
 138       (callPackage ../development/cuda-modules/cuda/extension.nix { inherit cudaVersion; })
 139       (import ../development/cuda-modules/cuda/overrides.nix)
 140       (callPackage ../development/cuda-modules/generic-builders/multiplex.nix {
 141         inherit cudaVersion flags mkVersionedPackageName;
 142         pname = "cudnn";
 143         releasesModule = ../development/cuda-modules/cudnn/releases.nix;
 144         shimsFn = ../development/cuda-modules/cudnn/shims.nix;
 145         fixupFn = ../development/cuda-modules/cudnn/fixup.nix;
 146       })
 147       (callPackage ../development/cuda-modules/cutensor/extension.nix {
 148         inherit cudaVersion flags mkVersionedPackageName;
 149       })
 150       (callPackage ../development/cuda-modules/generic-builders/multiplex.nix {
 151         inherit cudaVersion flags mkVersionedPackageName;
 152         pname = "tensorrt";
 153         releasesModule = ../development/cuda-modules/tensorrt/releases.nix;
 154         shimsFn = ../development/cuda-modules/tensorrt/shims.nix;
 155         fixupFn = ../development/cuda-modules/tensorrt/fixup.nix;
 156       })
 157       (callPackage ../development/cuda-modules/cuda-samples/extension.nix { inherit cudaVersion; })
 158       (callPackage ../development/cuda-modules/cuda-library-samples/extension.nix { })
 159     ]
 160     ++ lib.optionals config.allowAliases [ (import ../development/cuda-modules/aliases.nix) ]
 161   );
 162
 163   cudaPackages = customisation.makeScope newScope (
 164     fixedPoints.extends composedExtension passthruFunction
 165   );
 166 in
 167 # We want to warn users about the upcoming deprecation of old CUDA
 168 # versions, without breaking Nixpkgs CI with evaluation warnings. This
 169 # gross hack ensures that the warning only triggers if aliases are
 170 # enabled, which is true by default, but not for ofborg.
 171 lib.warnIf (cudaPackages.cudaOlder "12.0" && config.allowAliases)
 172   "CUDA versions older than 12.0 will be removed in Nixpkgs 25.05; see the 24.11 release notes for more information"
 173   cudaPackages