diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index f4898897435..d68bb928233 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -3,7 +3,7 @@ # ":platform" - Low-level and platform-specific Python code. load("//tensorflow:tensorflow.bzl", "py_strict_library") -load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_tpu", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cc_test", "tf_cuda_library", "tf_gen_op_wrapper_py") +load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_tpu", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cc_test", "tf_cuda_library", "tf_enable_mlir_bridge", "tf_gen_op_wrapper_py") # buildifier: disable=same-origin-load load("//tensorflow:tensorflow.bzl", "tf_monitoring_python_deps") @@ -5670,6 +5670,12 @@ cc_library( ], ) +py_library( + name = "global_test_configuration", + deps = if_mlir(["//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration"]) + + tf_enable_mlir_bridge(), +) + py_library( name = "util", srcs = glob( @@ -5693,6 +5699,10 @@ py_library( ], deps = [ ":_pywrap_tf32_execution", + # global_test_configuration is added here because all major tests depend on this + # library. It isn't possible to add these test dependencies via tensorflow.bzl's + # py_test because not all tensorflow tests use tensorflow.bzl's py_test. + ":global_test_configuration", ":tf_decorator", ":tf_export", ":tf_stack", @@ -5702,7 +5712,7 @@ py_library( "@six_archive//:six", "@wrapt", "//tensorflow/tools/compatibility:all_renames_v2", - ] + if_mlir(["//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration"]), + ], ) tf_py_test( diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index e5c29d26f27..6139fc9372a 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -2145,12 +2145,6 @@ def pywrap_tensorflow_macro( # This macro is for running python tests against system installed pip package # on Windows. # -# This macro can also enable testing with the experimental mlir bridge when -# enable_mlir_bridge is true. When it is enabled tests are run both with and without -# the mlir bridge. Support for enabling the mlir bridge is added here because -# it allows all tensorflow tests to be configured to be run with and without the -# mlir bridge. -# # py_test is built as an executable python zip file on Windows, which contains all # dependencies of the target. Because of the C++ extensions, it would be very # inefficient if the py_test zips all runfiles, plus we don't need them when running @@ -2168,19 +2162,12 @@ def py_test(deps = [], data = [], kernels = [], **kwargs): # Python version placeholder if kwargs.get("python_version", None) == "PY3": kwargs["tags"] = kwargs.get("tags", []) + ["no_oss_py2"] - deps = deps.to_list() if type(deps) == "depset" else deps native.py_test( # TODO(jlebar): Ideally we'd use tcmalloc here., deps = select({ - "//conditions:default": deps, - clean_dep("//tensorflow:no_tensorflow_py_deps"): [], - }) + - select({ - str(Label("//tensorflow:enable_mlir_bridge")): [ - "//tensorflow/python:is_mlir_bridge_test_true", - ], - "//conditions:default": [], - }), + "//conditions:default": deps, + clean_dep("//tensorflow:no_tensorflow_py_deps"): [], + }), data = data + select({ "//conditions:default": kernels, clean_dep("//tensorflow:no_tensorflow_py_deps"): ["//tensorflow/tools/pip_package:win_pip_package_marker"], @@ -2921,6 +2908,14 @@ def if_mlir(if_true, if_false = []): "//conditions:default": if_false, }) +def tf_enable_mlir_bridge(): + return select({ + str(Label("//tensorflow:enable_mlir_bridge")): [ + "//tensorflow/python:is_mlir_bridge_test_true", + ], + "//conditions:default": [], + }) + def if_tpu(if_true, if_false = []): """Shorthand for select()ing whether to build for TPUs.""" return select({ diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py index d06cee1e038..40d2cff56b4 100644 --- a/tensorflow/tools/pip_package/pip_smoke_test.py +++ b/tensorflow/tools/pip_package/pip_smoke_test.py @@ -75,7 +75,6 @@ PYTHON_TARGETS, PY_TEST_QUERY_EXPRESSION = BuildPyTestDependencies() DEPENDENCY_DENYLIST = [ "//tensorflow/python:extra_py_tests_deps", "//tensorflow/cc/saved_model:saved_model_half_plus_two", - "//tensorflow:enable_mlir_bridge", "//tensorflow:no_tensorflow_py_deps", "//tensorflow/tools/pip_package:win_pip_package_marker", "//tensorflow/python:test_ops_2",