cudaPackages.buildRedist: include removeStubsFromRunpathHook

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>

pkgs/development/cuda-modules/buildRedist/buildRedistHook.bash

@@ -29,6 +29,12 @@ buildRedistHookRegistration() {
 
   postFixupHooks+=(fixupCudaPropagatedBuildOutputsToOut)
   nixLog "added fixupCudaPropagatedBuildOutputsToOut to postFixupHooks"
+
+  # NOTE: We need to do this in postFixup since we don't write the dependency on removeStubsFromRunpathHook until
+  # postFixup -- recall recordPropagatedDependencies happens during fixupPhase.
+  # NOTE: Iff is shorthand for "if and only if" -- the logical biconditional.
+  postFixupHooks+=(checkCudaHasStubsIffIncludeRemoveStubsFromRunpathHook)
+  nixLog "added checkCudaHasStubsIffIncludeRemoveStubsFromRunpathHook to postFixupHooks"
 }
 
 buildRedistHookRegistration
@@ -142,6 +148,55 @@ checkCudaNonEmptyOutputs() {
   return 0
 }
 
+checkCudaHasStubsIffIncludeRemoveStubsFromRunpathHook() {
+  local outputName
+  local -i hasStubs
+  local -i hasRemoveStubsFromRunpathHook
+  local -a outputNamesWronglyExcludingHook=()
+  local -a outputNamesWronglyIncludingHook=()
+
+  for outputName in $(getAllOutputNames); do
+    hasStubs=0
+    if find "${!outputName:?}" -mindepth 1 -type d -name stubs -print -quit | grep --silent .; then
+      hasStubs=1
+    fi
+
+    # The dependency should be recorded in both propagated-native-build-inputs and propagated-build-inputs, so the
+    # hook is propagated regardless of which dependency array includes the stubs-providing output.
+    hasRemoveStubsFromRunpathHook=0
+    if
+      grep --silent --no-messages removeStubsFromRunpathHook "${!outputName:?}/nix-support/propagated-native-build-inputs" &&
+        grep --silent --no-messages removeStubsFromRunpathHook "${!outputName:?}/nix-support/propagated-build-inputs"
+    then
+      hasRemoveStubsFromRunpathHook=1
+    fi
+
+    if ((hasStubs && !hasRemoveStubsFromRunpathHook)); then
+      outputNamesWronglyExcludingHook+=("${outputName:?}")
+    elif ((!hasStubs && hasRemoveStubsFromRunpathHook)); then
+      outputNamesWronglyIncludingHook+=("${outputName:?}")
+    fi
+  done
+
+  if ((${#outputNamesWronglyExcludingHook[@]})); then
+    nixErrorLog "we detected outputs containing a stubs directory without a dependency on" \
+      "removeStubsFromRunpathHook: ${outputNamesWronglyExcludingHook[*]}"
+    nixErrorLog "ensure redistributables providing stubs set includeRemoveStubsFromRunpathHook to true"
+  fi
+
+  if ((${#outputNamesWronglyIncludingHook[@]})); then
+    nixErrorLog "we detected outputs without a stubs directory with a dependency on" \
+      "removeStubsFromRunpathHook: ${outputNamesWronglyIncludingHook[*]}"
+    nixErrorLog "ensure redistributables without stubs do not set includeRemoveStubsFromRunpathHook to true"
+  fi
+
+  if ((${#outputNamesWronglyExcludingHook[@]} || ${#outputNamesWronglyIncludingHook[@]})); then
+    exit 1
+  fi
+
+  return 0
+}
+
 # TODO(@connorbaker): https://github.com/NixOS/nixpkgs/issues/323126.
 # _multioutPropagateDev() currently expects a space-separated string rather than an array.
 # NOTE: Because _multioutPropagateDev is a postFixup hook, we correct it in preFixup.

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

@@ -13,6 +13,7 @@
   lib,
   manifests,
   markForCudatoolkitRootHook,
+  removeStubsFromRunpathHook,
   srcOnly,
   stdenv,
   stdenvNoCC,
@@ -22,6 +23,7 @@ let
   inherit (_cuda.lib) getNixSystems _mkCudaVariant mkRedistUrl;
   inherit (lib.attrsets)
     foldlAttrs
+    getDev
     hasAttr
     isAttrs
     attrNames
@@ -53,6 +55,7 @@ let
     ;
   inherit (lib.strings)
     concatMapStringsSep
+    optionalString
     toUpper
     stringLength
     substring
@@ -136,6 +139,8 @@ extendMkDerivation {
 
       # Fixups
       appendRunpaths ? [ ],
+      includeRemoveStubsFromRunpathHook ? elem "stubs" finalAttrs.outputs,
+      postFixup ? "",
 
       # Extra
       passthru ? { },
@@ -199,7 +204,10 @@ extendMkDerivation {
         outputPython = [ "python" ];
         outputSamples = [ "samples" ];
         outputStatic = [ "static" ];
-        outputStubs = [ "stubs" ];
+        outputStubs = [
+          "stubs"
+          "lib"
+        ];
       },
       ...
     }:
@@ -264,6 +272,9 @@ extendMkDerivation {
         # 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
+        # TODO(@connorbaker): Given we'll have stubs available, we can switch from autoPatchelfIgnoreMissingDeps to
+        # allowing autoPatchelf to find and link against the stub files and rely on removeStubsFromRunpathHook to
+        # automatically find and replace those references with ones to the driver link lib directory.
         autoAddDriverRunpath
         markForCudatoolkitRootHook
       ]
@@ -330,6 +341,18 @@ extendMkDerivation {
 
       inherit doInstallCheck;
       inherit allowFHSReferences;
+      inherit includeRemoveStubsFromRunpathHook;
+
+      postFixup =
+        postFixup
+        + optionalString finalAttrs.includeRemoveStubsFromRunpathHook ''
+          nixLog "installing stub removal runpath hook"
+          mkdir -p "''${!outputStubs:?}/nix-support"
+          printWords >>"''${!outputStubs:?}/nix-support/propagated-native-build-inputs" \
+            "${getDev removeStubsFromRunpathHook.__spliced.buildHost or removeStubsFromRunpathHook}"
+          printWords >>"''${!outputStubs:?}/nix-support/propagated-build-inputs" \
+            "${getDev removeStubsFromRunpathHook.__spliced.hostTarget or removeStubsFromRunpathHook}"
+        '';
 
       passthru = passthru // {
         inherit redistName release;

pkgs/development/cuda-modules/default.nix

@@ -150,6 +150,7 @@ let
           cudaNamePrefix
           manifests
           markForCudatoolkitRootHook
+          removeStubsFromRunpathHook
           ;
       };
 

pkgs/development/cuda-modules/packages/cuda_cudart.nix

@@ -22,6 +22,9 @@ buildRedist (finalAttrs: {
     "out"
   ];
 
+  # We have stubs but we don't have an explicit stubs output.
+  includeRemoveStubsFromRunpathHook = true;
+
   propagatedBuildOutputs =
     # required by CMake
     lib.optionals (lib.elem "static" finalAttrs.outputs) [ "static" ]

pkgs/development/cuda-modules/packages/libnvfatbin.nix

@@ -5,6 +5,9 @@ buildRedist {
 
   outputs = [ "out" ];
 
+  # Includes stubs.
+  includeRemoveStubsFromRunpathHook = true;
+
   meta = {
     description = "APIs which can be used at runtime to combine multiple CUDA objects into one CUDA fat binary (fatbin)";
     homepage = "https://docs.nvidia.com/cuda/nvfatbin";