• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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        str(Label("//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        str(Label("//third_party/mkl_dnn:build_with_mkl_dnn_only")): if_false,
45        str(Label("//third_party/mkl:build_with_mkl")): if_true,
46        "//conditions:default": if_false,
47    })
48
49def if_mkl_ml_only(if_true, if_false = []):
50    """Shorthand for select()'ing on whether we're building with MKL-ML only.
51
52    Args:
53      if_true: expression to evaluate if building with MKL-ML only.
54      if_false: expression to evaluate if building without MKL, or with MKL-DNN.
55
56    Returns:
57      a select evaluating to either if_true or if_false as appropriate.
58    """
59    return select({
60        str(Label("//third_party/mkl:build_with_mkl_ml_only")): if_true,
61        "//conditions:default": if_false,
62    })
63
64def if_mkl_lnx_x64(if_true, if_false = []):
65    """Shorthand to select() if building with MKL and the target is Linux x86-64.
66
67    Args:
68      if_true: expression to evaluate if building with MKL is enabled and the
69        target platform is Linux x86-64.
70      if_false: expression to evaluate if building without MKL or for a
71        different platform.
72
73    Returns:
74      a select evaluating to either if_true or if_false as appropriate.
75    """
76    return select({
77        str(Label("//third_party/mkl:build_with_mkl_lnx_x64")): if_true,
78        "//conditions:default": if_false,
79    })
80
81def if_enable_mkl(if_true, if_false = []):
82    """Shorthand to select() if we are building with MKL and MKL is enabled.
83
84    This is only effective when built with MKL.
85
86    Args:
87      if_true: expression to evaluate if building with MKL and MKL is enabled
88      if_false: expression to evaluate if building without MKL or MKL is not enabled.
89
90    Returns:
91      A select evaluating to either if_true or if_false as appropriate.
92    """
93    return select({
94        str(Label("//third_party/mkl:enable_mkl")): if_true,
95        "//conditions:default": if_false,
96    })
97
98def mkl_deps():
99    """Shorthand for select() to pull in the correct set of MKL library deps.
100
101    Can pull in MKL-ML, MKL-DNN, both, or neither depending on config settings.
102
103    Returns:
104      a select evaluating to a list of library dependencies, suitable for
105      inclusion in the deps attribute of rules.
106    """
107    return select({
108        str(Label("//third_party/mkl_dnn:build_with_mkl_dnn_only")): ["@mkl_dnn"],
109        str(Label("//third_party/mkl_dnn:build_with_mkl_dnn_v1_only")): ["@mkl_dnn_v1//:mkl_dnn"],
110        str(Label("//third_party/mkl:build_with_mkl_ml_only")): ["//third_party/mkl:intel_binary_blob"],
111        str(Label("//third_party/mkl:build_with_mkl")): [
112            "//third_party/mkl:intel_binary_blob",
113            "@mkl_dnn",
114        ],
115        "//conditions:default": [],
116    })
117
118def _enable_local_mkl(repository_ctx):
119    return _TF_MKL_ROOT in repository_ctx.os.environ
120
121def _mkl_autoconf_impl(repository_ctx):
122    """Implementation of the local_mkl_autoconf repository rule."""
123
124    if _enable_local_mkl(repository_ctx):
125        # Symlink lib and include local folders.
126        mkl_root = repository_ctx.os.environ[_TF_MKL_ROOT]
127        mkl_lib_path = "%s/lib" % mkl_root
128        repository_ctx.symlink(mkl_lib_path, "lib")
129        mkl_include_path = "%s/include" % mkl_root
130        repository_ctx.symlink(mkl_include_path, "include")
131        mkl_license_path = "%s/license.txt" % mkl_root
132        repository_ctx.symlink(mkl_license_path, "license.txt")
133    else:
134        # setup remote mkl repository.
135        repository_ctx.download_and_extract(
136            repository_ctx.attr.urls,
137            sha256 = repository_ctx.attr.sha256,
138            stripPrefix = repository_ctx.attr.strip_prefix,
139        )
140
141    # Also setup BUILD file.
142    repository_ctx.symlink(repository_ctx.attr.build_file, "BUILD")
143
144mkl_repository = repository_rule(
145    implementation = _mkl_autoconf_impl,
146    environ = [
147        _TF_MKL_ROOT,
148    ],
149    attrs = {
150        "build_file": attr.label(),
151        "urls": attr.string_list(default = []),
152        "sha256": attr.string(default = ""),
153        "strip_prefix": attr.string(default = ""),
154    },
155)
156