Allow MLIR bridge to run in the fallback mode if the user sets the tf_mlir_bridge_safe_mode flag.
PiperOrigin-RevId: 362409942
Change-Id: Id20fb30a52091bd3e86cdf619280468f92ae9114
diff --git a/RELEASE.md b/RELEASE.md
index b3e191b..6a5bbac 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -176,6 +176,11 @@
`tf.config.experimental.mlir_bridge_rollout` to enable a \"safe\" mode.
This runs the MLIR bridge only when an analysis of the graph only when
an analysis of the graph determines that it is safe to run.
+ * Add new enum value 'MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED' to
+ `tf.config.experimental.mlir_bridge_rollout` to enable a fallback for
+ the MLIR bridge in a \"safe\" mode. This runs the MLIR bridge in a
+ FallbackEnabled mode when an analysis of the graph determines
+ that the graph does not have unsupported features.
* Other
* Adding show_debug_info to mlir.convert_graph_def and
diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc
index 9112b8d..110678a 100644
--- a/tensorflow/compiler/jit/flags.cc
+++ b/tensorflow/compiler/jit/flags.cc
@@ -242,7 +242,10 @@
mlir_flags = new MlirCommonFlags;
if (!enable_mlir_bridge_is_explicit) {
mlir_flags->tf_mlir_enable_mlir_bridge =
- ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
+ (mlir_bridge_safe_mode)
+ ? ConfigProto::Experimental::
+ MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED
+ : ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
} else if (enable_mlir_bridge) {
mlir_flags->tf_mlir_enable_mlir_bridge =
(mlir_bridge_safe_mode)
diff --git a/tensorflow/core/protobuf/config.proto b/tensorflow/core/protobuf/config.proto
index 9b50d5e..d4d03bc 100644
--- a/tensorflow/core/protobuf/config.proto
+++ b/tensorflow/core/protobuf/config.proto
@@ -598,6 +598,13 @@
// supported by the MLIR bridge, the MLIR bridge will be used to run the
// graph.
MLIR_BRIDGE_ROLLOUT_SAFE_MODE_ENABLED = 3;
+ // Enable the MLIR bridge in a fallback mode on a per graph basis based
+ // on an analysis of the features used in the graph.
+ // Running the MLIR bridge in the fallback mode means that it is
+ // executed and it commits all the changes to the TF graph in case
+ // of success. And it does not in case of failures and let the old bridge
+ // to process the TF graph.
+ MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED = 4;
}
// This field is underdevelopment, for now use enable_mlir_bridge
// (b/166038521).
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.-experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.-experimental.pbtxt
index c070882..49759cf 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.-experimental.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.-experimental.pbtxt
@@ -108,22 +108,26 @@
}
enum_type {
name: "MlirBridgeRollout"
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_UNSPECIFIED"
number: 0
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_ENABLED"
number: 1
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_DISABLED"
number: 2
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_SAFE_MODE_ENABLED"
number: 3
}
+ value {
+ name: "MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED"
+ number: 4
+ }
}
reserved_range {
start: 2
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.pbtxt
index 34bf0d5..864a943 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-config-proto.pbtxt
@@ -237,22 +237,26 @@
}
enum_type {
name: "MlirBridgeRollout"
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_UNSPECIFIED"
number: 0
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_ENABLED"
number: 1
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_DISABLED"
number: 2
}
- value: {
+ value {
name: "MLIR_BRIDGE_ROLLOUT_SAFE_MODE_ENABLED"
number: 3
}
+ value {
+ name: "MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED"
+ number: 4
+ }
}
reserved_range {
start: 2
