diff --git a/.bazelrc b/.bazelrc index 590a87f5732093..f4d51e7f81ed58 100644 --- a/.bazelrc +++ b/.bazelrc @@ -35,6 +35,8 @@ build --define framework_shared_object=true # environment variable "TF_MKL_ROOT" every time before build. build:mkl --define=build_with_mkl=true --define=enable_mkl=true build:mkl --define=tensorflow_mkldnn_contraction_kernel=0 +build:mkl --define build_with_mkl_dnn_only=true +build:mkl --define mkl_dnn_threading=tf build:mkl -c opt # This config option is used to enable MKL-DNN open source library only, @@ -43,6 +45,7 @@ build:mkl_open_source_only --define=build_with_mkl_dnn_only=true build:mkl_open_source_only --define=build_with_mkl_dnn_v1_only=true build:mkl_open_source_only --define=build_with_mkl=true --define=enable_mkl=true build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=0 +build:mkl_open_source_only --define mkl_dnn_threading=tf build:download_clang --crosstool_top=@local_config_download_clang//:toolchain build:download_clang --define=using_clang=true diff --git a/tensorflow/core/common_runtime/local_device.cc b/tensorflow/core/common_runtime/local_device.cc index 84f3eca532df08..7a25dbaea0a9bb 100644 --- a/tensorflow/core/common_runtime/local_device.cc +++ b/tensorflow/core/common_runtime/local_device.cc @@ -30,6 +30,10 @@ limitations under the License. #include "tensorflow/core/public/session_options.h" #include "tensorflow/core/util/env_var.h" +#if defined(MKLDNN_TF_THREADING) +#include "mkldnn.h" +#endif + namespace tensorflow { namespace { @@ -155,6 +159,9 @@ LocalDevice::LocalDevice(const SessionOptions& options, tp_info = owned_tp_info_.get(); } set_tensorflow_cpu_worker_threads(&tp_info->eigen_worker_threads_); +#if defined(MKLDNN_TF_THREADING) + mkldnn_set_tensorflow_thread_pool(tp_info->eigen_worker_threads_.workers); +#endif set_eigen_cpu_device(tp_info->eigen_device_.get()); } diff --git a/tensorflow/core/kernels/mkl_conv_ops_test.cc b/tensorflow/core/kernels/mkl_conv_ops_test.cc index a055351337c7c6..a95f80f0a18180 100644 --- a/tensorflow/core/kernels/mkl_conv_ops_test.cc +++ b/tensorflow/core/kernels/mkl_conv_ops_test.cc @@ -28,7 +28,6 @@ limitations under the License. #include "tensorflow/core/public/session.h" #if defined(INTEL_MKL_DNN_ONLY) -#include "third_party/intel_mkl_dnn/include/mkldnn.h" #include "tensorflow/core/util/mkl_util.h" #endif diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 0de95b67542f30..4edc5c8719c9c7 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -45,6 +45,7 @@ load( "//third_party/mkl_dnn:build_defs.bzl", "if_mkl_open_source_only", "if_mkl_v1_open_source_only", + "if_mkl_dnn_uses_tf_threading", ) load( "//third_party/ngraph:build_defs.bzl", @@ -296,6 +297,7 @@ def tf_copts(android_optimization_level_override = "-O2", is_external = False): if_enable_mkl(["-DENABLE_MKL"]) + if_ngraph(["-DINTEL_NGRAPH=1"]) + if_mkl_lnx_x64(["-fopenmp"]) + + if_mkl_dnn_uses_tf_threading(["-fno-openmp -DMKLDNN_TF_THREADING"]) + if_android_arm(["-mfpu=neon"]) + if_linux_x86_64(["-msse3"]) + if_ios_x86_64(["-msse4.1"]) + diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 0303a49982d3b3..031ca981dde3c8 100755 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -133,11 +133,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""): tf_http_archive( name = "mkl_dnn", build_file = clean_dep("//third_party/mkl_dnn:mkldnn.BUILD"), - sha256 = "26f720ed912843ba293e8a1e0822fe5318e93c529d80c87af1cf555d68e642d0", - strip_prefix = "mkl-dnn-0.20.1", + sha256 = "d11c296c9f228c9e5c69d891daf2463d5acb2c5932cd7bce9f8654b206fa303d", + strip_prefix = "mkl-dnn-c5f6005a8220f9a928a698e54421b3d1342df64b", urls = [ - "http://mirror.tensorflow.org/github.com/intel/mkl-dnn/archive/v0.20.1.tar.gz", - "https://github.com/intel/mkl-dnn/archive/v0.20.1.tar.gz", + "https://mirror.bazel.build/github.com/rsdubtso/mkl-dnn/archive/c5f6005a8220f9a928a698e54421b3d1342df64b.tar.gz", + "https://github.com/rsdubtso/mkl-dnn/archive/c5f6005a8220f9a928a698e54421b3d1342df64b.tar.gz", ], ) diff --git a/third_party/mkl_dnn/BUILD b/third_party/mkl_dnn/BUILD index 5229dd5aa74ff3..788202047fa9ac 100644 --- a/third_party/mkl_dnn/BUILD +++ b/third_party/mkl_dnn/BUILD @@ -19,3 +19,12 @@ config_setting( }, visibility = ["//visibility:public"], ) + +config_setting( + name = "mkl_dnn_use_tf_threading", + define_values = { + "mkl_dnn_threading" : "tf", + "build_with_mkl_dnn_only": "true", + }, + visibility = ["//visibility:public"], +) diff --git a/third_party/mkl_dnn/build_defs.bzl b/third_party/mkl_dnn/build_defs.bzl index 384b528c2733aa..6aefc14d6a4a08 100644 --- a/third_party/mkl_dnn/build_defs.bzl +++ b/third_party/mkl_dnn/build_defs.bzl @@ -1,3 +1,6 @@ +def clean_dep(dep): + return str(Label(dep)) + def if_mkl_open_source_only(if_true, if_false = []): """Returns `if_true` if MKL-DNN v0.x is used. @@ -10,7 +13,13 @@ def if_mkl_open_source_only(if_true, if_false = []): """ return select({ - str(Label("//third_party/mkl_dnn:build_with_mkl_dnn_only")): if_true, + clean_dep("//third_party/mkl_dnn:build_with_mkl_dnn_only"): if_true, + "//conditions:default": if_false, + }) + +def if_mkl_dnn_uses_tf_threading(if_true, if_false = []): + return select({ + clean_dep("//third_party/mkl_dnn:mkl_dnn_use_tf_threading"): if_true, "//conditions:default": if_false, }) diff --git a/third_party/mkl_dnn/mkldnn.BUILD b/third_party/mkl_dnn/mkldnn.BUILD index 6331a108e502fa..2eca42e7e3a2cd 100644 --- a/third_party/mkl_dnn/mkldnn.BUILD +++ b/third_party/mkl_dnn/mkldnn.BUILD @@ -4,6 +4,7 @@ load( "@org_tensorflow//third_party/mkl_dnn:build_defs.bzl", "if_mkl_open_source_only", "if_mkl_v1_open_source_only", + "if_mkl_dnn_uses_tf_threading", ) load( "@org_tensorflow//third_party:common.bzl", @@ -50,20 +51,36 @@ template_rule( }, ) +cc_library( + name = "mkl_dnn_mkl_deps", + deps = select({ + "@org_tensorflow//tensorflow:linux_x86_64": [ + "@mkl_linux//:mkl_headers", + "@mkl_linux//:mkl_libs_linux", + ], + "@org_tensorflow//tensorflow:macos": [ + "@mkl_macos//:mkl_headers", + "@mkl_macos//:mkl_libs_macos", + ], + "@org_tensorflow//tensorflow:windows": [ + "@mkl_windows//:mkl_headers", + "@mkl_windows//:mkl_libs_windows", + ], + "//conditions:default": [], + }), +) + cc_library( name = "mkl_dnn", srcs = glob([ "src/common/*.cpp", "src/common/*.hpp", - "src/cpu/*.cpp", - "src/cpu/*.hpp", "src/cpu/**/*.cpp", "src/cpu/**/*.hpp", - "src/cpu/xbyak/*.h", + "src/cpu/**/*.c", + "src/cpu/**/*.h", ]) + if_mkl_v1_open_source_only([ ":mkldnn_config_h", - "src/cpu/jit_utils/jit_utils.cpp", - "src/cpu/jit_utils/jit_utils.hpp", ]) + [":mkldnn_version_h"], hdrs = glob(["include/*"]), copts = [ @@ -84,7 +101,10 @@ cc_library( # dependency. ":clang_linux_x86_64": [], "//conditions:default": [], - }), + }) + if_mkl_dnn_uses_tf_threading([ + "-DMKLDNN_THR=MKLDNN_THR_TENSORFLOW", + "-fno-openmp", + ]), includes = [ "include", "src", @@ -95,21 +115,14 @@ cc_library( ], nocopts = "-fno-exceptions", visibility = ["//visibility:public"], - deps = select({ - "@org_tensorflow//tensorflow:linux_x86_64": [ - "@mkl_linux//:mkl_headers", - "@mkl_linux//:mkl_libs_linux", + deps = if_mkl_dnn_uses_tf_threading( + [ + "@eigen_archive//:eigen", + "@com_google_protobuf//:protobuf_headers", + "@org_tensorflow//tensorflow/core:core_cpu_headers_lib", + "@org_tensorflow//tensorflow/core:framework_headers_lib", ], - "@org_tensorflow//tensorflow:macos": [ - "@mkl_darwin//:mkl_headers", - "@mkl_darwin//:mkl_libs_darwin", - ], - "@org_tensorflow//tensorflow:windows": [ - "@mkl_windows//:mkl_headers", - "@mkl_windows//:mkl_libs_windows", - ], - "//conditions:default": [], - }), + ["mkl_dnn_mkl_deps"]) ) cc_library( @@ -117,11 +130,10 @@ cc_library( srcs = glob([ "src/common/*.cpp", "src/common/*.hpp", - "src/cpu/*.cpp", - "src/cpu/*.hpp", "src/cpu/**/*.cpp", "src/cpu/**/*.hpp", - "src/cpu/xbyak/*.h", + "src/cpu/**/*.c", + "src/cpu/**/*.h", ]) + [":mkldnn_version_h"], hdrs = glob(["include/*"]), copts = [