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)