diff options
Diffstat (limited to 'tensorflow/tensorflow.bzl')
| -rw-r--r-- | tensorflow/tensorflow.bzl | 78 |
1 files changed, 49 insertions, 29 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index f51a628ca3..cad5de1b0c 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -17,18 +17,20 @@ load( ) load( "@local_config_cuda//cuda:build_defs.bzl", + "cuda_default_copts", "if_cuda", "if_cuda_is_configured", - "cuda_default_copts", ) load( "@local_config_rocm//rocm:build_defs.bzl", "if_rocm", "if_rocm_is_configured", + "rocm_copts", "rocm_default_copts", ) load( "//third_party/mkl:build_defs.bzl", + "if_enable_mkl", "if_mkl", "if_mkl_lnx_x64", "if_mkl_ml", @@ -45,6 +47,8 @@ load( def register_extension_info(**kwargs): pass +# if_cuda_is_configured def placeholder + # Given a source file, generate a test name. # i.e. "common_runtime/direct_session_test.cc" becomes # "common_runtime_direct_session_test" @@ -244,6 +248,7 @@ def tf_copts(android_optimization_level_override = "-O2", is_external = False): if_tensorrt(["-DGOOGLE_TENSORRT=1"]) + if_mkl(["-DINTEL_MKL=1", "-DEIGEN_USE_VML"]) + if_mkl_open_source_only(["-DINTEL_MKL_DNN_ONLY"]) + + if_enable_mkl(["-DENABLE_MKL"]) + if_ngraph(["-DINTEL_NGRAPH=1"]) + if_mkl_lnx_x64(["-fopenmp"]) + if_android_arm(["-mfpu=neon"]) + @@ -455,7 +460,7 @@ def tf_gen_op_wrapper_cc( tf_cc_binary( name = tool, copts = tf_copts(), - linkopts = if_not_windows(["-lm"]), + linkopts = if_not_windows(["-lm", "-Wl,-ldl"]), linkstatic = 1, # Faster to link this one-time-use binary dynamically deps = [op_gen] + deps, ) @@ -609,6 +614,7 @@ def tf_gen_op_wrappers_cc( # is invalid to specify both "hidden" and "op_whitelist". # cc_linkopts: Optional linkopts to be added to tf_cc_binary that contains the # specified ops. + def tf_gen_op_wrapper_py( name, out = None, @@ -630,7 +636,7 @@ def tf_gen_op_wrapper_py( deps = [str(Label("//tensorflow/core:" + name + "_op_lib"))] tf_cc_binary( name = tool_name, - linkopts = if_not_windows(["-lm"]) + cc_linkopts, + linkopts = if_not_windows(["-lm", "-Wl,-ldl"]) + cc_linkopts, copts = tf_copts(), linkstatic = 1, # Faster to link this one-time-use binary dynamically deps = ([ @@ -867,14 +873,16 @@ def tf_cuda_only_cc_test( srcs = srcs + tf_binary_additional_srcs(), size = size, args = args, - copts = _cuda_copts() + _rocm_copts() + tf_copts(), + copts = _cuda_copts() + rocm_copts() + tf_copts(), data = data + tf_binary_dynamic_kernel_dsos(kernels), deps = deps + tf_binary_dynamic_kernel_deps(kernels) + - if_cuda_is_configured([ - clean_dep("//tensorflow/core:cuda"), - clean_dep("//tensorflow/core:gpu_lib")]) + - if_rocm_is_configured([ - clean_dep("//tensorflow/core:gpu_lib")]), + if_cuda_is_configured([ + clean_dep("//tensorflow/core:cuda"), + clean_dep("//tensorflow/core:gpu_lib"), + ]) + + if_rocm_is_configured([ + clean_dep("//tensorflow/core:gpu_lib"), + ]), linkopts = if_not_windows(["-lpthread", "-lm"]) + linkopts + _rpath_linkopts(name), linkstatic = linkstatic or select({ # cc_tests with ".so"s in srcs incorrectly link on Darwin @@ -1027,20 +1035,6 @@ def _cuda_copts(opts = []): ]), }) + if_cuda_is_configured(opts) -def _rocm_copts(opts = []): - """Gets the appropriate set of copts for (maybe) ROCm compilation. - - If we're doing ROCm compilation, returns copts for our particular ROCm - compiler. If we're not doing ROCm compilation, returns an empty list. - - """ - return rocm_default_copts() + select({ - "//conditions:default": [], - "@local_config_rocm//rocm:using_hipcc": ([ - "", - ]) - }) + if_rocm_is_configured(opts) - # Build defs for TensorFlow kernels # When this target is built using --config=cuda, a cc_library is built @@ -1057,7 +1051,7 @@ def tf_gpu_kernel_library( deps = [], hdrs = [], **kwargs): - copts = copts + tf_copts() + _cuda_copts(opts = cuda_copts) + _rocm_copts(opts = cuda_copts) + copts = copts + tf_copts() + _cuda_copts(opts = cuda_copts) + rocm_copts(opts = cuda_copts) kwargs["features"] = kwargs.get("features", []) + ["-use_header_modules"] native.cc_library( @@ -1104,14 +1098,15 @@ def tf_cuda_library(deps = None, cuda_deps = None, copts = tf_copts(), **kwargs) kwargs["features"] = kwargs.get("features", []) + ["-use_header_modules"] native.cc_library( - deps = deps + if_cuda_is_configured(cuda_deps + [ + deps = deps + if_cuda(cuda_deps + [ clean_dep("//tensorflow/core:cuda"), - "@local_config_cuda//cuda:cuda_headers" + "@local_config_cuda//cuda:cuda_headers", ]) + if_rocm_is_configured(cuda_deps + [ - "@local_config_rocm//rocm:rocm_headers" + # rocm_header placeholder ]), copts = (copts + if_cuda(["-DGOOGLE_CUDA=1"]) + if_rocm(["-DTENSORFLOW_USE_ROCM=1"]) + if_mkl(["-DINTEL_MKL=1"]) + if_mkl_open_source_only(["-DINTEL_MKL_DNN_ONLY"]) + + if_enable_mkl(["-DENABLE_MKL"]) + if_tensorrt(["-DGOOGLE_TENSORRT=1"])), **kwargs ) @@ -1246,9 +1241,11 @@ def tf_mkl_kernel_library( if prefix: srcs = srcs + native.glob( [prefix + "*.cc"], + exclude = [prefix + "*test*"], ) hdrs = hdrs + native.glob( [prefix + "*.h"], + exclude = [prefix + "*test*"], ) # -fno-exceptions in nocopts breaks compilation if header modules are enabled. @@ -1501,7 +1498,7 @@ def tf_custom_op_library(name, srcs = [], gpu_srcs = [], deps = [], linkopts = [ srcs = gpu_srcs, copts = _cuda_copts() + if_tensorrt(["-DGOOGLE_TENSORRT=1"]), features = if_cuda(["-use_header_modules"]), - deps = deps + if_cuda_is_configured(cuda_deps) + if_rocm_is_configured(rocm_deps) + deps = deps + if_cuda_is_configured(cuda_deps) + if_rocm_is_configured(rocm_deps), ) cuda_deps.extend([":" + basename + "_gpu"]) rocm_deps.extend([":" + basename + "_gpu"]) @@ -1709,7 +1706,7 @@ def py_test(deps = [], data = [], kernels = [], **kwargs): deps = select({ "//conditions:default": deps, clean_dep("//tensorflow:no_tensorflow_py_deps"): [], - }) + tf_binary_dynamic_kernel_deps(kernels), + }), data = data + select({ "//conditions:default": [], clean_dep("//tensorflow:no_tensorflow_py_deps"): ["//tensorflow/tools/pip_package:win_pip_package_marker"], @@ -1722,6 +1719,29 @@ register_extension_info( label_regex_for_dep = "{extension_name}", ) +# Similar to py_test above, this macro is used to exclude dependencies for some py_binary +# targets in order to reduce the size of //tensorflow/tools/pip_package:simple_console_windows. +# See https://github.com/tensorflow/tensorflow/issues/22390 +def py_binary(name, deps = [], **kwargs): + # Add an extra target for dependencies to avoid nested select statement. + native.py_library( + name = name + "_deps", + deps = deps, + ) + native.py_binary( + name = name, + deps = select({ + "//conditions:default": [":" + name + "_deps"], + clean_dep("//tensorflow:no_tensorflow_py_deps"): [], + }), + **kwargs + ) + +register_extension_info( + extension_name = "py_binary", + label_regex_for_dep = "{extension_name}", +) + def tf_py_test( name, srcs, |