[Contrib] Add MKL DNN option (apache#4323)

* [Contrib] Add MKL DNN * update * update
zxy844288792 · Nov 15, 2019 · fd1c187 · fd1c187
1 parent d3a88e0
commit fd1c187
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 1 deletion.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -53,6 +53,7 @@ tvm_option(PICOJSON_PATH "Path to PicoJSON" "3rdparty/picojson")
 # 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_MKLDNN "Build with MKLDNN" OFF)
 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
@@ -115,6 +115,9 @@ 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
+set(USE_MKLDNN OFF)
+
 # Whether use OpenMP thread pool, choices: gnu, intel
 # Note: "gnu" uses gomp library, "intel" uses iomp5 library
 set(USE_OPENMP none)

diff --git a/cmake/modules/contrib/BLAS.cmake b/cmake/modules/contrib/BLAS.cmake
@@ -55,3 +55,10 @@ elseif(USE_BLAS STREQUAL "none")
 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})
+endif()
diff --git a/src/runtime/contrib/cblas/cblas.cc b/src/runtime/contrib/cblas/cblas.cc
@@ -31,6 +31,9 @@ extern "C" {
 #else
 #include <cblas.h>
 #endif
+#if USE_DNNL == 1
+#include <dnnl.h>
+#endif
 }
 
 namespace tvm {
@@ -40,12 +43,19 @@ using namespace runtime;
 
 inline CBLAS_TRANSPOSE BooleanToTranspose(bool trans) { return trans ? CblasTrans : CblasNoTrans; }
 
+inline char BooleanToTransposeChar(bool trans) { return trans ? 'T' : 'N'; }
+
 struct CblasSgemmOp {
   typedef float TDatatype;
   void operator()(bool ta, bool tb, int M, int N, int K, float alpha, float* A, int lda, float* B,
                   int ldb, float beta, float* C, int ldc) {
+#if USE_DNNL == 1
+    dnnl_sgemm(BooleanToTransposeChar(tb), BooleanToTransposeChar(ta), N, M, K, alpha, B,
+               ldb, A, lda, beta, C, ldc);
+#else
     cblas_sgemm(CblasColMajor, BooleanToTranspose(ta), BooleanToTranspose(tb), M, N, K, alpha, A,
                 lda, B, ldb, beta, C, ldc);
+#endif
   }
 };
 

diff --git a/topi/python/topi/x86/dense.py b/topi/python/topi/x86/dense.py
@@ -32,7 +32,7 @@ def _declaration_dense(cfg, data, weight, bias=None, out_dtype=None):
     if "cblas" in target.libs:
         C = cblas.matmul(data, weight, False, True)
         if bias is not None:
-            C = tvm.compute(C.shape, lambda i, j: C[i, j] + bias[j].astype(out_dtype),
+            C = tvm.compute(C.shape, lambda i, j: C[i, j] + bias[j],
                             tag=tag.BROADCAST)
         return C