diff --git a/tvm/cppbuild.sh b/tvm/cppbuild.sh index ce51dd93f6f..95d81c61644 100755 --- a/tvm/cppbuild.sh +++ b/tvm/cppbuild.sh @@ -7,6 +7,9 @@ if [[ -z "$PLATFORM" ]]; then exit fi +export CUDACXX="/usr/local/cuda/bin/nvcc" +export CUDA_HOME="/usr/local/cuda" +export CUDNN_HOME="/usr/local/cuda" export GPU_FLAGS= if [[ "$EXTENSION" == *gpu ]]; then GPU_FLAGS="-DUSE_CUDA=ON -DUSE_CUDNN=ON -DUSE_CUBLAS=ON" @@ -170,6 +173,11 @@ case $PLATFORM in for f in $(find ../ -iname '*.dylib'); do install_name_tool -add_rpath @loader_path/ -add_rpath @loader_path/../../ $f || true; done ;; windows-x86_64) + if [[ -n "${CUDA_PATH:-}" ]]; then + export CUDACXX="$CUDA_PATH/bin/nvcc" + export CUDA_HOME="$CUDA_PATH" + export CUDNN_HOME="$CUDA_PATH" + fi export CC="cl.exe" export CXX="cl.exe" $CMAKE -G "Ninja" -DCMAKE_INSTALL_PREFIX="$INSTALL_PATH" -DCMAKE_BUILD_TYPE=Release -DUSE_LIBBACKTRACE=OFF -DUSE_BLAS=openblas -DCMAKE_PREFIX_PATH=$OPENBLAS_PATH -DUSE_LLVM=$LLVM_PATH/bin/llvm-config -DUSE_MKLDNN=$MKLDNN_PATH -DCMAKE_LIBRARY_PATH=$MKLDNN_PATH/lib -DUSE_DNNL_CODEGEN=ON -DUSE_DNNL=$MKLDNN_PATH $GPU_FLAGS -DUSE_OPENCL=$OPENCL_PATH .