allow customize mkldnn library location (#4814)

cmake/config.cmake

@@ -132,7 +132,7 @@ set(USE_BLAS none)
 # 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

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()