diff --git a/cmake/config.cmake b/cmake/config.cmake index 22db32e66e40..dcec40e39d1c 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -132,7 +132,7 @@ set(USE_BLAS none) # set(USE_MKL_PATH ) if using `pip install mkl` set(USE_MKL_PATH none) -# Whether use MKLDNN library +# Whether use MKLDNN library, choices: ON, OFF, path to mkldnn library set(USE_MKLDNN OFF) # Whether use OpenMP thread pool, choices: gnu, intel diff --git a/cmake/modules/contrib/BLAS.cmake b/cmake/modules/contrib/BLAS.cmake index bd8c0d0c445f..c2e1fd65743b 100644 --- a/cmake/modules/contrib/BLAS.cmake +++ b/cmake/modules/contrib/BLAS.cmake @@ -56,9 +56,27 @@ else() message(FATAL_ERROR "Invalid option: USE_BLAS=" ${USE_BLAS}) endif() -if(USE_MKLDNN STREQUAL "ON") - find_library(BLAS_LIBRARY_MKLDNN dnnl) - list(APPEND TVM_RUNTIME_LINKER_LIBS ${BLAS_LIBRARY_MKLDNN}) - add_definitions(-DUSE_DNNL=1) - message(STATUS "Use MKLDNN library " ${BLAS_LIBRARY_MKLDNN}) +if(IS_DIRECTORY ${USE_MKLDNN}) + find_library(MKLDNN_LIBRARY NAMES dnnl HINTS ${USE_MKLDNN}/lib/) + if (MKLDNN_LIBRARY STREQUAL "MKLDNN_LIBRARY-NOTFOUND") + message(WARNING "Cannot find MKLDNN library at ${USE_MKLDNN}.") + else() + include_directories(${USE_MKLDNN}/include) + list(APPEND TVM_RUNTIME_LINKER_LIBS ${MKLDNN_LIBRARY}) + add_definitions(-DUSE_DNNL=1) + message(STATUS "Use MKLDNN library " ${MKLDNN_LIBRARY}) + endif() +elseif(USE_MKLDNN STREQUAL "ON") + find_library(MKLDNN_LIBRARY dnnl) + if (MKLDNN_LIBRARY STREQUAL "MKLDNN_LIBRARY-NOTFOUND") + message(WARNING "Cannot find MKLDNN library. Try to specify the path to MKLDNN library.") + else() + list(APPEND TVM_RUNTIME_LINKER_LIBS ${MKLDNN_LIBRARY}) + add_definitions(-DUSE_DNNL=1) + message(STATUS "Use MKLDNN library " ${MKLDNN_LIBRARY}) + endif() +elseif(USE_MKLDNN STREQUAL "OFF") + # pass +else() + message(FATAL_ERROR "Invalid option: USE_MKLDNN=" ${USE_MKLDNN}) endif()