You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
136 lines
4.8 KiB
136 lines
4.8 KiB
"""Skylark macros for MKL.
|
|
|
|
if_mkl is a conditional to check if we are building with MKL.
|
|
if_mkl_ml is a conditional to check if we are building with MKL-ML.
|
|
if_mkl_ml_only is a conditional to check for MKL-ML-only (no MKL-DNN) mode.
|
|
if_mkl_lnx_x64 is a conditional to check for MKL
|
|
if_enable_mkl is a conditional to check if building with MKL and MKL is enabled.
|
|
|
|
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.
|
|
|
|
Args:
|
|
if_true: expression to evaluate if building with MKL.
|
|
if_false: expression to evaluate if building without MKL.
|
|
|
|
Returns:
|
|
a select evaluating to either if_true or if_false as appropriate.
|
|
"""
|
|
return select({
|
|
"@org_tensorflow//third_party/mkl:build_with_mkl": if_true,
|
|
"//conditions:default": if_false,
|
|
})
|
|
|
|
def if_mkl_ml(if_true, if_false = []):
|
|
"""Shorthand for select()'ing on whether we're building with MKL-ML.
|
|
|
|
Args:
|
|
if_true: expression to evaluate if building with MKL-ML.
|
|
if_false: expression to evaluate if building without MKL-ML
|
|
(i.e. without MKL at all, or with MKL-DNN only).
|
|
|
|
Returns:
|
|
a select evaluating to either if_true or if_false as appropriate.
|
|
"""
|
|
return select({
|
|
"@org_tensorflow//third_party/mkl_dnn:build_with_mkl_opensource": if_false,
|
|
"@org_tensorflow//third_party/mkl:build_with_mkl": if_true,
|
|
"//conditions:default": if_false,
|
|
})
|
|
|
|
def if_mkl_lnx_x64(if_true, if_false = []):
|
|
"""Shorthand to select() if building with MKL and the target is Linux x86-64.
|
|
|
|
Args:
|
|
if_true: expression to evaluate if building with MKL is enabled and the
|
|
target platform is Linux x86-64.
|
|
if_false: expression to evaluate if building without MKL or for a
|
|
different platform.
|
|
|
|
Returns:
|
|
a select evaluating to either if_true or if_false as appropriate.
|
|
"""
|
|
return select({
|
|
"@org_tensorflow//third_party/mkl:build_with_mkl_lnx_x64": if_true,
|
|
"//conditions:default": if_false,
|
|
})
|
|
|
|
def if_enable_mkl(if_true, if_false = []):
|
|
"""Shorthand to select() if we are building with MKL and MKL is enabled.
|
|
|
|
This is only effective when built with MKL.
|
|
|
|
Args:
|
|
if_true: expression to evaluate if building with MKL and MKL is enabled
|
|
if_false: expression to evaluate if building without MKL or MKL is not enabled.
|
|
|
|
Returns:
|
|
A select evaluating to either if_true or if_false as appropriate.
|
|
"""
|
|
return select({
|
|
"@org_tensorflow//third_party/mkl:enable_mkl": if_true,
|
|
"//conditions:default": if_false,
|
|
})
|
|
|
|
def mkl_deps():
|
|
"""Shorthand for select() to pull in the correct set of MKL library deps.
|
|
|
|
Can pull in MKL-ML, MKL-DNN, both, or neither depending on config settings.
|
|
|
|
Returns:
|
|
a select evaluating to a list of library dependencies, suitable for
|
|
inclusion in the deps attribute of rules.
|
|
"""
|
|
return select({
|
|
"@org_tensorflow//third_party/mkl:build_with_mkl": ["@mkl_dnn_v1//:mkl_dnn"],
|
|
"@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": ["@mkl_dnn_v1//:mkl_dnn_aarch64"],
|
|
"//conditions:default": [],
|
|
})
|
|
|
|
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(),
|
|
"urls": attr.string_list(default = []),
|
|
"sha256": attr.string(default = ""),
|
|
"strip_prefix": attr.string(default = ""),
|
|
},
|
|
)
|