From 6badb7ae875fa4ecdcd0d8ebc99d3c9dc6c425eb Mon Sep 17 00:00:00 2001 From: Jon Soifer Date: Mon, 27 Jan 2020 16:27:55 -0600 Subject: [PATCH] [Build] Explicitly link to cublasLt if it exists (#4776) * Explicitly link to cublasLt * Only link cublasLt if it's found Co-authored-by: Jon Soifer --- cmake/modules/CUDA.cmake | 3 +++ cmake/util/FindCUDA.cmake | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index 9f815710dc3e..f057560fec6a 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 4cec01c04db2..3161262fc2a2 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)