Add plumbing for setting enable_mlir_bridge=False

In order to roll out the mlir_bridge, we need to support forcibly
enabling the bridge, forcibly disabling the bridge and a default state.
If the bridge is not forciby enabled or disabled, the new bridge will
decide whether or not it should be enabled. Overtime, the new bridge
will increasing select more models to run.

This cl adds front-end support for forcibly disabling the bridge.
For now, if the bride is not forcibly enabled, it is assumed to
be disabled.

PiperOrigin-RevId: 333767292
Change-Id: Icbc89dd6e7a49c98963d157bd030cda8e5b262ed

tensorflow/BUILD

@@ -497,13 +497,20 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
-# This flag enables experimental MLIR bridge support.
+# This flag forcibly enables experimental MLIR bridge support.
 config_setting(
     name = "enable_mlir_bridge",
     values = {"define": "enable_mlir_bridge=true"},
     visibility = ["//visibility:public"],
 )
 
+# This flag forcibly disables experimental MLIR bridge support.
+config_setting(
+    name = "disable_mlir_bridge",
+    values = {"define": "enable_mlir_bridge=false"},
+    visibility = ["//visibility:public"],
+)
+
 # This flag enables experimental TPU support
 config_setting(
     name = "with_tpu_support",

tensorflow/python/BUILD

@@ -2184,6 +2184,15 @@ py_library(
     visibility = visibility,
 )
 
+# Including this as a dependency will result in tests using
+# :framework_test_lib to NOT use MLIR.
+py_library(
+    name = "is_mlir_bridge_test_false",
+    srcs = ["framework/is_mlir_bridge_test_false.py"],
+    srcs_version = "PY2AND3",
+    visibility = visibility,
+)
+
 # Including this as a dependency will result in tests to use TFRT.
 # TODO(b/153582383): Move tf_ops_alwayslink dependency to c_api_tfrt instead.
 py_library(

tensorflow/python/framework/is_mlir_bridge_test_false.py

@@ -0,0 +1,30 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Including this as a dependency will result in tests NOT using MLIR bridge.
+
+This function is defined by default in test_util.py to None. The test_util then
+attempts to import this module. If this file is made available through the BUILD
+rule, then this function is overridden and will instead cause Tensorflow graphs
+to be always NOT be compiled with MLIR bridge.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+def is_mlir_bridge_enabled():
+  """Returns false if the MLIR bridge should be not be enabled for tests."""
+  return False

tensorflow/python/framework/is_mlir_bridge_test_true.py

@@ -26,5 +26,5 @@ from __future__ import print_function
 
 
 def is_mlir_bridge_enabled():
-  """Returns true to if MLIR bridge should be enabled for tests."""
+  """Returns true if MLIR bridge should be enabled for tests."""
   return True

tensorflow/python/framework/test_util.py

@@ -104,15 +104,19 @@ except Exception:  # pylint: disable=broad-except
   pass
 
 
-# Uses the same mechanism as above to selectively enable MLIR compilation.
+# Uses the same mechanism as above to selectively enable/disable MLIR
+# compilation.
 def is_mlir_bridge_enabled():
   return False
 
 
 try:
-  from tensorflow.python.framework.is_mlir_bridge_test_true import is_mlir_bridge_enabled  # pylint: disable=g-import-not-at-top, unused-import
-except Exception:  # pylint: disable=broad-except
-  pass
+  from tensorflow.python.framework.is_mlir_bridge_test_false import is_mlir_bridge_enabled  # pylint: disable=g-import-not-at-top, unused-import
+except ImportError:
+  try:
+    from tensorflow.python.framework.is_mlir_bridge_test_true import is_mlir_bridge_enabled  # pylint: disable=g-import-not-at-top, unused-import
+  except ImportError:
+    pass
 
 
 # Uses the same mechanism as above to selectively enable TFRT.

tensorflow/tensorflow.bzl

@@ -2874,6 +2874,9 @@ def tf_enable_mlir_bridge():
         str(Label("//tensorflow:enable_mlir_bridge")): [
             "//tensorflow/python:is_mlir_bridge_test_true",
         ],
+        str(Label("//tensorflow:disable_mlir_bridge")): [
+            "//tensorflow/python:is_mlir_bridge_test_false",
+        ],
         "//conditions:default": [],
     })