allow customize mkldnn library location (#4814)

apache · Feb 5, 2020 · 3e7bd70 · 3e7bd70
1 parent f9b46c4
commit 3e7bd70
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/cmake/config.cmake b/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

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