blob: 70c8c8eebe28ffcd7746266c4d6d28ac1f779c66 [file] [log] [blame]
# CUDA Module
find_cuda(${USE_CUDA})
if(CUDA_FOUND)
# always set the includedir when cuda is available
# avoid global retrigger of cmake
include_directories(${CUDA_INCLUDE_DIRS})
endif(CUDA_FOUND)
if(USE_CUDA)
if(NOT CUDA_FOUND)
message(FATAL_ERROR "Cannot find CUDA, USE_CUDA=" ${USE_CUDA})
endif()
message(STATUS "Build with CUDA support")
file(GLOB RUNTIME_CUDA_SRCS src/runtime/cuda/*.cc)
list(APPEND RUNTIME_SRCS ${RUNTIME_CUDA_SRCS})
list(APPEND COMPILER_SRCS src/codegen/opt/build_cuda_on.cc)
list(APPEND TVM_LINKER_LIBS ${CUDA_NVRTC_LIBRARY})
list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDART_LIBRARY})
list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDA_LIBRARY})
list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_NVRTC_LIBRARY})
if(USE_CUDNN)
message(STATUS "Build with cuDNN support")
file(GLOB CONTRIB_CUDNN_SRCS src/contrib/cudnn/*.cc)
list(APPEND RUNTIME_SRCS ${CONTRIB_CUDNN_SRCS})
list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIBRARY})
endif(USE_CUDNN)
if(USE_CUBLAS)
message(STATUS "Build with cuBLAS support")
file(GLOB CONTRIB_CUBLAS_SRCS src/contrib/cublas/*.cc)
list(APPEND RUNTIME_SRCS ${CONTRIB_CUBLAS_SRCS})
list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIBRARY})
endif(USE_CUBLAS)
else(USE_CUDA)
list(APPEND COMPILER_SRCS src/codegen/opt/build_cuda_off.cc)
endif(USE_CUDA)