Skip to content

Commit

Permalink
Adding CUDA_VERSION check for compiling BERT plugins (#257)
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinch-nv authored Dec 6, 2019
1 parent 98ab844 commit daacf15
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -188,14 +188,21 @@ else()

include_directories(
${CUDA_INCLUDE_DIRS}
${CUDNN_ROOT_DIR}/include
)
find_library(CUDNN_LIB cudnn HINTS
${CUDA_TOOLKIT_ROOT_DIR} ${CUDNN_ROOT_DIR} PATH_SUFFIXES lib64 lib)
find_library(CUBLAS_LIB cublas HINTS
${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib64 lib lib/stubs)

find_library(CUBLASLT_LIB cublasLt HINTS
${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib64 lib lib/stubs)
# CUBLASLT libraries are only available in CUDA versions > 10. Check for CUDA version here and
# remove dependency on the libarary and unset BERT_GENCODES.
if (CUDA_VERSION VERSION_LESS_EQUAL 10.0)
message(WARNING "Detected CUDA version is <= 10.0! Removing BERT plugins from compilation list.")
unset(BERT_GENCODES)
else()
find_library(CUBLASLT_LIB cublasLt HINTS
${CUDA_TOOLKIT_ROOT_DIR} PATH_SUFFIXES lib64 lib lib/stubs)
endif()
if(BUILD_PARSERS)
configure_protobuf(${PROTOBUF_VERSION})
endif()
Expand Down

0 comments on commit daacf15

Please sign in to comment.