github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/mkl/build_defs.bzl (about)

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