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>

2 files changed, 4 insertions, 4 deletions

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