cuda-modules: rewrite setup hooks

pkgs/development/cuda-modules/aliases.nix

@@ -6,7 +6,6 @@ let
     final.lib.warn "cudaPackages.${oldName} is deprecated, use ${newName} instead" newPkg;
 in
 {
-
   # Deprecated: an alias kept for compatibility. Consider removing after 24.05
   autoFixElfFiles = mkRenamed "autoFixElfFiles" "pkgs.autoFixElfFiles" final.pkgs.autoFixElfFiles; # Added 2024-03-30
   autoAddDriverRunpath =
@@ -15,4 +14,8 @@ in
   autoAddOpenGLRunpathHook =
     mkRenamed "autoAddOpenGLRunpathHook" "pkgs.autoAddDriverRunpath"
       final.pkgs.autoAddDriverRunpath; # Added 2024-03-30
+  markForCudatoolkitRootHook =
+    mkRenamed "markForCudatoolkitRootHook" "cudaPackages.markForCudatoolkitRoot"
+      final.markForCudatoolkitRoot; # Added 2024-04-19
+  setupCudaHook = mkRenamed "setupCudaHook" "cudaPackages.setupCuda" final.setupCuda; # Added 2024-04-19
 }

pkgs/development/cuda-modules/cuda/overrides.nix

@@ -170,7 +170,7 @@ filterAndCreateOverrides {
       backendStdenv,
       cuda_cudart,
       lib,
-      setupCudaHook,
+      setupCuda,
     }:
     prevAttrs: {
       # Patch the nvcc.profile.
@@ -214,7 +214,7 @@ filterAndCreateOverrides {
       # `propagatedNativeBuildInputs`, it stops being propagated to downstream packages during their build because
       # setup hooks in `propagatedNativeBuildInputs` are not designed to affect the runtime or build environment of
       # dependencies; they are only meant to affect the build environment of the package that directly includes them.
-      propagatedBuildInputs = (prevAttrs.propagatedBuildInputs or [ ]) ++ [ setupCudaHook ];
+      propagatedBuildInputs = (prevAttrs.propagatedBuildInputs or [ ]) ++ [ setupCuda ];
 
       postInstall =
         (prevAttrs.postInstall or "")

pkgs/development/cuda-modules/cudatoolkit/default.nix

@@ -20,15 +20,15 @@
   libkrb5,
   krb5,
   makeWrapper,
-  markForCudatoolkitRootHook,
+  markForCudatoolkitRoot,
   ncurses5,
   numactl,
   nss,
   patchelf,
   perl,
   python3, # FIXME: CUDAToolkit 10 may still need python27
   pulseaudio,
-  setupCudaHook,
+  setupCuda,
   stdenv,
   backendStdenv, # E.g. gcc11Stdenv, set in extension.nix
   unixODBC,
@@ -77,11 +77,11 @@ backendStdenv.mkDerivation rec {
       addOpenGLRunpath
       autoPatchelfHook
       autoAddDriverRunpath
-      markForCudatoolkitRootHook
+      markForCudatoolkitRoot
     ]
     ++ lib.optionals (lib.versionOlder version "11") [ libsForQt5.wrapQtAppsHook ]
     ++ lib.optionals (lib.versionAtLeast version "11.8") [ qt6Packages.wrapQtAppsHook ];
-  propagatedBuildInputs = [ setupCudaHook ];
+  propagatedBuildInputs = [ setupCuda ];
   buildInputs =
     lib.optionals (lib.versionOlder version "11") [
       libsForQt5.qt5.qtwebengine

pkgs/development/cuda-modules/generic-builders/manifest.nix

@@ -7,7 +7,7 @@
   fetchurl,
   lib,
   lndir,
-  markForCudatoolkitRootHook,
+  markForCudatoolkitRoot,
   flags,
   stdenv,
   # Builder-specific arguments
@@ -200,7 +200,7 @@ backendStdenv.mkDerivation (finalAttrs: {
       # directory to the rpath of all ELF binaries.
       # Check e.g. with `patchelf --print-rpath path/to/my/binary
       autoAddDriverRunpath
-      markForCudatoolkitRootHook
+      markForCudatoolkitRoot
     ]
     # autoAddCudaCompatRunpath depends on cuda_compat and would cause
     # infinite recursion if applied to `cuda_compat` itself (beside the fact

pkgs/development/cuda-modules/saxpy/default.nix

@@ -15,7 +15,6 @@ let
     cudatoolkit
     flags
     libcublas
-    setupCudaHook
     ;
   inherit (lib) getDev getLib getOutput;
 in

pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath/default.nix

@@ -0,0 +1,22 @@
+# autoAddCudaCompatRunpath hook must be added AFTER `setupCuda`. 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.libcudaPath = lib.optionalString flags.isJetsonBuild "${cuda_compat}/compat";
+  meta = {
+    broken = !flags.isJetsonBuild;
+    badPlatforms = lib.optionals (cuda_compat == null) lib.platforms.all;
+    platforms = cuda_compat.meta.platforms or [ ];
+  };
+} ./auto-add-cuda-compat-runpath.sh

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

@@ -1,55 +1,5 @@
 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
-  ) { };
+  autoAddCudaCompatRunpath = final.callPackage ./auto-add-cuda-compat-runpath { };
+  markForCudatoolkitRoot = final.callPackage ./mark-for-cudatoolkit-root { };
+  setupCuda = final.callPackage ./setup-cuda { };
 }

pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root/default.nix

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

pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root/mark-for-cudatoolkit-root.sh

@@ -0,0 +1,64 @@
+# shellcheck shell=bash
+
+# shellcheck disable=SC1091
+source "@logFromSetupHook@"
+
+# Guard helper function
+# Returns 0 (success) if the hook should be run, 1 (failure) otherwise.
+# This allows us to use short-circuit evaluation to avoid running the hook when it shouldn't be.
+markForCUDAToolkit_ROOTGuard() {
+    log() {
+        logFromSetupHook \
+            "${1:?}" \
+            "mark-for-cudatoolkit-root" \
+            "markForCUDAToolkit_ROOTGuard" \
+            "${2:?}"
+    }
+    local guard="Skipping"
+    local reason=""
+
+    # This hook is meant only to add a stub file to the nix-support directory of the package including it in its
+    # nativeBuildInputs, so that the setup hook propagated by cuda_nvcc, setup-cuda, can detect it and add the
+    # package to the CUDA toolkit root. Therefore, since it only modifies the package being built and will not be
+    # propagated, it should only ever be included in nativeBuildInputs.
+    if (( ${hostOffset:?} == -1 && ${targetOffset:?} == 0)); then
+        guard="Sourcing"
+        reason="because the hook is in nativeBuildInputs relative to the package being built"
+    fi
+
+    log "INFO" "$guard $reason"
+
+    # Recall that test commands return 0 for success and 1 for failure.
+    [[ "$guard" == Sourcing ]]
+    return $?
+}
+
+# Guard against calling the hook at the wrong time.
+markForCUDAToolkit_ROOTGuard || return 0
+
+markForCUDAToolkit_ROOT() {
+    log() {
+        logFromSetupHook \
+            "${1:?}" \
+            "mark-for-cudatoolkit-root" \
+            "markForCUDAToolkit_ROOT" \
+            "${2:?}"
+    }
+    log "INFO" "Running on ${prefix:?}"
+
+    local markerPath="$prefix/nix-support/include-in-cudatoolkit-root"
+    mkdir -p "$(dirname "$markerPath")"
+    if [[ -f "$markerPath" ]]; then
+        log "DEBUG" "$markerPath exists, skipping"
+        return 0
+    fi
+
+    # Always create the file, even if it's empty, since setup-cuda relies on its existence.
+    # However, only populate it if strictDeps is not set.
+    touch "$markerPath"
+    if [[ -z "${strictDeps-}" ]]; then
+        log "DEBUG" "Populating $markerPath"
+        echo "${pname:?}-${output:?}" > "$markerPath"
+    fi
+}
+fixupOutputHooks+=(markForCUDAToolkit_ROOT)