diff options
Diffstat (limited to 'third_party/mkl/build_defs.bzl')
-rw-r--r-- | third_party/mkl/build_defs.bzl | 57 |
1 files changed, 56 insertions, 1 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 = ""), + }, +) |