cudaPackages.buildRedist: include removeStubsFromRunpathHook

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

Author: ConnorBaker

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
@@ -140,6 +146,62 @@ checkCudaNonEmptyOutputs() {
   return 0
 }
 
+# Any redistributable providing stubs should set includeRemoveStubsFromRunpathHook to true -- since we don't track the
+# contents of the redistributables, it's only included by default if there is a stubs output.
+# This check additionally requires that any output which has a stubs directory includes a dependency on
+# includeRemoveStubsFromRunpathHook -- that way, if *any* of them are used, the hook is brought in as well.
+# Since includeRemoveStubsFromRunpathHook only adds the hook to whatever outputStubs resolves to, having stubs present
+# across multiple outputs will result in an error.
+checkCudaHasStubsIffIncludeRemoveStubsFromRunpathHook() {
+  local outputName
+  local -i hasStubs
+  local -i hasRemoveStubsFromRunpathHook
+  local -a outputNamesWronglyExcludingHook=()
+  local -a outputNamesWronglyIncludingHook=()
+
+  for outputName in $(getAllOutputNames); do
+    # Record the output if it contains a directory named "stubs" and doesn't include a dependency on
+    # removeStubsFromRunpathHook.
+    hasStubs=0
+    if find "${!outputName:?}" -mindepth 1 -type d -name stubs -print -quit | grep --silent .; then
+      hasStubs=1
+    fi
+
+    hasRemoveStubsFromRunpathHook=0
+    if
+      grep --silent --no-messages removeStubsFromRunpathHook "${!outputName:?}/nix-support/propagated-build-inputs"
+    then
+      hasRemoveStubsFromRunpathHook=1
+    fi
+
+    if ((hasStubs && !hasRemoveStubsFromRunpathHook)); then
+      # Outputs with stubs must include the hook.
+      outputNamesWronglyExcludingHook+=("${outputName:?}")
+    elif ((!hasStubs && hasRemoveStubsFromRunpathHook)); then
+      # Outputs without stubs cannot include the hook.
+      outputNamesWronglyIncludingHook+=("${outputName:?}")
+    fi
+  done
+
+  if ((${#outputNamesWronglyExcludingHook[@]})); then
+    nixErrorLog "includeRemoveStubsFromRunpathHook is false but we detected outputs containing a stubs" \
+      "directory: ${outputNamesWronglyExcludingHook[*]}"
+    nixErrorLog "ensure redistributables providing stubs set includeRemoveStubsFromRunpathHook to true"
+  fi
+
+  if ((${#outputNamesWronglyIncludingHook[@]})); then
+    nixErrorLog "includeRemoveStubsFromRunpathHook is true but we detected outputs without a stubs" \
+      "directory: ${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

@@ -6,15 +6,14 @@
   autoAddDriverRunpath,
   autoPatchelfHook,
   backendStdenv,
-  config,
   cudaMajorMinorVersion,
   cudaMajorVersion,
   cudaNamePrefix,
   fetchurl,
   lib,
   manifests,
   markForCudatoolkitRootHook,
-  setupCudaHook,
+  removeStubsFromRunpathHook,
   srcOnly,
   stdenv,
   stdenvNoCC,
@@ -24,6 +23,7 @@ let
   inherit (_cuda.lib) getNixSystems _mkCudaVariant mkRedistUrl;
   inherit (lib.attrsets)
     foldlAttrs
+    getDev
     hasAttr
     isAttrs
     attrNames
@@ -55,6 +55,7 @@ let
     ;
   inherit (lib.strings)
     concatMapStringsSep
+    optionalString
     toUpper
     stringLength
     substring
@@ -138,6 +139,8 @@ extendMkDerivation {
 
       # Fixups
       appendRunpaths ? [ ],
+      includeRemoveStubsFromRunpathHook ? elem "stubs" finalAttrs.outputs,
+      postFixup ? "",
 
       # Extra
       passthru ? { },
@@ -201,7 +204,10 @@ extendMkDerivation {
         outputPython = [ "python" ];
         outputSamples = [ "samples" ];
         outputStatic = [ "static" ];
-        outputStubs = [ "stubs" ];
+        outputStubs = [
+          "stubs"
+          "lib"
+        ];
       },
       ...
     }:
@@ -266,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
       ]
@@ -332,6 +341,20 @@ extendMkDerivation {
 
       inherit doInstallCheck;
       inherit allowFHSReferences;
+      inherit includeRemoveStubsFromRunpathHook;
+
+      postFixup =
+        postFixup
+        # Register removeStubsFromRunpathHook.
+        # This must happen in postFixup after recordPropagatedDependencies (fixupPhase).
+        + ''
+          if [[ -n "''${includeRemoveStubsFromRunpathHook:-}" ]] ; then
+            nixLog "installing stub removal runpath hook"
+            mkdir -p "''${!outputStubs:?}/nix-support"
+            printWords >>"''${!outputStubs:?}/nix-support/propagated-build-inputs" \
+              "${getDev removeStubsFromRunpathHook}"
+          fi
+        '';
 
       passthru = passthru // {
         inherit redistName release;

pkgs/development/cuda-modules/default.nix

@@ -132,7 +132,6 @@ let
       buildRedist = import ./buildRedist {
         inherit
           _cuda
-          config
           lib
           ;
         inherit (pkgs)
@@ -151,7 +150,7 @@ let
           cudaNamePrefix
           manifests
           markForCudatoolkitRootHook
-          setupCudaHook
+          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";