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