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 )