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

PiperOrigin-RevId: 204530874

tensorflow/tensorflow.bzl

@@ -24,7 +24,10 @@ load(
     "if_mkl",
     "if_mkl_lnx_x64"
 )
-
+load(
+    "//third_party/mkl_dnn:build_defs.bzl",
+    "if_mkl_open_source_only",
+)
 def register_extension_info(**kwargs):
     pass
 
@@ -214,6 +217,7 @@ def tf_copts(android_optimization_level_override="-O2", is_external=False):
       + if_cuda(["-DGOOGLE_CUDA=1"])
       + if_tensorrt(["-DGOOGLE_TENSORRT=1"])
       + if_mkl(["-DINTEL_MKL=1", "-DEIGEN_USE_VML"])
+      + if_mkl_open_source_only(["-DDO_NOT_USE_ML"])
       + if_mkl_lnx_x64(["-fopenmp"])
       + if_android_arm(["-mfpu=neon"])
       + if_linux_x86_64(["-msse3"])

tensorflow/tools/lib_package/BUILD

@@ -143,6 +143,7 @@ genrule(
         "@zlib_archive//:zlib.h",
     ] + if_mkl([
         "//third_party/mkl:LICENSE",
+        "//third_party/mkl_dnn:LICENSE",
     ]),
     outs = ["include/tensorflow/c/LICENSE"],
     cmd = "$(location :concat_licenses.sh) $(SRCS) >$@",
@@ -182,6 +183,7 @@ genrule(
         "@zlib_archive//:zlib.h",
     ] + if_mkl([
         "//third_party/mkl:LICENSE",
+        "//third_party/mkl_dnn:LICENSE",
     ]),
     outs = ["include/tensorflow/jni/LICENSE"],
     cmd = "$(location :concat_licenses.sh) $(SRCS) >$@",

tensorflow/tools/pip_package/BUILD

@@ -169,6 +169,7 @@ filegroup(
         "@org_python_pypi_backports_weakref//:LICENSE",
     ] + if_mkl([
         "//third_party/mkl:LICENSE",
+        "//third_party/mkl_dnn:LICENSE",
     ]) + tf_additional_license_deps(),
 )
 

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

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

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

tools/bazel.rc

@@ -27,6 +27,10 @@ build --define framework_shared_object=true
 build:mkl --define=using_mkl=true
 build:mkl -c opt
 
+# This config option is used to enable MKL-DNN open source library only,
+# without depending on MKL binary version.
+build:mkl_open_source_only --define=using_mkl_dnn_only=true
+
 build:download_clang --crosstool_top=@local_config_download_clang//:toolchain
 build:download_clang --define=using_clang=true