diff options
author | Connor Baker <connor.baker@tweag.io> | 2023-12-07 10:59:46 -0500 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-12-07 10:59:46 -0500 |
commit | c94fdf82d4f498d20be2291f730dbc693fa2efd5 (patch) | |
tree | c01a0d3eaf52b6ab10d30dab90809865e9d8f85e /pkgs/development/compilers | |
parent | e816589ee37c8429af497424a7f91ced729ac237 (diff) | |
parent | 2df7ccfa1498f5038b15acd50bc9277ad768dcbf (diff) |
Merge pull request #271078 from SomeoneSerge/feat/torch-propagated-cuda
cudaPackages.setupCudaHook: propagate deps and the hook
Diffstat (limited to 'pkgs/development/compilers')
4 files changed, 94 insertions, 40 deletions
diff --git a/pkgs/development/compilers/cudatoolkit/extension.nix b/pkgs/development/compilers/cudatoolkit/extension.nix index d75d288f5577e..016675fa07015 100644 --- a/pkgs/development/compilers/cudatoolkit/extension.nix +++ b/pkgs/development/compilers/cudatoolkit/extension.nix @@ -47,35 +47,21 @@ final: prev: let ./hooks/mark-for-cudatoolkit-root-hook.sh) { }); - # Normally propagated by cuda_nvcc or cudatoolkit through their depsHostHostPropagated + # 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++"; - - # Required by cmake's enable_language(CUDA) to build a test program - # When implementing cross-compilation support: this is - # final.pkgs.targetPackages.cudaPackages.cuda_cudart - # Given the multiple-outputs each CUDA redist has, we can specify the exact components we - # need from the package. CMake requires: - # - the cuda_runtime.h header, which is in the dev output - # - the dynamic library, which is in the lib output - # - the static library, which is in the static output - substitutions.cudartFlags = let cudart = final.cuda_cudart; in - builtins.concatStringsSep " " (final.lib.optionals (final ? cuda_cudart) ([ - "-I${final.lib.getDev cudart}/include" - "-L${final.lib.getLib cudart}/lib" - ] ++ final.lib.optionals (builtins.elem "static" cudart.outputs) [ - "-L${cudart.static}/lib" - ])); } ./hooks/setup-cuda-hook.sh) { }); diff --git a/pkgs/development/compilers/cudatoolkit/hooks/mark-for-cudatoolkit-root-hook.sh b/pkgs/development/compilers/cudatoolkit/hooks/mark-for-cudatoolkit-root-hook.sh index 5c18760a3a2b0..ba04c2e0806af 100644 --- a/pkgs/development/compilers/cudatoolkit/hooks/mark-for-cudatoolkit-root-hook.sh +++ b/pkgs/development/compilers/cudatoolkit/hooks/mark-for-cudatoolkit-root-hook.sh @@ -1,8 +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" - touch "${prefix}/nix-support/include-in-cudatoolkit-root" + [[ -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/compilers/cudatoolkit/hooks/nvcc-setup-hook.sh b/pkgs/development/compilers/cudatoolkit/hooks/nvcc-setup-hook.sh deleted file mode 100644 index e75a84a9550e7..0000000000000 --- a/pkgs/development/compilers/cudatoolkit/hooks/nvcc-setup-hook.sh +++ /dev/null @@ -1,5 +0,0 @@ -# shellcheck shell=bash - -# CMake's enable_language(CUDA) runs a compiler test and it doesn't account for -# CUDAToolkit_ROOT. We have to help it locate libcudart -export NVCC_APPEND_FLAGS+=" -L@cudartLib@/lib -L@cudartStatic@/lib -I@cudartInclude@/include" diff --git a/pkgs/development/compilers/cudatoolkit/hooks/setup-cuda-hook.sh b/pkgs/development/compilers/cudatoolkit/hooks/setup-cuda-hook.sh index 0fa8883081c50..7b7b3bdde80e3 100644 --- a/pkgs/development/compilers/cudatoolkit/hooks/setup-cuda-hook.sh +++ b/pkgs/development/compilers/cudatoolkit/hooks/setup-cuda-hook.sh @@ -3,19 +3,57 @@ # Only run the hook from nativeBuildInputs (( "$hostOffset" == -1 && "$targetOffset" == 0)) || return 0 -echo Sourcing setup-cuda-hook >&2 +guard=Sourcing +reason= -extendCUDAToolkit_ROOT() { - if [[ -f "$1/nix-support/include-in-cudatoolkit-root" ]] ; then - addToSearchPathWithCustomDelimiter ";" CUDAToolkit_ROOT "$1" +[[ -n ${cudaSetupHookOnce-} ]] && guard=Skipping && reason=" because the hook has been propagated more than once" - if [[ -d "$1/include" ]] ; then - addToSearchPathWithCustomDelimiter ";" CUDAToolkit_INCLUDE_DIR "$1/include" - fi - fi +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 -addEnvHooks "$targetOffset" extendCUDAToolkit_ROOT + 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 @@ -58,15 +96,44 @@ setupCUDAToolkitCompilers() { # CMake's enable_language(CUDA) runs a compiler test and it doesn't account for # CUDAToolkit_ROOT. We have to help it locate libcudart - local cudartFlags="@cudartFlags@" - if [[ -z "${nvccDontPrependCudartFlags-}" ]] && [[ -n "${cudartFlags:-}" ]] ; then - export NVCC_APPEND_FLAGS+=" $cudartFlags" + 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) -setupCMakeCUDAToolkit_ROOT() { - export cmakeFlags+=" -DCUDAToolkit_INCLUDE_DIR=$CUDAToolkit_INCLUDE_DIR -DCUDAToolkit_ROOT=$CUDAToolkit_ROOT" -} +propagateCudaLibraries() { + (( "${NIX_DEBUG:-0}" >= 1 )) && echo "propagateCudaLibraries: cudaPropagateToOutput=$cudaPropagateToOutput cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2 -postHooks+=(setupCUDAToolkitCompilers) -preConfigureHooks+=(setupCMakeCUDAToolkit_ROOT) + [[ -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) |