1"""Skylark macros for MKL.
2
3if_mkl is a conditional to check if we are building with MKL.
4if_mkl_ml is a conditional to check if we are building with MKL-ML.
5if_mkl_ml_only is a conditional to check for MKL-ML-only (no MKL-DNN) mode.
6if_mkl_lnx_x64 is a conditional to check for MKL
7if_enable_mkl is a conditional to check if building with MKL and MKL is enabled.
8
9mkl_repository is a repository rule for creating MKL repository rule that can
10be pointed to either a local folder, or download it from the internet.
11mkl_repository depends on the following environment variables:
12  * `TF_MKL_ROOT`: The root folder where a copy of libmkl is located.
13"""
14
15_TF_MKL_ROOT = "TF_MKL_ROOT"
16
17def if_mkl(if_true, if_false = []):
18    """Shorthand for select()'ing on whether we're building with MKL.
19
20    Args:
21      if_true: expression to evaluate if building with MKL.
22      if_false: expression to evaluate if building without MKL.
23
24    Returns:
25      a select evaluating to either if_true or if_false as appropriate.
26    """
27    return select({
28        "@org_tensorflow//third_party/mkl:build_with_mkl": if_true,
29        "//conditions:default": if_false,
30    })
31
32def if_mkl_ml(if_true, if_false = []):
33    """Shorthand for select()'ing on whether we're building with MKL-ML.
34
35    Args:
36      if_true: expression to evaluate if building with MKL-ML.
37      if_false: expression to evaluate if building without MKL-ML
38        (i.e. without MKL at all, or with MKL-DNN only).
39
40    Returns:
41      a select evaluating to either if_true or if_false as appropriate.
42    """
43    return select({
44        "@org_tensorflow//third_party/mkl_dnn:build_with_mkl_opensource": if_false,
45        "@org_tensorflow//third_party/mkl:build_with_mkl": if_true,
46        "//conditions:default": if_false,
47    })
48
49def if_mkl_lnx_x64(if_true, if_false = []):
50    """Shorthand to select() if building with MKL and the target is Linux x86-64.
51
52    Args:
53      if_true: expression to evaluate if building with MKL is enabled and the
54        target platform is Linux x86-64.
55      if_false: expression to evaluate if building without MKL or for a
56        different platform.
57
58    Returns:
59      a select evaluating to either if_true or if_false as appropriate.
60    """
61    return select({
62        "@org_tensorflow//third_party/mkl:build_with_mkl_lnx_x64": if_true,
63        "//conditions:default": if_false,
64    })
65
66def if_enable_mkl(if_true, if_false = []):
67    """Shorthand to select() if we are building with MKL and MKL is enabled.
68
69    This is only effective when built with MKL.
70
71    Args:
72      if_true: expression to evaluate if building with MKL and MKL is enabled
73      if_false: expression to evaluate if building without MKL or MKL is not enabled.
74
75    Returns:
76      A select evaluating to either if_true or if_false as appropriate.
77    """
78    return select({
79        "@org_tensorflow//third_party/mkl:enable_mkl": if_true,
80        "//conditions:default": if_false,
81    })
82
83def mkl_deps():
84    """Shorthand for select() to pull in the correct set of MKL library deps.
85
86    Can pull in MKL-ML, MKL-DNN, both, or neither depending on config settings.
87
88    Returns:
89      a select evaluating to a list of library dependencies, suitable for
90      inclusion in the deps attribute of rules.
91    """
92    return select({
93        "@org_tensorflow//third_party/mkl:build_with_mkl": ["@mkl_dnn_v1//:mkl_dnn"],
94        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": ["@mkl_dnn_v1//:mkl_dnn_aarch64"],
95        "//conditions:default": [],
96    })
97
98def _enable_local_mkl(repository_ctx):
99    return _TF_MKL_ROOT in repository_ctx.os.environ
100
101def _mkl_autoconf_impl(repository_ctx):
102    """Implementation of the local_mkl_autoconf repository rule."""
103
104    if _enable_local_mkl(repository_ctx):
105        # Symlink lib and include local folders.
106        mkl_root = repository_ctx.os.environ[_TF_MKL_ROOT]
107        mkl_lib_path = "%s/lib" % mkl_root
108        repository_ctx.symlink(mkl_lib_path, "lib")
109        mkl_include_path = "%s/include" % mkl_root
110        repository_ctx.symlink(mkl_include_path, "include")
111        mkl_license_path = "%s/license.txt" % mkl_root
112        repository_ctx.symlink(mkl_license_path, "license.txt")
113    else:
114        # setup remote mkl repository.
115        repository_ctx.download_and_extract(
116            repository_ctx.attr.urls,
117            sha256 = repository_ctx.attr.sha256,
118            stripPrefix = repository_ctx.attr.strip_prefix,
119        )
120
121    # Also setup BUILD file.
122    repository_ctx.symlink(repository_ctx.attr.build_file, "BUILD")
123
124mkl_repository = repository_rule(
125    implementation = _mkl_autoconf_impl,
126    environ = [
127        _TF_MKL_ROOT,
128    ],
129    attrs = {
130        "build_file": attr.label(),
131        "urls": attr.string_list(default = []),
132        "sha256": attr.string(default = ""),
133        "strip_prefix": attr.string(default = ""),
134    },
135)
136