Add the kernels option to also the python binary macros defined in tensorflow.bzl

PiperOrigin-RevId: 210859798

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,
     )