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:
Marissa Ikonomidis 2020-10-05 14:27:17 -07:00 committed by TensorFlower Gardener
parent 122655424e
commit 4bf9a8cc0f
4 changed files with 67 additions and 1 deletions

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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