tjh.dev
1# shellcheck shell=bash
2
3# Only run the hook from nativeBuildInputs
4(( "$hostOffset" == -1 && "$targetOffset" == 0)) || return 0
5
6guard=Sourcing
7reason=
8
9[[ -n ${cudaSetupHookOnce-} ]] && guard=Skipping && reason=" because the hook has been propagated more than once"
10
11if (( "${NIX_DEBUG:-0}" >= 1 )) ; then
12 echo "$guard hostOffset=$hostOffset targetOffset=$targetOffset setup-cuda-hook$reason" >&2
13else
14 echo "$guard setup-cuda-hook$reason" >&2
15fi
16
17[[ "$guard" = Sourcing ]] || return 0
18
19declare -g cudaSetupHookOnce=1
20declare -Ag cudaHostPathsSeen=()
21declare -Ag cudaOutputToPath=()
22
23extendcudaHostPathsSeen() {
24 (( "${NIX_DEBUG:-0}" >= 1 )) && echo "extendcudaHostPathsSeen $1" >&2
25
26 local markerPath="$1/nix-support/include-in-cudatoolkit-root"
27 [[ ! -f "${markerPath}" ]] && return 0
28 [[ -v cudaHostPathsSeen[$1] ]] && return 0
29
30 cudaHostPathsSeen["$1"]=1
31
32 # E.g. cuda_cudart-lib
33 local cudaOutputName
34 # Fail gracefully if the file is empty.
35 # One reason the file may be empty: the package was built with strictDeps set, but the current build does not have
36 # strictDeps set.
37 read -r cudaOutputName < "$markerPath" || return 0
38
39 [[ -z "$cudaOutputName" ]] && return 0
40
41 local oldPath="${cudaOutputToPath[$cudaOutputName]-}"
42 [[ -n "$oldPath" ]] && echo "extendcudaHostPathsSeen: warning: overwriting $cudaOutputName from $oldPath to $1" >&2
43 cudaOutputToPath["$cudaOutputName"]="$1"
44}
45addEnvHooks "$targetOffset" extendcudaHostPathsSeen
46
47setupCUDAToolkit_ROOT() {
48 (( "${NIX_DEBUG:-0}" >= 1 )) && echo "setupCUDAToolkit_ROOT: cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2
49
50 for path in "${!cudaHostPathsSeen[@]}" ; do
51 addToSearchPathWithCustomDelimiter ";" CUDAToolkit_ROOT "$path"
52 if [[ -d "$path/include" ]] ; then
53 addToSearchPathWithCustomDelimiter ";" CUDAToolkit_INCLUDE_DIR "$path/include"
54 fi
55 done
56
57 appendToVar cmakeFlags "-DCUDAToolkit_INCLUDE_DIR=$CUDAToolkit_INCLUDE_DIR"
58 appendToVar cmakeFlags "-DCUDAToolkit_ROOT=$CUDAToolkit_ROOT"
59}
60preConfigureHooks+=(setupCUDAToolkit_ROOT)
61
62setupCUDAToolkitCompilers() {
63 echo Executing setupCUDAToolkitCompilers >&2
64
65 if [[ -n "${dontSetupCUDAToolkitCompilers-}" ]] ; then
66 return 0
67 fi
68
69 # Point NVCC at a compatible compiler
70
71 # For CMake-based projects:
72 # https://cmake.org/cmake/help/latest/module/FindCUDA.html#input-variables
73 # https://cmake.org/cmake/help/latest/envvar/CUDAHOSTCXX.html
74 # https://cmake.org/cmake/help/latest/variable/CMAKE_CUDA_HOST_COMPILER.html
75
76 appendToVar cmakeFlags "-DCUDA_HOST_COMPILER=@ccFullPath@"
77 appendToVar cmakeFlags "-DCMAKE_CUDA_HOST_COMPILER=@ccFullPath@"
78
79 # For non-CMake projects:
80 # We prepend --compiler-bindir to nvcc flags.
81 # Downstream packages can override these, because NVCC
82 # uses the last --compiler-bindir it gets on the command line.
83 # FIXME: this results in "incompatible redefinition" warnings.
84 # https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html#compiler-bindir-directory-ccbin
85 if [ -z "${CUDAHOSTCXX-}" ]; then
86 export CUDAHOSTCXX="@ccFullPath@";
87 fi
88
89 appendToVar NVCC_PREPEND_FLAGS "--compiler-bindir=@ccRoot@/bin"
90
91 # NOTE: We set -Xfatbin=-compress-all, which reduces the size of the compiled
92 # binaries. If binaries grow over 2GB, they will fail to link. This is a problem for us, as
93 # the default set of CUDA capabilities we build can regularly cause this to occur (for
94 # example, with Magma).
95 #
96 # @SomeoneSerge: original comment was made by @ConnorBaker in .../cudatoolkit/common.nix
97 if [[ -z "${dontCompressFatbin-}" ]]; then
98 appendToVar NVCC_PREPEND_FLAGS "-Xfatbin=-compress-all"
99 fi
100}
101preConfigureHooks+=(setupCUDAToolkitCompilers)
102
103propagateCudaLibraries() {
104 (( "${NIX_DEBUG:-0}" >= 1 )) && echo "propagateCudaLibraries: cudaPropagateToOutput=$cudaPropagateToOutput cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2
105
106 [[ -z "${cudaPropagateToOutput-}" ]] && return 0
107
108 mkdir -p "${!cudaPropagateToOutput}/nix-support"
109 # One'd expect this should be propagated-bulid-build-deps, but that doesn't seem to work
110 echo "@setupCudaHook@" >> "${!cudaPropagateToOutput}/nix-support/propagated-native-build-inputs"
111
112 local propagatedBuildInputs=( "${!cudaHostPathsSeen[@]}" )
113 for output in $(getAllOutputNames) ; do
114 if [[ ! "$output" = "$cudaPropagateToOutput" ]] ; then
115 appendToVar propagatedBuildInputs "${!output}"
116 fi
117 break
118 done
119
120 # One'd expect this should be propagated-host-host-deps, but that doesn't seem to work
121 printWords "${propagatedBuildInputs[@]}" >> "${!cudaPropagateToOutput}/nix-support/propagated-build-inputs"
122}
123postFixupHooks+=(propagateCudaLibraries)