From 9b8cb1b669467329a89bdb53b3b2f9d3594ad3a9 Mon Sep 17 00:00:00 2001
From: Tao Lv <tao.a.lv@intel.com>
Date: Tue, 26 Jun 2018 00:57:37 +0800
Subject: [PATCH] Contrib: add mkl blas support (#1336)
---
CMakeLists.txt | 1 +
cmake/config.cmake | 7 ++++++-
cmake/modules/contrib/BLAS.cmake | 10 ++++++++++
src/contrib/cblas/cblas.cc | 4 ++++
4 files changed, 21 insertions(+), 1 deletion(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index e9699b1f6..028c16aa5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -42,6 +42,7 @@ tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
# Contrib library options
tvm_option(USE_BLAS "The blas library to be linked" none)
+tvm_option(USE_MKL_PATH "MKL root path when use MKL blas" none)
tvm_option(USE_CUDNN "Build with cuDNN" OFF)
tvm_option(USE_CUBLAS "Build with cuBLAS" OFF)
tvm_option(USE_MIOPEN "Build with ROCM:MIOpen" OFF)
diff --git a/cmake/config.cmake b/cmake/config.cmake
index 1f02bfd7d..d2576ab30 100644
--- a/cmake/config.cmake
+++ b/cmake/config.cmake
@@ -77,9 +77,14 @@ set(USE_LLVM OFF)
#---------------------------------------------
# Contrib libraries
#---------------------------------------------
-# Whether use BLAS, choices: openblas, atlas, blas, apple
+# Whether use BLAS, choices: openblas, mkl, atlas, apple
set(USE_BLAS none)
+# /path/to/mkl: mkl root path when use mkl blas library
+# set(USE_MKL_PATH /opt/intel/mkl) for UNIX
+# set(USE_MKL_PATH ../IntelSWTools/compilers_and_libraries_2018/windows/mkl) for WIN32
+set(USE_MKL_PATH none)
+
# Whether use contrib.random in runtime
set(USE_RANDOM OFF)
diff --git a/cmake/modules/contrib/BLAS.cmake b/cmake/modules/contrib/BLAS.cmake
index 423cb9cd5..45269a207 100644
--- a/cmake/modules/contrib/BLAS.cmake
+++ b/cmake/modules/contrib/BLAS.cmake
@@ -6,6 +6,16 @@ if(USE_BLAS STREQUAL "openblas")
list(APPEND TVM_RUNTIME_LINKER_LIBS ${BLAS_LIBRARY})
list(APPEND RUNTIME_SRCS ${CBLAS_CONTRIB_SRC})
message(STATUS "Use BLAS library " ${BLAS_LIBRARY})
+elseif(USE_BLAS STREQUAL "mkl")
+ if(NOT IS_DIRECTORY ${USE_MKL_PATH})
+ set(USE_MKL_PATH /opt/intel/mkl)
+ endif()
+ find_library(BLAS_LIBRARY mkl_rt ${USE_MKL_PATH}/lib/ ${USE_MKL_PATH}/lib/intel64)
+ include_directories(${USE_MKL_PATH}/include)
+ list(APPEND TVM_RUNTIME_LINKER_LIBS ${BLAS_LIBRARY})
+ list(APPEND RUNTIME_SRCS ${CBLAS_CONTRIB_SRC})
+ add_definitions(-DUSE_MKL_BLAS=1)
+ message(STATUS "Use BLAS library " ${BLAS_LIBRARY})
elseif(USE_BLAS STREQUAL "atlas" OR USE_BLAS STREQUAL "blas")
find_library(BLAS_LIBRARY cblas)
list(APPEND TVM_RUNTIME_LINKER_LIBS ${BLAS_LIBRARY})
diff --git a/src/contrib/cblas/cblas.cc b/src/contrib/cblas/cblas.cc
index c2e981ba5..24ed9deb9 100644
--- a/src/contrib/cblas/cblas.cc
+++ b/src/contrib/cblas/cblas.cc
@@ -7,7 +7,11 @@
#include <dmlc/logging.h>
extern "C" {
+#if USE_MKL_BLAS == 1
+#include <mkl_cblas.h>
+#else
#include <cblas.h>
+#endif
}
namespace tvm {
--
GitLab