about summary refs log tree commit diff
path: root/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
blob: 7b7b3bdde80e3045233b0bdc3138787ae8980b86 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
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)