Add new mlir_bridge_rollout field to Config proto.
The exisiting enable_mlir_bridge field is insufficient because it is a boolean. Either the MLIR bridge is enabled or disabled. To slowly and safely rollout the MLIR bridge we need a third category: unspecified. If the field is unspecified then we can selectively choose to enable the bridge. This will allow us to slowly enable the bridge for increasingly more models. PiperOrigin-RevId: 335499348 Change-Id: I53dc4789c5e2aca9203e8c7dc4b0b5ca77c27f26
This commit is contained in:
parent
122655424e
commit
4bf9a8cc0f
@ -311,7 +311,10 @@ h# Release 2.4.0
|
||||
* We have replaced uses of "whitelist" and "blacklist" with "allowlist"
|
||||
and "denylist" where possible. Please see
|
||||
https://developers.google.com/style/word-list#blacklist for more
|
||||
context. <ADD RELEASE NOTES HERE>
|
||||
context.
|
||||
* Add `tf.config.experimental.mlir_bridge_rollout` which will help us
|
||||
rollout the new MLIR TPU bridge.
|
||||
* <ADD RELEASE NOTES HERE>
|
||||
|
||||
## Thanks to our Contributors
|
||||
|
||||
|
@ -568,6 +568,9 @@ message ConfigProto {
|
||||
// Session::Extend() may not be supported.
|
||||
bool optimize_for_static_graph = 12;
|
||||
|
||||
// This field will eventually be deprecated and replaced by
|
||||
// mlir_bridge_rollout (b/166038521).
|
||||
//
|
||||
// Whether to enable the MLIR-based TF->XLA bridge.
|
||||
//
|
||||
// This is a replacement to the existing bridge, and not ready for
|
||||
@ -581,6 +584,22 @@ message ConfigProto {
|
||||
// to lower the encapsulated graph to a particular device.
|
||||
bool enable_mlir_bridge = 13;
|
||||
|
||||
// An enum that describes the state of the MLIR bridge rollout.
|
||||
enum MlirBridgeRollout {
|
||||
// If this field is left unspecified, the MLIR bridge may be selectively
|
||||
// enabled on a per graph basis.
|
||||
MLIR_BRIDGE_ROLLOUT_UNSPECIFIED = 0;
|
||||
// Enabling the MLIR bridge enables it for all graphs in this session.
|
||||
MLIR_BRIDGE_ROLLOUT_ENABLED = 1;
|
||||
// Disabling the MLIR bridge disables it for all graphs in this session.
|
||||
MLIR_BRIDGE_ROLLOUT_DISABLED = 2;
|
||||
}
|
||||
// This field is underdevelopment, for now use enable_mlir_bridge
|
||||
// (b/166038521).
|
||||
//
|
||||
// Whether to enable the MLIR-based TF->XLA bridge.
|
||||
MlirBridgeRollout mlir_bridge_rollout = 17;
|
||||
|
||||
// Whether to enable the MLIR-based Graph optimizations.
|
||||
//
|
||||
// This will become a part of standard Tensorflow graph optimization
|
||||
|
@ -75,6 +75,13 @@ tf_proto {
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_BOOL
|
||||
}
|
||||
field {
|
||||
name: "mlir_bridge_rollout"
|
||||
number: 17
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_ENUM
|
||||
type_name: ".tensorflow.ConfigProto.Experimental.MlirBridgeRollout"
|
||||
}
|
||||
field {
|
||||
name: "enable_mlir_graph_optimization"
|
||||
number: 16
|
||||
@ -93,6 +100,21 @@ tf_proto {
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_INT64
|
||||
}
|
||||
enum_type {
|
||||
name: "MlirBridgeRollout"
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_UNSPECIFIED"
|
||||
number: 0
|
||||
}
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_ENABLED"
|
||||
number: 1
|
||||
}
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_DISABLED"
|
||||
number: 2
|
||||
}
|
||||
}
|
||||
reserved_range {
|
||||
start: 2
|
||||
end: 3
|
||||
|
@ -204,6 +204,13 @@ tf_proto {
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_BOOL
|
||||
}
|
||||
field {
|
||||
name: "mlir_bridge_rollout"
|
||||
number: 17
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_ENUM
|
||||
type_name: ".tensorflow.ConfigProto.Experimental.MlirBridgeRollout"
|
||||
}
|
||||
field {
|
||||
name: "enable_mlir_graph_optimization"
|
||||
number: 16
|
||||
@ -222,6 +229,21 @@ tf_proto {
|
||||
label: LABEL_OPTIONAL
|
||||
type: TYPE_INT64
|
||||
}
|
||||
enum_type {
|
||||
name: "MlirBridgeRollout"
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_UNSPECIFIED"
|
||||
number: 0
|
||||
}
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_ENABLED"
|
||||
number: 1
|
||||
}
|
||||
value: {
|
||||
name: "MLIR_BRIDGE_ROLLOUT_DISABLED"
|
||||
number: 2
|
||||
}
|
||||
}
|
||||
reserved_range {
|
||||
start: 2
|
||||
end: 3
|
||||
|
Loading…
Reference in New Issue
Block a user