diff options
author | Gunhan Gulsoy <gunan@google.com> | 2018-08-30 00:09:05 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-30 00:12:48 -0700 |
commit | d004f08ee6102b2081a4ae14420e9eb76b1eb669 (patch) | |
tree | 1ed4cca466d111d975ba0f3dac0d219838ae67c5 /tensorflow | |
parent | c74d0aeecac3f969f07d67763de6a48755db1844 (diff) |
Add the kernels option to also the python binary macros defined in tensorflow.bzl
PiperOrigin-RevId: 210859798
Diffstat (limited to 'tensorflow')
-rw-r--r-- | tensorflow/tensorflow.bzl | 33 |
1 files changed, 28 insertions, 5 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 3562a5192d..b9f26898c4 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -329,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 = [], ) @@ -796,6 +801,7 @@ def tf_cuda_cc_test( extra_copts = [], linkstatic = 0, args = [], + kernels = [], linkopts = []): tf_cc_test( name = name, @@ -808,6 +814,7 @@ def tf_cuda_cc_test( linkstatic = linkstatic, linkopts = linkopts, args = args, + kernels = kernels, ) tf_cc_test( name = name, @@ -829,6 +836,7 @@ def tf_cuda_cc_test( extra_copts = extra_copts, linkopts = linkopts, args = args, + kernels = kernels, ) register_extension_info( @@ -884,6 +892,7 @@ def tf_cc_tests( size = "medium", args = None, linkopts = [], + kernels = [], nocopts = None): for src in srcs: tf_cc_test( @@ -896,6 +905,7 @@ def tf_cc_tests( args = args, linkopts = linkopts, nocopts = nocopts, + kernels = kernels, ) def tf_cc_test_mkl( @@ -943,8 +953,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, @@ -954,6 +965,7 @@ def tf_cuda_cc_tests( size = "medium", linkstatic = 0, args = None, + kernels = [], linkopts = []): for src in srcs: tf_cuda_cc_test( @@ -964,6 +976,7 @@ def tf_cuda_cc_tests( size = size, linkstatic = linkstatic, args = args, + kernels = kernels, linkopts = linkopts, ) @@ -1654,17 +1667,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 ) @@ -1683,6 +1696,7 @@ def tf_py_test( tags = [], shard_count = 1, additional_deps = [], + kernels = [], flaky = 0, xla_enabled = False, grpc_enabled = False): @@ -1699,6 +1713,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"), @@ -1722,6 +1737,7 @@ def cuda_py_test( args = [], shard_count = 1, additional_deps = [], + kernels = [], tags = [], flaky = 0, xla_enabled = False, @@ -1737,6 +1753,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, @@ -1756,6 +1773,7 @@ def sycl_py_test( args = [], shard_count = 1, additional_deps = [], + kernels = [], tags = [], flaky = 0, xla_enabled = False, @@ -1771,6 +1789,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, @@ -1786,6 +1805,7 @@ def py_tests( srcs, size = "medium", additional_deps = [], + kernels = [], data = [], tags = [], shard_count = 1, @@ -1805,6 +1825,7 @@ def py_tests( shard_count = shard_count, data = data, additional_deps = additional_deps, + kernels = kernels, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, ) @@ -1814,6 +1835,7 @@ def cuda_py_tests( srcs, size = "medium", additional_deps = [], + kernels = [], data = [], shard_count = 1, tags = [], @@ -1830,6 +1852,7 @@ def cuda_py_tests( tags = test_tags, shard_count = shard_count, prefix = prefix, + kernels = kernels, xla_enabled = xla_enabled, grpc_enabled = grpc_enabled, ) |