diff options
Diffstat (limited to 'third_party/mkl')
| -rw-r--r-- | third_party/mkl/build_defs.bzl | 52 |
1 files changed, 24 insertions, 28 deletions
diff --git a/third_party/mkl/build_defs.bzl b/third_party/mkl/build_defs.bzl index 9970a772fe..b645c0fc5c 100644 --- a/third_party/mkl/build_defs.bzl +++ b/third_party/mkl/build_defs.bzl @@ -11,10 +11,8 @@ 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. @@ -45,7 +43,7 @@ def if_mkl_ml(if_true, if_false = []): str(Label("//third_party/mkl_dnn:using_mkl_dnn_only")): if_false, str(Label("//third_party/mkl:using_mkl")): if_true, "//conditions:default": if_false, - }) + }) def if_mkl_ml_only(if_true, if_false = []): """Shorthand for select()'ing on whether we're building with MKL-ML only. @@ -96,35 +94,33 @@ def mkl_deps(): "@mkl_dnn", ], "//conditions:default": [], - }) + }) def _enable_local_mkl(repository_ctx): - return _TF_MKL_ROOT in repository_ctx.os.environ - + 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") - + """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, |