# 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)
# 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)
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()
#else
#include <cblas.h>
#endif
+#if USE_DNNL == 1
+#include <dnnl.h>
+#endif
}
namespace tvm {
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
}
};
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