Feat/cuda stub setup hook

pkgs/build-support/setup-hooks/arrayUtilities/getRunpathEntries/getRunpathEntries.bash

@@ -0,0 +1,49 @@
+# shellcheck shell=bash
+
+# getRunpathEntries
+# Append the runpath entries of the dynamically linked ELF file at path to the indexed array referenced by
+# outputArrRef.
+#
+# NOTE: This function does not check if path is a valid ELF file.
+#
+# Arguments:
+# - path: a path to an ELF file with a dynamic section
+# - outputArrRef: a reference to an indexed array (mutated only by appending)
+#
+# Returns 0 if the file is dynamically linked and the runpath was appended to the output array.
+# Returns 1 if patchelf fails (e.g., the ELF file is not dynamically linked, so patchelf fails to print the rpath).
+getRunpathEntries() {
+  if (($# != 2)); then
+    nixErrorLog "expected two arguments!"
+    nixErrorLog "usage: getRunpathEntries path outputArrRef"
+    exit 1
+  fi
+
+  local -r path="$1"
+  local -rn outputArrRef="$2"
+
+  if [[ ! -f $path ]]; then
+    nixErrorLog "first argument path $path is not a file"
+    exit 1
+  elif ! isDeclaredArray "${!outputArrRef}"; then
+    nixErrorLog "second argument outputArrRef must be a reference to an indexed array"
+    exit 1
+  fi
+
+  # Declare runpath separately to avoid masking the return value of patchelf.
+  local runpath
+  # Files that are not dynamically linked cause patchelf to exit with a non-zero status and print to stderr.
+  # If patchelf fails to print the rpath, we assume the file is not dynamically linked.
+  runpath="$(patchelf --print-rpath "$path" 2>/dev/null)" || return 1
+
+  # If the runpath is empty and we feed it to mapfile, it gives us a singleton array with an empty string.
+  # We want to avoid that, so we check if the runpath is empty before trying to populate runpathEntries.
+  local -a runpathEntries=()
+  if [[ -n $runpath ]]; then
+    mapfile -d ':' -t runpathEntries < <(echo -n "$runpath")
+  fi
+
+  outputArrRef+=("${runpathEntries[@]}")
+
+  return 0
+}

pkgs/build-support/setup-hooks/arrayUtilities/getRunpathEntries/package.nix

@@ -0,0 +1,15 @@
+{
+  callPackages,
+  isDeclaredArray,
+  makeSetupHook,
+  patchelf,
+}:
+makeSetupHook {
+  name = "getRunpathEntries";
+  propagatedBuildInputs = [
+    isDeclaredArray
+    patchelf
+  ];
+  passthru.tests = callPackages ./tests.nix { };
+  meta.description = "Appends runpath entries of a file to an array";
+} ./getRunpathEntries.bash

pkgs/build-support/setup-hooks/arrayUtilities/getRunpathEntries/tests.nix

@@ -0,0 +1,101 @@
+# NOTE: Tests related to getRunpathEntries go here.
+{
+  emptyFile,
+  getRunpathEntries,
+  hello,
+  lib,
+  pkgsStatic,
+  stdenv,
+  testers,
+}:
+let
+  inherit (lib.attrsets) recurseIntoAttrs;
+  inherit (testers)
+    shellcheck
+    shfmt
+    testBuildFailure'
+    testEqualArrayOrMap
+    ;
+
+  check =
+    {
+      name,
+      elfFile,
+      runpathEntries,
+    }:
+    (testEqualArrayOrMap {
+      inherit name;
+      expectedArray = runpathEntries;
+      script = ''
+        set -eu
+        nixLog "running getRunpathEntries with ''${elfFile@Q} to populate actualArray"
+        getRunpathEntries "$elfFile" actualArray || {
+          nixErrorLog "getRunpathEntries failed"
+          exit 1
+        }
+      '';
+    }).overrideAttrs
+      (prevAttrs: {
+        inherit elfFile;
+        nativeBuildInputs = prevAttrs.nativeBuildInputs or [ ] ++ [ getRunpathEntries ];
+        meta = prevAttrs.meta or { } // {
+          platforms = lib.platforms.linux;
+        };
+      });
+in
+recurseIntoAttrs {
+  shellcheck = shellcheck {
+    name = "getRunpathEntries";
+    src = ./getRunpathEntries.bash;
+  };
+
+  shfmt = shfmt {
+    name = "getRunpathEntries";
+    src = ./getRunpathEntries.bash;
+  };
+}
+# Only tested on Linux.
+// lib.optionalAttrs stdenv.hostPlatform.isLinux {
+  # Not an ELF file
+  notElfFileFails = testBuildFailure' {
+    name = "notElfFileFails";
+    drv = check {
+      name = "notElfFile";
+      elfFile = emptyFile;
+      runpathEntries = [ ];
+    };
+    expectedBuilderLogEntries = [
+      "getRunpathEntries failed"
+    ];
+  };
+
+  # Not a dynamic ELF file
+  staticElfFileFails = testBuildFailure' {
+    name = "staticElfFileFails";
+    drv = check {
+      name = "staticElfFile";
+      elfFile = lib.getExe pkgsStatic.hello;
+      runpathEntries = [ ];
+    };
+    expectedBuilderLogEntries = [
+      "getRunpathEntries failed"
+    ];
+  };
+
+  hello = check {
+    name = "hello";
+    elfFile = lib.getExe hello;
+    runpathEntries = [
+      "${lib.getLib stdenv.cc.libc}/lib"
+    ];
+  };
+
+  libstdcplusplus = check {
+    name = "libstdcplusplus";
+    elfFile = "${lib.getLib stdenv.cc.cc}/lib/libstdc++.so";
+    runpathEntries = [
+      "${lib.getLib stdenv.cc.cc}/lib"
+      "${lib.getLib stdenv.cc.libc}/lib"
+    ];
+  };
+}

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
@@ -108,7 +114,7 @@ checkCudaFhsRefs() {
   local -a outputPaths=()
   local firstMatches
 
-  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
+  mapfile -t outputPaths < <(for outputName in $(getAllOutputNames); do echo "${!outputName:?}"; done)
   firstMatches="$(grep --max-count=5 --recursive --exclude=LICENSE /usr/ "${outputPaths[@]}")" || true
   if [[ -n $firstMatches ]]; then
     nixErrorLog "detected references to /usr: $firstMatches"
@@ -119,27 +125,83 @@ checkCudaFhsRefs() {
 }
 
 checkCudaNonEmptyOutputs() {
-  local output
+  local outputName
   local dirs
-  local -a failingOutputs=()
+  local -a failingOutputNames=()
 
-  for output in $(getAllOutputNames); do
-    [[ ${!output:?} == "out" || ${!output:?} == "${!outputDev:?}" ]] && continue
-    dirs="$(find "${!output:?}" -mindepth 1 -maxdepth 1)" || true
-    if [[ -z $dirs || $dirs == "${!output:?}/nix-support" ]]; then
-      failingOutputs+=("$output")
+  for outputName in $(getAllOutputNames); do
+    [[ ${outputName:?} == "out" || ${outputName:?} == "${outputDev:?}" ]] && continue
+    dirs="$(find "${!outputName:?}" -mindepth 1 -maxdepth 1)" || true
+    if [[ -z $dirs || $dirs == "${!outputName:?}/nix-support" ]]; then
+      failingOutputNames+=("${outputName:?}")
     fi
   done
 
-  if ((${#failingOutputs[@]})); then
-    nixErrorLog "detected empty (excluding nix-support) outputs: ${failingOutputs[*]}"
+  if ((${#failingOutputNames[@]})); then
+    nixErrorLog "detected empty (excluding nix-support) outputs: ${failingOutputNames[*]}"
     nixErrorLog "this typically indicates a failure in packaging or moveToOutput ordering"
     exit 1
   fi
 
   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" ]