cudaPackages: move derivations to cuda-modules & support aarch64

cudaPackages.cuda_compat: ignore missing libs provided at runtime

cudaPackages.gpus: Jetson should never build by default

cudaPackages.flags: don't build Jetson capabilities by default

cudaPackages: re-introduce filter for pre-existing CUDA redist packages in overrides

cudaPackages: only recurseIntoAttrs for the latest of each major version

cudaPackages.nvccCompatabilities: use GCC 10 through CUDA 11.5 to avoid a GLIBC incompatability

cudaPackages.cutensor: acquire libcublas through cudatoolkit prior to 11.4

cudaPackages.cuda_compat: mark as broken on aarch64-linux if not targeting Jetson

cudaPackages.cutensor_1_4: fix build

cudaPackages: adjust use of autoPatchelfIgnoreMissingDeps

cudaPackages.cuda_nvprof: remove unecessary override to add addOpenGLRunpath

cudaPackages: use getExe' to avoid patchelf warning about missing meta.mainProgram

cudaPackages: fix evaluation with Nix 2.3

cudaPackages: fix platform detection for Jetson/non-Jetson aarch64-linux

python3Packages.tensorrt: mark as broken if required packages are missing

Note: evaluating the name of the derivation will fail if tensorrt is not present,
which is why we wrap the value in `lib.optionalString`.

cudaPackages.flags.getNixSystem: add guard based on jetsonTargets

cudaPackages.cudnn: use explicit path to patchelf

cudaPackages.tensorrt: use explicit path to patchelf

4 files changed, 228 insertions, 0 deletions

diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
new file mode 100644
index 0000000000000..f50a5f6c25c66
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
@@ -0,0 +1,28 @@
+# shellcheck shell=bash
+# Run addOpenGLRunpath on all dynamically linked, ELF files
+echo "Sourcing auto-add-opengl-runpath-hook"
+
+elfHasDynamicSection() {
+    patchelf --print-rpath "$1" >& /dev/null
+}
+
+autoAddOpenGLRunpathPhase() (
+  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 "autoAddOpenGLRunpathHook: patching $f"
+        addOpenGLRunpath "$f"
+      elif (( "${NIX_DEBUG:-0}" >= 1 )) ; then
+        echo "autoAddOpenGLRunpathHook: skipping a statically-linked ELF file $f"
+      fi
+    fi
+  done
+)
+
+if [ -z "${dontUseAutoAddOpenGLRunpath-}" ]; then
+  echo "Using autoAddOpenGLRunpathPhase"
+  postFixupHooks+=(autoAddOpenGLRunpathPhase)
+fi
diff --git a/pkgs/development/cuda-modules/setup-hooks/extension.nix b/pkgs/development/cuda-modules/setup-hooks/extension.nix
new file mode 100644
index 0000000000000..762dad9ea8765
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/extension.nix
@@ -0,0 +1,47 @@
+final: _: {
+  # Internal hook, used by cudatoolkit and cuda redist packages
+  # to accommodate automatic CUDAToolkit_ROOT construction
+  markForCudatoolkitRootHook =
+    final.callPackage
+      (
+        {makeSetupHook}:
+        makeSetupHook {name = "mark-for-cudatoolkit-root-hook";} ./mark-for-cudatoolkit-root-hook.sh
+      )
+      {};
+
+  # Currently propagated by cuda_nvcc or cudatoolkit, rather than used directly
+  setupCudaHook =
+    (final.callPackage
+      (
+        {makeSetupHook, backendStdenv}:
+        makeSetupHook
+          {
+            name = "setup-cuda-hook";
+
+            substitutions.setupCudaHook = placeholder "out";
+
+            # Point NVCC at a compatible compiler
+            substitutions.ccRoot = "${backendStdenv.cc}";
+
+            # Required in addition to ccRoot as otherwise bin/gcc is looked up
+            # when building CMakeCUDACompilerId.cu
+            substitutions.ccFullPath = "${backendStdenv.cc}/bin/${backendStdenv.cc.targetPrefix}c++";
+          }
+          ./setup-cuda-hook.sh
+      )
+      {}
+    );
+
+  autoAddOpenGLRunpathHook =
+    final.callPackage
+      (
+        {addOpenGLRunpath, makeSetupHook}:
+        makeSetupHook
+          {
+            name = "auto-add-opengl-runpath-hook";
+            propagatedBuildInputs = [addOpenGLRunpath];
+          }
+          ./auto-add-opengl-runpath-hook.sh
+      )
+      {};
+}
diff --git a/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
new file mode 100644
index 0000000000000..ba04c2e0806af
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
@@ -0,0 +1,14 @@
+# shellcheck shell=bash
+
+# Should we mimick cc-wrapper's "hygiene"?
+[[ -z ${strictDeps-} ]] || (( "$hostOffset" < 0 )) || return 0
+
+echo "Sourcing mark-for-cudatoolkit-root-hook" >&2
+
+markForCUDAToolkit_ROOT() {
+    mkdir -p "${prefix}/nix-support"
+    [[ -f "${prefix}/nix-support/include-in-cudatoolkit-root" ]] && return
+    echo "$pname-$output" > "${prefix}/nix-support/include-in-cudatoolkit-root"
+}
+
+fixupOutputHooks+=(markForCUDAToolkit_ROOT)
diff --git a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
new file mode 100644
index 0000000000000..7b7b3bdde80e3
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
@@ -0,0 +1,139 @@
+# shellcheck shell=bash
+
+# Only run the hook from nativeBuildInputs
+(( "$hostOffset" == -1 && "$targetOffset" == 0)) || return 0
+
+guard=Sourcing
+reason=
+
+[[ -n ${cudaSetupHookOnce-} ]] && guard=Skipping && reason=" because the hook has been propagated more than once"
+
+if (( "${NIX_DEBUG:-0}" >= 1 )) ; then
+    echo "$guard hostOffset=$hostOffset targetOffset=$targetOffset setupCudaHook$reason" >&2
+else
+    echo "$guard setup-cuda-hook$reason" >&2
+fi
+
+[[ "$guard" = Sourcing ]] || return 0
+
+declare -g cudaSetupHookOnce=1
+declare -Ag cudaHostPathsSeen=()
+declare -Ag cudaOutputToPath=()
+
+extendcudaHostPathsSeen() {
+    (( "${NIX_DEBUG:-0}" >= 1 )) && echo "extendcudaHostPathsSeen $1" >&2
+
+    local markerPath="$1/nix-support/include-in-cudatoolkit-root"
+    [[ ! -f "${markerPath}" ]] && return
+    [[ -v cudaHostPathsSeen[$1] ]] && return
+
+    cudaHostPathsSeen["$1"]=1
+
+    # E.g. cuda_cudart-lib
+    local cudaOutputName
+    read -r cudaOutputName < "$markerPath"
+
+    [[ -z "$cudaOutputName" ]] && return
+
+    local oldPath="${cudaOutputToPath[$cudaOutputName]-}"
+    [[ -n "$oldPath" ]] && echo "extendcudaHostPathsSeen: warning: overwriting $cudaOutputName from $oldPath to $1" >&2
+    cudaOutputToPath["$cudaOutputName"]="$1"
+}
+addEnvHooks "$targetOffset" extendcudaHostPathsSeen
+
+setupCUDAToolkit_ROOT() {
+    (( "${NIX_DEBUG:-0}" >= 1 )) && echo "setupCUDAToolkit_ROOT: cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2
+
+    for path in "${!cudaHostPathsSeen[@]}" ; do
+        addToSearchPathWithCustomDelimiter ";" CUDAToolkit_ROOT "$path"
+        if [[ -d "$path/include" ]] ; then
+            addToSearchPathWithCustomDelimiter ";" CUDAToolkit_INCLUDE_DIR "$path/include"
+        fi
+    done
+
+    export cmakeFlags+=" -DCUDAToolkit_INCLUDE_DIR=$CUDAToolkit_INCLUDE_DIR -DCUDAToolkit_ROOT=$CUDAToolkit_ROOT"
+}
+preConfigureHooks+=(setupCUDAToolkit_ROOT)
+
+setupCUDAToolkitCompilers() {
+    echo Executing setupCUDAToolkitCompilers >&2
+
+    if [[ -n "${dontSetupCUDAToolkitCompilers-}" ]] ; then
+        return
+    fi
+
+    # Point NVCC at a compatible compiler
+
+    # For CMake-based projects:
+    # https://cmake.org/cmake/help/latest/module/FindCUDA.html#input-variables
+    # https://cmake.org/cmake/help/latest/envvar/CUDAHOSTCXX.html
+    # https://cmake.org/cmake/help/latest/variable/CMAKE_CUDA_HOST_COMPILER.html
+
+    export cmakeFlags+=" -DCUDA_HOST_COMPILER=@ccFullPath@"
+    export cmakeFlags+=" -DCMAKE_CUDA_HOST_COMPILER=@ccFullPath@"
+
+    # For non-CMake projects:
+    # We prepend --compiler-bindir to nvcc flags.
+    # Downstream packages can override these, because NVCC
+    # uses the last --compiler-bindir it gets on the command line.
+    # FIXME: this results in "incompatible redefinition" warnings.
+    # https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html#compiler-bindir-directory-ccbin
+    if [ -z "${CUDAHOSTCXX-}" ]; then
+      export CUDAHOSTCXX="@ccFullPath@";
+    fi
+
+    export NVCC_PREPEND_FLAGS+=" --compiler-bindir=@ccRoot@/bin"
+
+    # NOTE: We set -Xfatbin=-compress-all, which reduces the size of the compiled
+    #   binaries. If binaries grow over 2GB, they will fail to link. This is a problem for us, as
+    #   the default set of CUDA capabilities we build can regularly cause this to occur (for
+    #   example, with Magma).
+    #
+    # @SomeoneSerge: original comment was made by @ConnorBaker in .../cudatoolkit/common.nix
+    if [[ -z "${dontCompressFatbin-}" ]]; then
+        export NVCC_PREPEND_FLAGS+=" -Xfatbin=-compress-all"
+    fi
+
+    # CMake's enable_language(CUDA) runs a compiler test and it doesn't account for
+    # CUDAToolkit_ROOT. We have to help it locate libcudart
+    if [[ -z "${nvccDontPrependCudartFlags-}" ]] ; then
+        if [[ ! -v cudaOutputToPath["cuda_cudart-out"] ]] ; then
+            echo "setupCUDAToolkitCompilers: missing cudaPackages.cuda_cudart. This may become an an error in the future" >&2
+            # exit 1
+        fi
+        for pkg in "${!cudaOutputToPath[@]}" ; do
+            [[ ! "$pkg" = cuda_cudart* ]] && continue
+
+            local path="${cudaOutputToPath[$pkg]}"
+            if [[ -d "$path/include" ]] ; then
+                export NVCC_PREPEND_FLAGS+=" -I$path/include"
+            fi
+            if [[ -d "$path/lib" ]] ; then
+                export NVCC_PREPEND_FLAGS+=" -L$path/lib"
+            fi
+        done
+    fi
+}
+preConfigureHooks+=(setupCUDAToolkitCompilers)
+
+propagateCudaLibraries() {
+    (( "${NIX_DEBUG:-0}" >= 1 )) && echo "propagateCudaLibraries: cudaPropagateToOutput=$cudaPropagateToOutput cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2
+
+    [[ -z "${cudaPropagateToOutput-}" ]] && return
+
+    mkdir -p "${!cudaPropagateToOutput}/nix-support"
+    # One'd expect this should be propagated-bulid-build-deps, but that doesn't seem to work
+    echo "@setupCudaHook@" >> "${!cudaPropagateToOutput}/nix-support/propagated-native-build-inputs"
+
+    local propagatedBuildInputs=( "${!cudaHostPathsSeen[@]}" )
+    for output in $(getAllOutputNames) ; do
+        if [[ ! "$output" = "$cudaPropagateToOutput" ]] ; then
+            propagatedBuildInputs+=( "${!output}" )
+        fi
+        break
+    done
+
+    # One'd expect this should be propagated-host-host-deps, but that doesn't seem to work
+    printWords "${propagatedBuildInputs[@]}" >> "${!cudaPropagateToOutput}/nix-support/propagated-build-inputs"
+}
+postFixupHooks+=(propagateCudaLibraries)