diff options
Diffstat (limited to 'tensorflow/tensorflow.bzl')
| -rw-r--r-- | tensorflow/tensorflow.bzl | 49 |
1 files changed, 39 insertions, 10 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 741145216c..16f7b217b4 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -31,6 +31,10 @@ load( "//third_party/mkl_dnn:build_defs.bzl", "if_mkl_open_source_only", ) +load( + "//third_party/ngraph:build_defs.bzl", + "if_ngraph", +) def register_extension_info(**kwargs): pass @@ -233,6 +237,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_ngraph(["-DINTEL_NGRAPH=1"]) + if_mkl_lnx_x64(["-fopenmp"]) + if_android_arm(["-mfpu=neon"]) + if_linux_x86_64(["-msse3"]) + @@ -324,11 +329,16 @@ def tf_binary_additional_srcs(): ], ) +def _linux_kernel_dso_name(kernel_build_target): + """Given a build target, construct the dso name for linux.""" + parts = kernel_build_target.split(":") + return "%s:libtfkernel_%s.so" % (parts[0], parts[1]) + # Helper functions to add kernel dependencies to tf binaries when using dynamic # kernel linking. def tf_binary_dynamic_kernel_dsos(kernels): return if_dynamic_kernels( - extra_deps = ["libtfkernel_%s.so" % clean_dep(k) for k in kernels], + extra_deps = [_linux_kernel_dso_name(k) for k in kernels], otherwise = [], ) @@ -391,7 +401,7 @@ def tf_cc_binary( srcs = srcs + tf_binary_additional_srcs(), deps = deps + tf_binary_dynamic_kernel_deps(kernels) + if_mkl_ml( [ - "//third_party/mkl:intel_binary_blob", + clean_dep("//third_party/mkl:intel_binary_blob"), ], ), data = data + tf_binary_dynamic_kernel_dsos(kernels), @@ -438,7 +448,7 @@ def tf_gen_op_wrapper_cc( tf_cc_binary( name = tool, copts = tf_copts(), - linkopts = if_not_windows(["-lm","-Wl,-ldl"]), + linkopts = if_not_windows(["-lm", "-Wl,-ldl"]), linkstatic = 1, # Faster to link this one-time-use binary dynamically deps = [op_gen] + deps, ) @@ -614,7 +624,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","-Wl,-ldl"]) + 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 = ([ @@ -730,7 +740,7 @@ def tf_cc_test( }) + linkopts + _rpath_linkopts(name), deps = deps + tf_binary_dynamic_kernel_deps(kernels) + if_mkl_ml( [ - "//third_party/mkl:intel_binary_blob", + clean_dep("//third_party/mkl:intel_binary_blob"), ], ), data = data + tf_binary_dynamic_kernel_dsos(kernels), @@ -792,6 +802,7 @@ def tf_cuda_cc_test( extra_copts = [], linkstatic = 0, args = [], + kernels = [], linkopts = []): tf_cc_test( name = name, @@ -804,6 +815,7 @@ def tf_cuda_cc_test( linkstatic = linkstatic, linkopts = linkopts, args = args, + kernels = kernels, ) tf_cc_test( name = name, @@ -825,6 +837,7 @@ def tf_cuda_cc_test( extra_copts = extra_copts, linkopts = linkopts, args = args, + kernels = kernels, ) register_extension_info( @@ -880,6 +893,7 @@ def tf_cc_tests( size = "medium", args = None, linkopts = [], + kernels = [], nocopts = None): for src in srcs: tf_cc_test( @@ -892,6 +906,7 @@ def tf_cc_tests( args = args, linkopts = linkopts, nocopts = nocopts, + kernels = kernels, ) def tf_cc_test_mkl( @@ -939,8 +954,9 @@ def tf_cc_tests_gpu( linkstatic = 0, tags = [], size = "medium", + kernels = [], args = None): - tf_cc_tests(srcs, deps, linkstatic, tags = tags, size = size, args = args) + tf_cc_tests(srcs, deps, linkstatic, tags = tags, size = size, kernels = kernels, args = args) def tf_cuda_cc_tests( srcs, @@ -950,6 +966,7 @@ def tf_cuda_cc_tests( size = "medium", linkstatic = 0, args = None, + kernels = [], linkopts = []): for src in srcs: tf_cuda_cc_test( @@ -960,6 +977,7 @@ def tf_cuda_cc_tests( size = size, linkstatic = linkstatic, args = args, + kernels = kernels, linkopts = linkopts, ) @@ -1348,12 +1366,13 @@ def transitive_hdrs(name, deps = [], **kwargs): # Create a header only library that includes all the headers exported by # the libraries in deps. -def cc_header_only_library(name, deps = [], includes = [], **kwargs): +def cc_header_only_library(name, deps = [], includes = [], extra_deps = [], **kwargs): _transitive_hdrs(name = name + "_gather", deps = deps) native.cc_library( name = name, hdrs = [":" + name + "_gather"], includes = includes, + deps = extra_deps, **kwargs ) @@ -1650,17 +1669,17 @@ def tf_py_wrap_cc( # Note that this only works on Windows. See the definition of # //third_party/tensorflow/tools/pip_package:win_pip_package_marker for specific reasons. # 2. When --define=no_tensorflow_py_deps=false (by default), it's a normal py_test. -def py_test(deps = [], data = [], **kwargs): +def py_test(deps = [], data = [], kernels = [], **kwargs): native.py_test( # TODO(jlebar): Ideally we'd use tcmalloc here., 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"], - }), + }) + tf_binary_dynamic_kernel_dsos(kernels), **kwargs ) @@ -1679,6 +1698,7 @@ def tf_py_test( tags = [], shard_count = 1, additional_deps = [], + kernels = [], flaky = 0, xla_enabled = False, grpc_enabled = False): @@ -1695,6 +1715,7 @@ def tf_py_test( tags = tags, visibility = [clean_dep("//tensorflow:internal")], shard_count = shard_count, + kernels = kernels, data = data, deps = [ clean_dep("//tensorflow/python:extra_py_tests_deps"), @@ -1718,6 +1739,7 @@ def cuda_py_test( args = [], shard_count = 1, additional_deps = [], + kernels = [], tags = [], flaky = 0, xla_enabled = False, @@ -1733,6 +1755,7 @@ def cuda_py_test( tags = test_tags, shard_count = shard_count, additional_deps = additional_deps, + kernels = kernels, flaky = flaky, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, @@ -1752,6 +1775,7 @@ def sycl_py_test( args = [], shard_count = 1, additional_deps = [], + kernels = [], tags = [], flaky = 0, xla_enabled = False, @@ -1767,6 +1791,7 @@ def sycl_py_test( tags = test_tags, shard_count = shard_count, additional_deps = additional_deps, + kernels = kernels, flaky = flaky, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, @@ -1782,6 +1807,7 @@ def py_tests( srcs, size = "medium", additional_deps = [], + kernels = [], data = [], tags = [], shard_count = 1, @@ -1801,6 +1827,7 @@ def py_tests( shard_count = shard_count, data = data, additional_deps = additional_deps, + kernels = kernels, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, ) @@ -1810,6 +1837,7 @@ def cuda_py_tests( srcs, size = "medium", additional_deps = [], + kernels = [], data = [], shard_count = 1, tags = [], @@ -1826,6 +1854,7 @@ def cuda_py_tests( tags = test_tags, shard_count = shard_count, prefix = prefix, + kernels = kernels, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, ) |