cudaPackages: generalize and refactor setup hook

This PR refactor CUDA setup hooks, and in particular
autoAddOpenGLRunpath and autoAddCudaCompatRunpathHook, that were using a
lot of code in common (in fact, I introduced the latter by copy pasting
most of the bash script of the former). This is not satisfying for
maintenance, as a recent patch showed, because we need to duplicate
changes to both hooks.

This commit abstract the common part in a single shell script that
applies a generic patch action to every elf file in the output. For
autoAddOpenGLRunpath the action is just addOpenGLRunpath (now
addDriverRunpath), and is few line function for
autoAddCudaCompatRunpathHook.

Doing so, we also takes the occasion to use the newer addDriverRunpath
instead of the previous addOpenGLRunpath, and rename the CUDA hook to
reflect that as well.

Co-Authored-By: Connor Baker <connor.baker@tweag.io>

17 files changed, 139 insertions, 78 deletions

diff --git a/pkgs/development/cuda-modules/cuda-samples/generic.nix b/pkgs/development/cuda-modules/cuda-samples/generic.nix
index fb3d7cc99da9..3d1dac015e16 100644
--- a/pkgs/development/cuda-modules/cuda-samples/generic.nix
+++ b/pkgs/development/cuda-modules/cuda-samples/generic.nix
@@ -1,5 +1,5 @@
 {
-  autoAddOpenGLRunpathHook,
+  autoAddDriverRunpath,
   backendStdenv,
   cmake,
   cudatoolkit,
@@ -31,7 +31,7 @@ backendStdenv.mkDerivation (
 
     nativeBuildInputs =
       [
-        autoAddOpenGLRunpathHook
+        autoAddDriverRunpath
         pkg-config
       ]
       # CMake has to run as a native, build-time dependency for libNVVM samples.
diff --git a/pkgs/development/cuda-modules/cudatoolkit/default.nix b/pkgs/development/cuda-modules/cudatoolkit/default.nix
index aca0c7ad0b78..231a153bf7e6 100644
--- a/pkgs/development/cuda-modules/cudatoolkit/default.nix
+++ b/pkgs/development/cuda-modules/cudatoolkit/default.nix
@@ -2,7 +2,7 @@
   cudaVersion,
   runPatches ? [],
   autoPatchelfHook,
-  autoAddOpenGLRunpathHook,
+  autoAddDriverRunpath,
   addOpenGLRunpath,
   alsa-lib,
   curlMinimal,
@@ -76,7 +76,7 @@ backendStdenv.mkDerivation rec {
       rsync
       addOpenGLRunpath
       autoPatchelfHook
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
       markForCudatoolkitRootHook
     ]
     ++ lib.optionals (lib.versionOlder version "11") [libsForQt5.wrapQtAppsHook]
diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index 5e837fa36b5e..4f40b7f01dc2 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -1,7 +1,7 @@
 {
   # General callPackage-supplied arguments
-  autoAddOpenGLRunpathHook,
-  autoAddCudaCompatRunpathHook,
+  autoAddDriverRunpath,
+  autoAddCudaCompatRunpath,
   autoPatchelfHook,
   backendStdenv,
   fetchurl,
@@ -193,16 +193,16 @@ backendStdenv.mkDerivation (
       # in typically /lib/opengl-driver by adding that
       # directory to the rpath of all ELF binaries.
       # Check e.g. with `patchelf --print-rpath path/to/my/binary
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
       markForCudatoolkitRootHook
     ]
-    # autoAddCudaCompatRunpathHook depends on cuda_compat and would cause
+    # autoAddCudaCompatRunpath depends on cuda_compat and would cause
     # infinite recursion if applied to `cuda_compat` itself (beside the fact
     # that it doesn't make sense in the first place)
     ++ lib.optionals (pname != "cuda_compat" && flags.isJetsonBuild) [
-      # autoAddCudaCompatRunpathHook must appear AFTER autoAddOpenGLRunpathHook.
+      # autoAddCudaCompatRunpath must appear AFTER autoAddDriverRunpath.
       # See its documentation in ./setup-hooks/extension.nix.
-      autoAddCudaCompatRunpathHook
+      autoAddCudaCompatRunpath
     ];
 
     buildInputs =
diff --git a/pkgs/development/cuda-modules/nccl/default.nix b/pkgs/development/cuda-modules/nccl/default.nix
index 25296c21365d..b8bdc69bba4f 100644
--- a/pkgs/development/cuda-modules/nccl/default.nix
+++ b/pkgs/development/cuda-modules/nccl/default.nix
@@ -12,7 +12,7 @@
 }:
 let
   inherit (cudaPackages)
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     backendStdenv
     cuda_cccl
     cuda_cudart
@@ -44,7 +44,7 @@ backendStdenv.mkDerivation (
     nativeBuildInputs =
       [
         which
-        autoAddOpenGLRunpathHook
+        autoAddDriverRunpath
         python3
       ]
       ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index 68c60e2d8446..bc299dea006f 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -5,7 +5,7 @@
 }:
 let
   inherit (cudaPackages)
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     backendStdenv
     cuda_cccl
     cuda_cudart
@@ -29,7 +29,7 @@ backendStdenv.mkDerivation {
   nativeBuildInputs =
     [
       cmake
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
     ]
     ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
     ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [cuda_nvcc];
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
index ae25cebaca6f..fc41024f1551 100644
--- a/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
@@ -3,25 +3,25 @@
 # coming from the cuda_compat package by adding it to the RUNPATH.
 echo "Sourcing auto-add-cuda-compat-runpath-hook"
 
-elfHasDynamicSection() {
-    patchelf --print-rpath "$1" >& /dev/null
-}
+addCudaCompatRunpath() {
+  local libPath
+  local origRpath
+
+  if [[ $# -eq 0 ]]; then
+    echo "addCudaCompatRunpath: no library path provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "addCudaCompatRunpath: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "addCudaCompatRunpath: empty library path" >&2
+    exit 1
+  else
+    libPath="$1"
+  fi
 
-autoAddCudaCompatRunpathPhase() (
-  local outputPaths
-  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
-  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
-    if isELF "$f"; then
-      # patchelf returns an error on statically linked ELF files
-      if elfHasDynamicSection "$f" ; then
-        echo "autoAddCudaCompatRunpathHook: patching $f"
-        local origRpath="$(patchelf --print-rpath "$f")"
-        patchelf --set-rpath "@libcudaPath@:$origRpath" "$f"
-      elif (( "${NIX_DEBUG:-0}" >= 1 )) ; then
-        echo "autoAddCudaCompatRunpathHook: skipping a statically-linked ELF file $f"
-      fi
-    fi
-  done
-)
+  origRpath="$(patchelf --print-rpath "$libPath")"
+  patchelf --set-rpath "@libcudaPath@:$origRpath" "$libPath"
+}
 
-postFixupHooks+=(autoAddCudaCompatRunpathPhase)
+postFixupHooks+=("autoFixElfFiles addCudaCompatRunpath")
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh
new file mode 100644
index 000000000000..ecff2a032d64
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh
@@ -0,0 +1,8 @@
+# shellcheck shell=bash
+# Run addDriverRunpath on all dynamically linked ELF files
+echo "Sourcing auto-add-driver-runpath-hook"
+
+if [ -z "${dontUseAutoAddDriverRunpath-}" ]; then
+  echo "Using autoAddDriverRunpath"
+  postFixupHooks+=("autoFixElfFiles addDriverRunpath")
+fi
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
deleted file mode 100644
index a6eeef7c7699..000000000000
--- a/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
+++ /dev/null
@@ -1,28 +0,0 @@
-# shellcheck shell=bash
-# Run addOpenGLRunpath on all dynamically linked, ELF files
-echo "Sourcing auto-add-opengl-runpath-hook"
-
-elfHasDynamicSection() {
-    patchelf --print-rpath "$1" >& /dev/null
-}
-
-autoAddOpenGLRunpathPhase() (
-  local outputPaths
-  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
-  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
-    if isELF "$f"; then
-      # patchelf returns an error on statically linked ELF files
-      if elfHasDynamicSection "$f" ; then
-        echo "autoAddOpenGLRunpathHook: patching $f"
-        addOpenGLRunpath "$f"
-      elif (( "${NIX_DEBUG:-0}" >= 1 )) ; then
-        echo "autoAddOpenGLRunpathHook: skipping a statically-linked ELF file $f"
-      fi
-    fi
-  done
-)
-
-if [ -z "${dontUseAutoAddOpenGLRunpath-}" ]; then
-  echo "Using autoAddOpenGLRunpathPhase"
-  postFixupHooks+=(autoAddOpenGLRunpathPhase)
-fi
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh b/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh
new file mode 100644
index 000000000000..1d57dfb17a66
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh
@@ -0,0 +1,64 @@
+# shellcheck shell=bash
+# List all dynamically linked ELF files in the outputs and apply a generic fix
+# action provided as a parameter (currently used to add the CUDA or the
+# cuda_compat driver to the runpath of binaries)
+echo "Sourcing cuda/fix-elf-files.sh"
+
+# Returns the exit code of patchelf --print-rpath.
+# A return code of 0 (success) means the ELF file has a dynamic section, while
+# a non-zero return code means the ELF file is statically linked (or is not an
+# ELF file).
+elfHasDynamicSection() {
+  local libPath
+
+  if [[ $# -eq 0 ]]; then
+    echo "elfHasDynamicSection: no library path provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "elfHasDynamicSection: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "elfHasDynamicSection: empty library path" >&2
+    exit 1
+  else
+    libPath="$1"
+    shift 1
+  fi
+
+  patchelf --print-rpath "$libPath" >& /dev/null
+  return $?
+}
+
+# Run a fix action on all dynamically linked ELF files in the outputs.
+autoFixElfFiles() {
+  local fixAction
+  local outputPaths
+
+  if [[ $# -eq 0 ]]; then
+    echo "autoFixElfFiles: no fix action provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "autoFixElfFiles: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "autoFixElfFiles: empty fix action" >&2
+    exit 1
+  else
+    fixAction="$1"
+  fi
+
+  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
+
+  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
+    if ! isELF "$f"; then
+      continue
+    elif elfHasDynamicSection "$f"; then
+      # patchelf returns an error on statically linked ELF files, and in
+      # practice fixing actions all involve patchelf
+      echo "autoFixElfFiles: using $fixAction to fix $f" >&2
+      $fixAction "$f"
+    elif (( "${NIX_DEBUG:-0}" >= 1 )); then
+      echo "autoFixElfFiles: skipping a statically-linked ELF file $f"
+    fi
+  done
+}
diff --git a/pkgs/development/cuda-modules/setup-hooks/extension.nix b/pkgs/development/cuda-modules/setup-hooks/extension.nix
index 9e352bd5b3af..0e14ce709a64 100644
--- a/pkgs/development/cuda-modules/setup-hooks/extension.nix
+++ b/pkgs/development/cuda-modules/setup-hooks/extension.nix
@@ -1,4 +1,19 @@
 final: _: {
+  # Helper hook used in both autoAddCudaCompatRunpath and
+  # autoAddDriverRunpath that applies a generic patching action to all elf
+  # files with a dynamic linking section.
+  autoFixElfFiles =
+    final.callPackage
+      (
+        {makeSetupHook}:
+         makeSetupHook
+          {
+            name = "auto-fix-elf-files";
+          }
+          ./auto-fix-elf-files.sh
+      )
+      {};
+
   # Internal hook, used by cudatoolkit and cuda redist packages
   # to accommodate automatic CUDAToolkit_ROOT construction
   markForCudatoolkitRootHook =
@@ -32,31 +47,33 @@ final: _: {
       {}
     );
 
-  autoAddOpenGLRunpathHook =
+  autoAddDriverRunpath =
     final.callPackage
       (
-        {addOpenGLRunpath, makeSetupHook}:
+        {addDriverRunpath, autoFixElfFiles, makeSetupHook}:
         makeSetupHook
           {
             name = "auto-add-opengl-runpath-hook";
-            propagatedBuildInputs = [addOpenGLRunpath];
+            propagatedBuildInputs = [addDriverRunpath autoFixElfFiles];
           }
-          ./auto-add-opengl-runpath-hook.sh
+          ./auto-add-driver-runpath-hook.sh
       )
       {};
 
-  # autoAddCudaCompatRunpathHook hook must be added AFTER `setupCudaHook`. Both
+  # 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.
-  autoAddCudaCompatRunpathHook =
+  autoAddCudaCompatRunpath =
     final.callPackage
       (
-        {makeSetupHook, cuda_compat ? null }:
+        {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;
diff --git a/pkgs/development/libraries/ucx/default.nix b/pkgs/development/libraries/ucx/default.nix
index 3b923d8efdd2..8e0772479ccb 100644
--- a/pkgs/development/libraries/ucx/default.nix
+++ b/pkgs/development/libraries/ucx/default.nix
@@ -36,7 +36,7 @@ stdenv.mkDerivation rec {
   ]
   ++ lib.optionals enableCuda [
     cudaPackages.cuda_nvcc
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = [
diff --git a/pkgs/development/libraries/xgboost/default.nix b/pkgs/development/libraries/xgboost/default.nix
index b700dd2581c4..67190d7cac76 100644
--- a/pkgs/development/libraries/xgboost/default.nix
+++ b/pkgs/development/libraries/xgboost/default.nix
@@ -57,7 +57,7 @@ stdenv.mkDerivation rec {
 
   nativeBuildInputs = [ cmake ]
     ++ lib.optionals stdenv.isDarwin [ llvmPackages.openmp ]
-    ++ lib.optionals cudaSupport [ cudaPackages.autoAddOpenGLRunpathHook ]
+    ++ lib.optionals cudaSupport [ cudaPackages.autoAddDriverRunpath ]
     ++ lib.optionals rLibrary [ R ];
 
   buildInputs = [ gtest ] ++ lib.optional cudaSupport cudaPackages.cudatoolkit
diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix
index 199b352b2aa2..e2260aec4b14 100644
--- a/pkgs/development/python-modules/jaxlib/bin.nix
+++ b/pkgs/development/python-modules/jaxlib/bin.nix
@@ -23,7 +23,7 @@
 }:
 
 let
-  inherit (cudaPackagesGoogle) autoAddOpenGLRunpathHook cudaVersion;
+  inherit (cudaPackagesGoogle) autoAddDriverRunpath cudaVersion;
 
   version = "0.4.24";
 
@@ -180,7 +180,7 @@ buildPythonPackage {
   # Prebuilt wheels are dynamically linked against things that nix can't find.
   # Run `autoPatchelfHook` to automagically fix them.
   nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ]
-    ++ lib.optionals cudaSupport [ autoAddOpenGLRunpathHook ];
+    ++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
   # Dynamic link dependencies
   buildInputs = [ stdenv.cc.cc.lib ];
 
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
index 657a6e52c084..cbce801888e2 100644
--- a/pkgs/development/python-modules/jaxlib/default.nix
+++ b/pkgs/development/python-modules/jaxlib/default.nix
@@ -51,7 +51,7 @@
 }@inputs:
 
 let
-  inherit (cudaPackagesGoogle) autoAddOpenGLRunpathHook cudaFlags cudaVersion cudnn nccl;
+  inherit (cudaPackagesGoogle) autoAddDriverRunpath cudaFlags cudaVersion cudnn nccl;
 
   pname = "jaxlib";
   version = "0.4.24";
@@ -420,7 +420,7 @@ buildPythonPackage {
     done
   '';
 
-  nativeBuildInputs = lib.optionals cudaSupport [ autoAddOpenGLRunpathHook ];
+  nativeBuildInputs = lib.optionals cudaSupport [ autoAddDriverRunpath ];
 
   propagatedBuildInputs = [
     absl-py
diff --git a/pkgs/development/python-modules/tensorrt/default.nix b/pkgs/development/python-modules/tensorrt/default.nix
index e108b1a773cb..58a3e64d5610 100644
--- a/pkgs/development/python-modules/tensorrt/default.nix
+++ b/pkgs/development/python-modules/tensorrt/default.nix
@@ -22,7 +22,7 @@ buildPythonPackage rec {
   nativeBuildInputs = [
     unzip
     autoPatchelfHook
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   preUnpack = ''
diff --git a/pkgs/development/python-modules/torch/bin.nix b/pkgs/development/python-modules/torch/bin.nix
index 0bb415574e39..bee32b616345 100644
--- a/pkgs/development/python-modules/torch/bin.nix
+++ b/pkgs/development/python-modules/torch/bin.nix
@@ -40,7 +40,7 @@ in buildPythonPackage {
   nativeBuildInputs = lib.optionals stdenv.isLinux [
     addOpenGLRunpath
     autoPatchelfHook
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = lib.optionals stdenv.isLinux (with cudaPackages; [
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index d6c51904bd9d..10eecd1de99b 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -338,7 +338,7 @@ in buildPythonPackage rec {
     pythonRelaxDepsHook
     removeReferencesTo
   ] ++ lib.optionals cudaSupport (with cudaPackages; [
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     cuda_nvcc
   ])
   ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];