xref: /aosp_15_r20/external/tensorflow/third_party/mkl/build_defs.bzl (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1"""Starlark 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 oneDNN.
19
20      OneDNN gets built if we are building on platforms that support oneDNN
21      (x86 linux/windows) or if specifcially configured to use oneDNN.
22
23    Args:
24      if_true: expression to evaluate if building with oneDNN.
25      if_false: expression to evaluate if building without oneDNN.
26
27    Returns:
28      a select evaluating to either if_true or if_false as appropriate.
29
30    TODO(intel-tf):
31      the first "if_true" line is kept because non-x86 platforms (e.g., ARM)
32      may need it. It may be deleted in future with refactoring.
33    """
34    return select({
35        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": if_true,
36        "@org_tensorflow//tensorflow:linux_x86_64": if_true,
37        "@org_tensorflow//tensorflow:windows": if_true,
38        "//conditions:default": if_false,
39    })
40
41def if_mkl_ml(if_true, if_false = []):
42    """Shorthand for select()'ing on whether we're building with MKL-ML.
43
44    Args:
45      if_true: expression to evaluate if building with MKL-ML.
46      if_false: expression to evaluate if building without MKL-ML
47        (i.e. without MKL at all, or with MKL-DNN only).
48
49    Returns:
50      a select evaluating to either if_true or if_false as appropriate.
51    """
52    return select({
53        "@org_tensorflow//third_party/mkl_dnn:build_with_mkl_opensource": if_false,
54        "@org_tensorflow//third_party/mkl:build_with_mkl": if_true,
55        "//conditions:default": if_false,
56    })
57
58def if_mkl_lnx_x64(if_true, if_false = []):
59    """Shorthand to select() if building with MKL and the target is Linux x86-64.
60
61    Args:
62      if_true: expression to evaluate if building with MKL is enabled and the
63        target platform is Linux x86-64.
64      if_false: expression to evaluate if building without MKL or for a
65        different platform.
66
67    Returns:
68      a select evaluating to either if_true or if_false as appropriate.
69    """
70    return select({
71        "@org_tensorflow//third_party/mkl:build_with_mkl_lnx_x64": if_true,
72        "//conditions:default": if_false,
73    })
74
75def if_enable_mkl(if_true, if_false = []):
76    """Shorthand to select() if we are building with MKL and MKL is enabled.
77
78    This is only effective when built with MKL.
79
80    Args:
81      if_true: expression to evaluate if building with MKL and MKL is enabled
82      if_false: expression to evaluate if building without MKL or MKL is not enabled.
83
84    Returns:
85      A select evaluating to either if_true or if_false as appropriate.
86    """
87    return select({
88        "@org_tensorflow//third_party/mkl:enable_mkl": if_true,
89        "//conditions:default": if_false,
90    })
91
92def mkl_deps():
93    """Returns the correct set of oneDNN library dependencies.
94
95      Shorthand for select() to pull in the correct set of oneDNN library deps
96      depending on the platform. x86 Linux/Windows with or without --config=mkl
97      will always build with oneDNN library.
98
99    Returns:
100      a select evaluating to a list of library dependencies, suitable for
101      inclusion in the deps attribute of rules.
102    """
103    return select({
104        "@org_tensorflow//third_party/mkl:build_with_mkl_aarch64": ["@mkl_dnn_acl_compatible//:mkl_dnn_acl"],
105        "@org_tensorflow//tensorflow:linux_x86_64": ["@mkl_dnn_v1//:mkl_dnn"],
106        "@org_tensorflow//tensorflow:windows": ["@mkl_dnn_v1//:mkl_dnn"],
107        "//conditions:default": [],
108    })
109
110def _enable_local_mkl(repository_ctx):
111    return _TF_MKL_ROOT in repository_ctx.os.environ
112
113def _mkl_autoconf_impl(repository_ctx):
114    """Implementation of the local_mkl_autoconf repository rule."""
115
116    if _enable_local_mkl(repository_ctx):
117        # Symlink lib and include local folders.
118        mkl_root = repository_ctx.os.environ[_TF_MKL_ROOT]
119        mkl_lib_path = "%s/lib" % mkl_root
120        repository_ctx.symlink(mkl_lib_path, "lib")
121        mkl_include_path = "%s/include" % mkl_root
122        repository_ctx.symlink(mkl_include_path, "include")
123        mkl_license_path = "%s/license.txt" % mkl_root
124        repository_ctx.symlink(mkl_license_path, "license.txt")
125    else:
126        # setup remote mkl repository.
127        repository_ctx.download_and_extract(
128            repository_ctx.attr.urls,
129            sha256 = repository_ctx.attr.sha256,
130            stripPrefix = repository_ctx.attr.strip_prefix,
131        )
132
133    # Also setup BUILD file.
134    repository_ctx.symlink(repository_ctx.attr.build_file, "BUILD")
135
136mkl_repository = repository_rule(
137    implementation = _mkl_autoconf_impl,
138    environ = [
139        _TF_MKL_ROOT,
140    ],
141    attrs = {
142        "build_file": attr.label(),
143        "urls": attr.string_list(default = []),
144        "sha256": attr.string(default = ""),
145        "strip_prefix": attr.string(default = ""),
146    },
147)
148