cudaPackages: introduce and use lib.packagesFromDirectoryRecursive

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>

pkgs/development/cuda-modules/packages/autoAddCudaCompatRunpath/package.nix

@@ -0,0 +1,27 @@
+# autoAddCudaCompatRunpath hook must be added AFTER `setupCudaHook`. Both
+# hooks prepend a path with `libcuda.so` to the `DT_RUNPATH` section of
+# patched elf files, but `cuda_compat` path must take precedence (otherwise,
+# it doesn't have any effect) and thus appear first. Meaning this hook must be
+# executed last.
+{
+  autoFixElfFiles,
+  cuda_compat ? null,
+  flags,
+  lib,
+  makeSetupHook,
+}:
+makeSetupHook {
+  name = "auto-add-cuda-compat-runpath-hook";
+  propagatedBuildInputs = [ autoFixElfFiles ];
+
+  substitutions = {
+    # Hotfix Ofborg evaluation
+    libcudaPath = if flags.isJetsonBuild then "${cuda_compat}/compat" else null;
+  };
+
+  meta.broken = !flags.isJetsonBuild;
+
+  # Pre-cuda_compat CUDA release:
+  meta.badPlatforms = lib.optionals (cuda_compat == null) lib.platforms.all;
+  meta.platforms = cuda_compat.meta.platforms or [ ];
+} ./auto-add-cuda-compat-runpath.sh

pkgs/development/cuda-modules/packages/backendStdenv.nix

@@ -1,3 +1,11 @@
+# Exposed as cudaPackages.backendStdenv.
+# This is what nvcc uses as a backend,
+# and it has to be an officially supported one (e.g. gcc11 for cuda11).
+#
+# It, however, propagates current stdenv's libstdc++ to avoid "GLIBCXX_* not found errors"
+# when linked with other C++ libraries.
+# E.g. for cudaPackages_11_8 we use gcc11 with gcc12's libstdc++
+# Cf. https://github.com/NixOS/nixpkgs/pull/218265 for context
 {
   cudaVersion,
   lib,

pkgs/development/cuda-modules/packages/markForCudatoolkitRootHook/package.nix

@@ -0,0 +1,4 @@
+# Internal hook, used by cudatoolkit and cuda redist packages
+# to accommodate automatic CUDAToolkit_ROOT construction
+{ makeSetupHook }:
+makeSetupHook { name = "mark-for-cudatoolkit-root-hook"; } ./mark-for-cudatoolkit-root-hook.sh

pkgs/development/cuda-modules/packages/setupCudaHook/package.nix

@@ -0,0 +1,14 @@
+# Currently propagated by cuda_nvcc or cudatoolkit, rather than used directly
+{ makeSetupHook, backendStdenv }:
+makeSetupHook {
+  name = "setup-cuda-hook";
+
+  substitutions.setupCudaHook = placeholder "out";
+
+  # Point NVCC at a compatible compiler
+  substitutions.ccRoot = "${backendStdenv.cc}";
+
+  # Required in addition to ccRoot as otherwise bin/gcc is looked up
+  # when building CMakeCUDACompilerId.cu
+  substitutions.ccFullPath = "${backendStdenv.cc}/bin/${backendStdenv.cc.targetPrefix}c++";
+} ./setup-cuda-hook.sh

pkgs/development/cuda-modules/setup-hooks/extension.nix

@@ -1,55 +0,0 @@
-final: _: {
-  # Internal hook, used by cudatoolkit and cuda redist packages
-  # to accommodate automatic CUDAToolkit_ROOT construction
-  markForCudatoolkitRootHook = final.callPackage (
-    { makeSetupHook }:
-    makeSetupHook { name = "mark-for-cudatoolkit-root-hook"; } ./mark-for-cudatoolkit-root-hook.sh
-  ) { };
-
-  # Currently propagated by cuda_nvcc or cudatoolkit, rather than used directly
-  setupCudaHook = (
-    final.callPackage (
-      { makeSetupHook, backendStdenv }:
-      makeSetupHook {
-        name = "setup-cuda-hook";
-
-        substitutions.setupCudaHook = placeholder "out";
-
-        # Point NVCC at a compatible compiler
-        substitutions.ccRoot = "${backendStdenv.cc}";
-
-        # Required in addition to ccRoot as otherwise bin/gcc is looked up
-        # when building CMakeCUDACompilerId.cu
-        substitutions.ccFullPath = "${backendStdenv.cc}/bin/${backendStdenv.cc.targetPrefix}c++";
-      } ./setup-cuda-hook.sh
-    ) { }
-  );
-
-  # autoAddCudaCompatRunpath hook must be added AFTER `setupCudaHook`. Both
-  # hooks prepend a path with `libcuda.so` to the `DT_RUNPATH` section of
-  # patched elf files, but `cuda_compat` path must take precedence (otherwise,
-  # it doesn't have any effect) and thus appear first. Meaning this hook must be
-  # executed last.
-  autoAddCudaCompatRunpath = final.callPackage (
-    {
-      makeSetupHook,
-      autoFixElfFiles,
-      cuda_compat ? null,
-    }:
-    makeSetupHook {
-      name = "auto-add-cuda-compat-runpath-hook";
-      propagatedBuildInputs = [ autoFixElfFiles ];
-
-      substitutions = {
-        # Hotfix Ofborg evaluation
-        libcudaPath = if final.flags.isJetsonBuild then "${cuda_compat}/compat" else null;
-      };
-
-      meta.broken = !final.flags.isJetsonBuild;
-
-      # Pre-cuda_compat CUDA release:
-      meta.badPlatforms = final.lib.optionals (cuda_compat == null) final.lib.platforms.all;
-      meta.platforms = cuda_compat.meta.platforms or [ ];
-    } ./auto-add-cuda-compat-runpath.sh
-  ) { };
-}

pkgs/top-level/cuda-packages.nix

@@ -35,14 +35,13 @@ let
     fixedPoints
     lists
     strings
-    trivial
     versions
     ;
   # Backbone
   gpus = builtins.import ../development/cuda-modules/gpus.nix;
   nvccCompatibilities = builtins.import ../development/cuda-modules/nvcc-compatibilities.nix;
   flags = callPackage ../development/cuda-modules/flags.nix { inherit cudaVersion gpus; };
-  passthruFunction = final: ({
+  passthruFunction = final: {
     inherit cudaVersion lib pkgs;
     inherit gpus nvccCompatibilities flags;
     cudaMajorVersion = versions.major cudaVersion;
@@ -63,27 +62,15 @@ let
     # TODO(@connorbaker): `cudaFlags` is an alias for `flags` which should be removed in the future.
     cudaFlags = flags;
 
-    # Exposed as cudaPackages.backendStdenv.
-    # This is what nvcc uses as a backend,
-    # and it has to be an officially supported one (e.g. gcc11 for cuda11).
-    #
-    # It, however, propagates current stdenv's libstdc++ to avoid "GLIBCXX_* not found errors"
-    # when linked with other C++ libraries.
-    # E.g. for cudaPackages_11_8 we use gcc11 with gcc12's libstdc++
-    # Cf. https://github.com/NixOS/nixpkgs/pull/218265 for context
-    backendStdenv = final.callPackage ../development/cuda-modules/backend-stdenv.nix { };
-
     # Loose packages
+    # Barring packages which share a home (e.g., cudatoolkit and cudatoolkit-legacy-runfile), new packages
+    # should be added to ../development/cuda-modules/packages in "by-name" style, where they will be automatically
+    # discovered and added to the package set.
 
     # TODO: Move to aliases.nix once all Nixpkgs has migrated to the splayed CUDA packages
     cudatoolkit = final.callPackage ../development/cuda-modules/cudatoolkit/redist-wrapper.nix { };
     cudatoolkit-legacy-runfile = final.callPackage ../development/cuda-modules/cudatoolkit { };
 
-    cudnn-frontend = final.callPackage ../development/cuda-modules/cudnn-frontend/default.nix { };
-    saxpy = final.callPackage ../development/cuda-modules/saxpy { };
-    nccl = final.callPackage ../development/cuda-modules/nccl { };
-    nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { };
-
     tests =
       let
         bools = [
@@ -122,9 +109,7 @@ let
           };
       in
       attrsets.listToAttrs (attrsets.mapCartesianProduct builder configs);
-
-    writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-test-python.nix { };
-  });
+  };
 
   mkVersionedPackageName =
     name: version:
@@ -135,7 +120,13 @@ let
 
   composedExtension = fixedPoints.composeManyExtensions (
     [
-      (import ../development/cuda-modules/setup-hooks/extension.nix)
+      (
+        final: _:
+        lib.packagesFromDirectoryRecursive {
+          inherit (final) callPackage;
+          directory = ../development/cuda-modules/packages;
+        }
+      )
       (callPackage ../development/cuda-modules/cuda/extension.nix { inherit cudaVersion; })
       (import ../development/cuda-modules/cuda/overrides.nix)
       (callPackage ../development/cuda-modules/generic-builders/multiplex.nix {