Re-enable using local MKL directory.

Setting TF_MKL_ROOT environment variable (together with "--config=mkl")
before build will enforce using a local copy of MKL during build.

PiperOrigin-RevId: 162787130

2 files changed, 65 insertions, 2 deletions

diff --git a/third_party/mkl/build_defs.bzl b/third_party/mkl/build_defs.bzl
index 9a28b312c2..533c0766c7 100644
--- a/third_party/mkl/build_defs.bzl
+++ b/third_party/mkl/build_defs.bzl
@@ -1,4 +1,16 @@
-# Macros for building MKL code.
+# -*- Python -*-
+"""Skylark macros for MKL.
+if_mkl is a conditional to check if MKL is enabled or not.
+
+mkl_repository is a repository rule for creating MKL repository rule that can
+be pointed to either a local folder, or download it from the internet.
+mkl_repository depends on the following environment variables:
+  * `TF_MKL_ROOT`: The root folder where a copy of libmkl is located.
+"""
+
+
+_TF_MKL_ROOT = "TF_MKL_ROOT"
+
 
 def if_mkl(if_true, if_false = []):
     """Shorthand for select()'ing on whether we're building with MKL.
@@ -11,3 +23,46 @@ def if_mkl(if_true, if_false = []):
         "//third_party/mkl:using_mkl": if_true,
         "//conditions:default": if_false
     })
+
+
+def _enable_local_mkl(repository_ctx):
+  return _TF_MKL_ROOT in repository_ctx.os.environ
+
+
+def _mkl_autoconf_impl(repository_ctx):
+  """Implementation of the local_mkl_autoconf repository rule."""
+
+  if _enable_local_mkl(repository_ctx):
+    # Symlink lib and include local folders.
+    mkl_root = repository_ctx.os.environ[_TF_MKL_ROOT]
+    mkl_lib_path = "%s/lib" % mkl_root
+    repository_ctx.symlink(mkl_lib_path, "lib")
+    mkl_include_path = "%s/include" % mkl_root
+    repository_ctx.symlink(mkl_include_path, "include")
+    mkl_license_path = "%s/license.txt" % mkl_root
+    repository_ctx.symlink(mkl_license_path, "license.txt")
+  else:
+    # setup remote mkl repository.
+    repository_ctx.download_and_extract(
+        repository_ctx.attr.urls,
+        sha256=repository_ctx.attr.sha256,
+        stripPrefix=repository_ctx.attr.strip_prefix,
+    )
+
+  # Also setup BUILD file.
+  repository_ctx.symlink(repository_ctx.attr.build_file, "BUILD")
+
+
+mkl_repository = repository_rule(
+    implementation = _mkl_autoconf_impl,
+    environ = [
+        _TF_MKL_ROOT,
+    ],
+    attrs = {
+        "build_file": attr.label(),
+        "repository": attr.string(),
+        "urls": attr.string_list(default = []),
+        "sha256": attr.string(default = ""),
+        "strip_prefix": attr.string(default = ""),
+    },
+)
diff --git a/third_party/mkl/mkl.BUILD b/third_party/mkl/mkl.BUILD
index 37daa52419..8db97232e1 100644
--- a/third_party/mkl/mkl.BUILD
+++ b/third_party/mkl/mkl.BUILD
@@ -1,6 +1,14 @@
 licenses(["notice"])  # 3-Clause BSD
 
-exports_files(["LICENSE"])
+exports_files(["license.txt"])
+
+filegroup(
+    name = "LICENSE",
+    srcs = [
+        "license.txt",
+    ],
+    visibility = ["//visibility:public"],
+)
 
 cc_library(
     name = "mkl_headers",