diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index 785fa945cd87..250a65c591e1 100644 --- a/cmake/modules/CUDA.cmake +++ b/cmake/modules/CUDA.cmake @@ -50,6 +50,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 15ae2ccbd8b4..f971c87f20a8 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)