From d0efe91c4acbaef66548b6f9e5c12fe35d96451e Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Wed, 3 Apr 2024 22:09:29 +0000 Subject: [PATCH] cudaPackages.setupCuda: rewrite and rename --- pkgs/development/cuda-modules/aliases.nix | 2 +- .../setup-hooks/setup-cuda/default.nix | 1 + .../setup-hooks/setup-cuda/setup-cuda.sh | 160 +++++++++++++----- 3 files changed, 117 insertions(+), 46 deletions(-) diff --git a/pkgs/development/cuda-modules/aliases.nix b/pkgs/development/cuda-modules/aliases.nix index 23c400b75bbf51d..8a8c7e8bb6ae56e 100644 --- a/pkgs/development/cuda-modules/aliases.nix +++ b/pkgs/development/cuda-modules/aliases.nix @@ -19,6 +19,6 @@ in mkRenamed "markForCudatoolkitRootHook" "cudaPackages.markForCudatoolkitRoot" final.markForCudatoolkitRoot; # Added 2024-04-19 setupCudaHook = - mkRenamed "setupCudaHook" "setupCuda" + mkRenamed "setupCudaHook" "cudaPackages.setupCuda" final.setupCuda; # Added 2024-04-19 } diff --git a/pkgs/development/cuda-modules/setup-hooks/setup-cuda/default.nix b/pkgs/development/cuda-modules/setup-hooks/setup-cuda/default.nix index 528439a62725357..259595ca80fdaa6 100644 --- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda/default.nix +++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda/default.nix @@ -11,6 +11,7 @@ makeSetupHook { ccFullPath = "${cc}/bin/${cc.targetPrefix}c++"; # Point NVCC at a compatible compiler ccRoot = "${cc}"; + logFromSetupHook = ../utilities/log-from-setup-hook.sh; setupCuda = placeholder "out"; }; } ./setup-cuda.sh diff --git a/pkgs/development/cuda-modules/setup-hooks/setup-cuda/setup-cuda.sh b/pkgs/development/cuda-modules/setup-hooks/setup-cuda/setup-cuda.sh index 2c7aa7d4c11d573..0bfa7791e6b553c 100644 --- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda/setup-cuda.sh +++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda/setup-cuda.sh @@ -1,33 +1,74 @@ # shellcheck shell=bash -# Only run the hook from nativeBuildInputs -(( "$hostOffset" == -1 && "$targetOffset" == 0)) || return 0 - -guard=Sourcing -reason= +# shellcheck disable=SC1091 +source "@logFromSetupHook@" + +# Guard helper function +# Returns 0 (success) if the hook should be run, 1 (failure) otherwise. +# This allows us to use short-circuit evaluation to avoid running the hook when it shouldn't be. +setupCudaGuard() { + log() { + logFromSetupHook \ + "${1:?}" \ + "setup-cuda" \ + "setupCudaGuard" \ + "${2:?}" + } + local guard="Skipping" + local reason="" + + if (( ${hostOffset:?} == -1 && ${targetOffset:?} == 0)); then + guard="Sourcing" + reason="because the hook is in nativeBuildInputs relative to the package being built" + elif [[ -n "${cudaSetupOnce-}" ]]; then + guard=Skipping + reason="because the hook has been propagated more than once" + fi -[[ -n ${cudaSetupOnce-} ]] && guard=Skipping && reason=" because the hook has been propagated more than once" + log "INFO" "$guard $reason" -if (( "${NIX_DEBUG:-0}" >= 1 )) ; then - echo "$guard hostOffset=$hostOffset targetOffset=$targetOffset setup-cuda$reason" >&2 -else - echo "$guard setup-cuda$reason" >&2 -fi + # Recall that test commands return 0 for success and 1 for failure. + [[ "$guard" == Sourcing ]] + return $? +} -[[ "$guard" = Sourcing ]] || return 0 +# Guard against calling the hook at the wrong time. +setupCudaGuard || return 0 +# NOTE: While it would appear that we are in the global scope, we are not! As such, we must declare these variables as +# global to ensure they are accessible in the global scope. declare -g cudaSetupOnce=1 -declare -Ag cudaHostPathsSeen=() -declare -Ag cudaOutputToPath=() +declare -gA cudaHostPathsSeen=() +declare -gA cudaOutputToPath=() extendcudaHostPathsSeen() { - (( "${NIX_DEBUG:-0}" >= 1 )) && echo "extendcudaHostPathsSeen $1" >&2 + log() { + logFromSetupHook \ + "${1:?}" \ + "setup-cuda" \ + "extendcudaHostPathsSeen" \ + "${2:?}" + } + local markerPath="${1:?}/nix-support/include-in-cudatoolkit-root" + + log "DEBUG" "checking for existence of $markerPath" + + if [[ ! -f "$markerPath" ]]; then + log "DEBUG" "skipping since $markerPath does not exist" + return 0 + fi - local markerPath="$1/nix-support/include-in-cudatoolkit-root" - [[ ! -f "${markerPath}" ]] && return 0 - [[ -v cudaHostPathsSeen[$1] ]] && return 0 + if [[ -v cudaHostPathsSeen["$1"] ]]; then + log "DEBUG" "skipping since $1 has already been seen" + return 0 + fi cudaHostPathsSeen["$1"]=1 + log "DEBUG" "added $1 to cudaHostPathsSeen" + + # Only attempt to read the file referenced by markerPath if strictDeps is not set; otherwise it is blank and we + # don't need to read it. + [[ -n "${strictDeps-}" ]] && return 0 # E.g. cuda_cudart-lib local cudaOutputName @@ -39,30 +80,52 @@ extendcudaHostPathsSeen() { [[ -z "$cudaOutputName" ]] && return 0 local oldPath="${cudaOutputToPath[$cudaOutputName]-}" - [[ -n "$oldPath" ]] && echo "extendcudaHostPathsSeen: warning: overwriting $cudaOutputName from $oldPath to $1" >&2 + if [[ -n "$oldPath" ]]; then + log "WARNING" "overwriting $cudaOutputName from $oldPath to $1" + fi cudaOutputToPath["$cudaOutputName"]="$1" } addEnvHooks "$targetOffset" extendcudaHostPathsSeen setupCUDAToolkit_ROOT() { - (( "${NIX_DEBUG:-0}" >= 1 )) && echo "setupCUDAToolkit_ROOT: cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2 - - for path in "${!cudaHostPathsSeen[@]}" ; do + log() { + logFromSetupHook \ + "${1:?}" \ + "setup-cuda" \ + "setupCUDAToolkit_ROOT" \ + "${2:?}" + } + log "DEBUG" "cudaHostPathsSeen is ${!cudaHostPathsSeen[*]}" + + for path in "${!cudaHostPathsSeen[@]}"; do addToSearchPathWithCustomDelimiter ";" CUDAToolkit_ROOT "$path" - if [[ -d "$path/include" ]] ; then - addToSearchPathWithCustomDelimiter ";" CUDAToolkit_INCLUDE_DIR "$path/include" - fi + [[ -d "$path/include" ]] && addToSearchPathWithCustomDelimiter ";" CUDAToolkit_INCLUDE_DIR "$path/include" done - export cmakeFlags+=" -DCUDAToolkit_INCLUDE_DIR=$CUDAToolkit_INCLUDE_DIR -DCUDAToolkit_ROOT=$CUDAToolkit_ROOT" + # NOTE: Due to the way CMake flags are structured within Nixpkgs, there is a distinction between flags set by + # cmakeFlags, and those set by cmakeFlagsArray (the former is a string, the latter is an array). + # The setup hook in Nixpkgs, as of this comment, interpolates the cmakeFlags string before expanding the array. + # Since we want these flags to have the lowest priority (and thus, be overridden), we must set them with + # cmakeFlags so all user flags will override them. + export cmakeFlags+=" -DCUDAToolkit_INCLUDE_DIR=\"${CUDAToolkit_INCLUDE_DIR:-}\"" + export cmakeFlags+=" -DCUDAToolkit_ROOT=\"${CUDAToolkit_ROOT:-}\"" } preConfigureHooks+=(setupCUDAToolkit_ROOT) setupCUDAToolkitCompilers() { - echo Executing setupCUDAToolkitCompilers >&2 - - if [[ -n "${dontSetupCUDAToolkitCompilers-}" ]] ; then + log() { + logFromSetupHook \ + "${1:?}" \ + "setup-cuda" \ + "setupCUDAToolkitCompilers" \ + "${2:?}" + } + + if [[ -n "${dontSetupCUDAToolkitCompilers-}" ]]; then + log "INFO" "skipping setup of CUDA toolkit compilers as requested" return 0 + else + log "INFO" "setting up CUDA toolkit compilers" fi # Point NVCC at a compatible compiler @@ -72,8 +135,9 @@ setupCUDAToolkitCompilers() { # 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@" + # NOTE: See note in setupCUDAToolkit_ROOT about why we use cmakeFlags over cmakeFlagsArray. + export cmakeFlags+=" -DCUDA_HOST_COMPILER=\"@ccFullPath@\"" + export cmakeFlags+=" -DCMAKE_CUDA_HOST_COMPILER=\"@ccFullPath@\"" # For non-CMake projects: # We prepend --compiler-bindir to nvcc flags. @@ -81,11 +145,9 @@ setupCUDAToolkitCompilers() { # 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 + [[ -z "${CUDAHOSTCXX-}" ]] && export CUDAHOSTCXX="@ccFullPath@" - export NVCC_PREPEND_FLAGS+=" --compiler-bindir=@ccRoot@/bin" + 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 @@ -93,27 +155,35 @@ setupCUDAToolkitCompilers() { # 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 + [[ -z "${dontCompressFatbin-}" ]] && export NVCC_PREPEND_FLAGS+=" -Xfatbin=-compress-all" } preConfigureHooks+=(setupCUDAToolkitCompilers) propagateCudaLibraries() { - (( "${NIX_DEBUG:-0}" >= 1 )) && echo "propagateCudaLibraries: cudaPropagateToOutput=$cudaPropagateToOutput cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2 + log() { + logFromSetupHook \ + "${1:?}" \ + "setup-cuda" \ + "propagateCudaLibraries" \ + "${2:?}" + } + + if [[ -n "${cudaPropagateToOutput-}" ]]; then + log "INFO" "propagating CUDA libraries to $cudaPropagateToOutput" + else + log "DEBUG" "skipping propagation of CUDA libraries since cudaPropagateToOutput is not set" + return 0 + fi - [[ -z "${cudaPropagateToOutput-}" ]] && return 0 + log "DEBUG" "cudaHostPathsSeen is ${!cudaHostPathsSeen[*]}" mkdir -p "${!cudaPropagateToOutput}/nix-support" # One'd expect this should be propagated-bulid-build-deps, but that doesn't seem to work echo "@setupCuda@" >> "${!cudaPropagateToOutput}/nix-support/propagated-native-build-inputs" local propagatedBuildInputs=( "${!cudaHostPathsSeen[@]}" ) - for output in $(getAllOutputNames) ; do - if [[ ! "$output" = "$cudaPropagateToOutput" ]] ; then - propagatedBuildInputs+=( "${!output}" ) - fi - break + for output in $(getAllOutputNames); do + [[ "$output" != "$cudaPropagateToOutput" ]] && propagatedBuildInputs+=( "${!output}" ) && break done # One'd expect this should be propagated-host-host-deps, but that doesn't seem to work