Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Contrib] Add MKL DNN option #4323

Merged
merged 3 commits into from
Nov 15, 2019
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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_MKL_DNN "Build with MKL DNN" 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)
Expand Down
3 changes: 3 additions & 0 deletions cmake/config.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -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 MKL DNN library
set(USE_MKL_DNN OFF)

# Whether use OpenMP thread pool, choices: gnu, intel
# Note: "gnu" uses gomp library, "intel" uses iomp5 library
set(USE_OPENMP none)
Expand Down
8 changes: 8 additions & 0 deletions cmake/modules/contrib/BLAS.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -55,3 +55,11 @@ elseif(USE_BLAS STREQUAL "none")
else()
message(FATAL_ERROR "Invalid option: USE_BLAS=" ${USE_BLAS})
endif()

if(USE_MKL_DNN STREQUAL "ON")
find_library(BLAS_LIBRARY_MKLDNN dnnl)
list(APPEND TVM_RUNTIME_LINKER_LIBS ${BLAS_LIBRARY_MKLDNN})
list(APPEND RUNTIME_SRCS ${CBLAS_CONTRIB_SRC})
add_definitions(-DUSE_MKL_DNN=1)
message(STATUS "Use MKL DNN library " ${BLAS_LIBRARY_MKLDNN})
endif()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the strategy if both USE_BLAS AND USE_MKL_DNN are set?

Copy link
Member Author

@icemelon icemelon Nov 13, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the logic here. MKL DNN will only be used when USE_BLAS is not none. When both are set, MKL DNN will be only used in the sgemm op as the library has limited support for the BLAS operators. And I find that MKL DNN kernel achieves better performance than MKL in TVM.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another minor nit is to change USE_MKL_DNN to either USE_DNNL or USE_MKLDNN. The first one aligns with the renaming trend of the library. The second one follows the coding convention in MKL-DNN before renaming and in other projects like MXNet.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the quick turnaround.

10 changes: 10 additions & 0 deletions src/runtime/contrib/cblas/cblas.cc
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ extern "C" {
#else
#include <cblas.h>
#endif
#if USE_MKL_DNN == 1
Copy link
Contributor

@ZhennanQin ZhennanQin Nov 13, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think code logic here need a small change as well:

#if USE_MKL_BLAS == 1
#include <mkl_cblas.h>
#elif USE_MKL_DNN == 1
#include <dnnl.h>
#else
#include <cblas.h>
#endif

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ZhennanQin As @icemelon9 mentioned here, we need both cblas and dnnl because the latter is used for sgemm only.

#include <dnnl.h>
#endif
}

namespace tvm {
Expand All @@ -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_MKL_DNN == 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
}
};

Expand Down
2 changes: 1 addition & 1 deletion topi/python/topi/x86/dense.py
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change doesn't seem to be related to mkl dnn?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, this is to fix the bug when using cblas library

tag=tag.BROADCAST)
return C

Expand Down