diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index 9f815710dc3ec..f057560fec6aa 100644 --- a/cmake/modules/CUDA.cmake +++ b/cmake/modules/CUDA.cmake @@ -51,6 +51,9 @@ if(USE_CUDA) file(GLOB CONTRIB_CUBLAS_SRCS src/runtime/contrib/cublas/*.cc) list(APPEND RUNTIME_SRCS ${CONTRIB_CUBLAS_SRCS}) list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIBRARY}) + if(NOT CUDA_CUBLASLT_LIBRARY STREQUAL "CUDA_CUBLASLT_LIBRARY-NOTFOUND") + list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLASLT_LIBRARY}) + endif() endif(USE_CUBLAS) else(USE_CUDA) diff --git a/cmake/util/FindCUDA.cmake b/cmake/util/FindCUDA.cmake index 4cec01c04db2d..3161262fc2a27 100644 --- a/cmake/util/FindCUDA.cmake +++ b/cmake/util/FindCUDA.cmake @@ -70,6 +70,9 @@ macro(find_cuda use_cuda) find_library(CUDA_CUBLAS_LIBRARY cublas ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64 ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32) + find_library(CUDA_CUBLASLT_LIBRARY cublaslt + ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64 + ${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32) else(MSVC) find_library(_CUDA_CUDA_LIBRARY cuda PATHS ${CUDA_TOOLKIT_ROOT_DIR} @@ -88,6 +91,9 @@ macro(find_cuda use_cuda) find_library(CUDA_CUBLAS_LIBRARY cublas ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib) + find_library(CUDA_CUBLASLT_LIBRARY cublaslt + ${CUDA_TOOLKIT_ROOT_DIR}/lib64 + ${CUDA_TOOLKIT_ROOT_DIR}/lib) endif(MSVC) message(STATUS "Found CUDA_TOOLKIT_ROOT_DIR=" ${CUDA_TOOLKIT_ROOT_DIR}) message(STATUS "Found CUDA_CUDA_LIBRARY=" ${CUDA_CUDA_LIBRARY}) @@ -95,5 +101,6 @@ macro(find_cuda use_cuda) message(STATUS "Found CUDA_NVRTC_LIBRARY=" ${CUDA_NVRTC_LIBRARY}) message(STATUS "Found CUDA_CUDNN_LIBRARY=" ${CUDA_CUDNN_LIBRARY}) message(STATUS "Found CUDA_CUBLAS_LIBRARY=" ${CUDA_CUBLAS_LIBRARY}) + message(STATUS "Found CUDA_CUBLASLT_LIBRARY=" ${CUDA_CUBLASLT_LIBRARY}) endif(CUDA_FOUND) endmacro(find_cuda)