Merge pull request #20576 from Intel-tensorflow:mabuzain/avx-performance-fix

PiperOrigin-RevId: 204530874

3 files changed, 51 insertions, 1 deletions

diff --git a/third_party/mkl_dnn/BUILD b/third_party/mkl_dnn/BUILD
index 5b01f6e3e4..d075809ee9 100644
--- a/third_party/mkl_dnn/BUILD
+++ b/third_party/mkl_dnn/BUILD
@@ -1 +1,11 @@
 licenses(["notice"])
+
+exports_files(["LICENSE"])
+
+config_setting(
+    name = "using_mkl_dnn_only",
+    values = {
+        "define": "using_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
new file mode 100644
index 0000000000..7ce2a7d9b0
--- /dev/null
+++ b/third_party/mkl_dnn/build_defs.bzl
@@ -0,0 +1,13 @@
+def if_mkl_open_source_only(if_true, if_false = []):
+    """Shorthand for select()'ing on whether we're building with
+    MKL-DNN open source lib only, without depending on MKL binary form.
+
+    Returns a select statement which evaluates to if_true if we're building
+    with MKL-DNN open source lib only. Otherwise,
+    the select statement evaluates to if_false.
+
+    """
+    return select({
+        str(Label("//third_party/mkl_dnn:using_mkl_dnn_only")): if_true,
+        "//conditions:default": if_false,
+    })
diff --git a/third_party/mkl_dnn/mkldnn.BUILD b/third_party/mkl_dnn/mkldnn.BUILD
index 68f24aabae..57d2e1292b 100644
--- a/third_party/mkl_dnn/mkldnn.BUILD
+++ b/third_party/mkl_dnn/mkldnn.BUILD
@@ -1,5 +1,10 @@
 exports_files(["LICENSE"])
 
+load(
+    "@org_tensorflow//third_party/mkl_dnn:build_defs.bzl",
+    "if_mkl_open_source_only",
+)
+
 config_setting(
     name = "clang_linux_x86_64",
     values = {
@@ -15,7 +20,14 @@ cc_library(
         "src/cpu/*.cpp",
     ]),
     hdrs = glob(["include/*"]),
-    copts = ["-fexceptions"] + select({
+    copts = [
+        "-fexceptions",
+        "-DUSE_MKL",
+        "-DUSE_CBLAS",
+    ] + if_mkl_open_source_only([
+        "-UUSE_MKL",
+        "-UUSE_CBLAS",
+    ]) + select({
         "@org_tensorflow//tensorflow:linux_x86_64": [
             "-fopenmp",  # only works with gcc
         ],
@@ -33,4 +45,19 @@ 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",
+        ],
+        "@org_tensorflow//tensorflow:darwin": [
+            "@mkl_darwin//:mkl_headers",
+            "@mkl_darwin//:mkl_libs_darwin",
+        ],
+        "@org_tensorflow//tensorflow:windows": [
+            "@mkl_windows//:mkl_headers",
+            "@mkl_windows//:mkl_libs_windows",
+        ],
+        "//conditions:default": [],
+    }),
 )