# set(USE_MKL_PATH <path to venv or site-packages directory>) 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
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()