PoC for using MKL-DNN with TensorFlow thread pool

The .bazelrc updated:
- TF threading is now default for MKL-DNN
- linking with binary MKL is now disabled by default

Limitations:
- only a single session without inter-op parallelism is supported
- XLA not covered

.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

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());
 }
 

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
 

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"]) +

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",
         ],
     )
 

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"],
+)

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,
     })
 

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,33 +115,25 @@ 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(
     name = "mkldnn_single_threaded",
     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 = [