Merge pull request #267247 from yannham/feat/cuda-compat-jetson

Use cuda_compat drivers when available

3 files changed, 56 insertions, 0 deletions

diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index 7ddecb7ed723..c7ff0f53bb74 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -1,6 +1,7 @@
 {
   # General callPackage-supplied arguments
   autoAddOpenGLRunpathHook,
+  autoAddCudaCompatRunpathHook,
   autoPatchelfHook,
   backendStdenv,
   fetchurl,
@@ -126,6 +127,14 @@ backendStdenv.mkDerivation (
       # Check e.g. with `patchelf --print-rpath path/to/my/binary
       autoAddOpenGLRunpathHook
       markForCudatoolkitRootHook
+    ]
+    # autoAddCudaCompatRunpathHook 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.
+      # See its documentation in ./setup-hooks/extension.nix.
+      autoAddCudaCompatRunpathHook
     ];
 
     buildInputs =
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
new file mode 100644
index 000000000000..537daad2f00e
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
@@ -0,0 +1,27 @@
+# shellcheck shell=bash
+# Patch all dynamically linked, ELF files with the CUDA driver (libcuda.so)
+# 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
+}
+
+autoAddCudaCompatRunpathPhase() (
+  local outputPaths
+  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
+  find "${outputPaths[@]}" -type f -executable -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
+)
+
+postFixupHooks+=(autoAddCudaCompatRunpathPhase)
diff --git a/pkgs/development/cuda-modules/setup-hooks/extension.nix b/pkgs/development/cuda-modules/setup-hooks/extension.nix
index 762dad9ea876..964c383a735d 100644
--- a/pkgs/development/cuda-modules/setup-hooks/extension.nix
+++ b/pkgs/development/cuda-modules/setup-hooks/extension.nix
@@ -44,4 +44,24 @@ final: _: {
           ./auto-add-opengl-runpath-hook.sh
       )
       {};
+
+  # autoAddCudaCompatRunpathHook 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 =
+    final.callPackage
+      (
+        {makeSetupHook, cuda_compat}:
+        makeSetupHook
+          {
+            name = "auto-add-cuda-compat-runpath-hook";
+            substitutions = {
+              libcudaPath = "${cuda_compat}/compat";
+            };
+          }
+          ./auto-add-cuda-compat-runpath.sh
+      )
+      {};
 }