Fix CUDA builds for TVM

bytedeco · May 1, 2024 · 7d5e6a0 · 7d5e6a0
1 parent c27ab7d
commit 7d5e6a0
Showing 1 changed file with 8 additions and 0 deletions.
diff --git 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 .