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.

Sign up for GitHub

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

Jump to bottom

[Contrib] Add MKL DNN option #4323

Merged
merged 3 commits into from
Nov 15, 2019
Merged

[Contrib] Add MKL DNN option #4323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
914b5ca
[Contrib] Add MKL DNN
icemelon Oct 1, 2019
b8a0843
update
icemelon Nov 13, 2019
e5e0774
update
icemelon Nov 14, 2019
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
View file
Open in desktop
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_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)
Expand Down
3 changes: 3 additions & 0 deletions cmake/config.cmake
View file
Open in desktop
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 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)
Expand Down
7 changes: 7 additions & 0 deletions cmake/modules/contrib/BLAS.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
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 •
edited
Loading

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
View file
Open in desktop
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_DNNL == 1
#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_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
}
};

Expand Down
2 changes: 1 addition & 1 deletion topi/python/topi/x86/dense.py
View file
Open in desktop
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