diff options
Diffstat (limited to 'third_party/mkl/build_defs.bzl')
-rw-r--r-- | third_party/mkl/build_defs.bzl | 83 |
1 files changed, 74 insertions, 9 deletions
diff --git a/third_party/mkl/build_defs.bzl b/third_party/mkl/build_defs.bzl index 53e02769da..06a8c3518c 100644 --- a/third_party/mkl/build_defs.bzl +++ b/third_party/mkl/build_defs.bzl @@ -1,6 +1,9 @@ # -*- Python -*- """Skylark macros for MKL. if_mkl is a conditional to check if MKL is enabled or not. +if_mkl_ml is a conditional to check if MKL-ML is enabled. +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 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. @@ -15,27 +18,89 @@ _TF_MKL_ROOT = "TF_MKL_ROOT" def if_mkl(if_true, if_false = []): """Shorthand for select()'ing on whether we're building with MKL. - Returns a select statement which evaluates to if_true if we're building - with MKL enabled. Otherwise, the select statement evaluates to if_false. + 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({ - str(Label("//third_party/mkl:using_mkl")): if_true, - "//conditions:default": if_false + "//third_party/mkl:using_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({ + "//third_party/mkl_dnn:using_mkl_dnn_only": + if_false, + "//third_party/mkl:using_mkl": if_true, + "//conditions:default": if_false, + }) + +def if_mkl_ml_only(if_true, if_false = []): + """Shorthand for select()'ing on whether we're building with MKL-ML only. + + Args: + if_true: expression to evaluate if building with MKL-ML only. + if_false: expression to evaluate if building without MKL, or with MKL-DNN. + + Returns: + a select evaluating to either if_true or if_false as appropriate. + """ + return select({ + "//third_party/mkl:using_mkl_ml_only": if_true, + "//conditions:default": if_false, }) def if_mkl_lnx_x64(if_true, if_false = []): - """Shorthand for select()'ing on whether we're building with MKL. + """Shorthand to select() on if MKL is on and the target is Linux x86-64. - Returns a select statement which evaluates to if_true if we're building - with MKL enabled. Otherwise, the select statement evaluates to if_false. + 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({ - str(Label("//third_party/mkl:using_mkl_lnx_x64")): if_true, - "//conditions:default": if_false + "//third_party/mkl:using_mkl_lnx_x64": 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({ + "//third_party/mkl_dnn:using_mkl_dnn_only": + ["@mkl_dnn"], + "//third_party/mkl:using_mkl_ml_only": + ["//third_party/mkl:intel_binary_blob"], + "//third_party/mkl:using_mkl": + [ + "//third_party/mkl:intel_binary_blob", + "@mkl_dnn" + ], + "//conditions:default": [] + }) def _enable_local_mkl(repository_ctx): return _TF_MKL_ROOT in repository_ctx.os.environ |