experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Do not fail when TPUReplicateMetadata op is missing.

Browse Source 
The rationale behind this change is that functionalizing control flow in TF v1 models adds functions to function library with nodes that have '_tpu_replicate' attrribute. The TPU cluster formation pass throws error on these functions when it does not find the TPUReplicateMetadata op associated with a _tpu_replicate attribute.

PiperOrigin-RevId: 292986754
Change-Id: I203c4e62db96c835bdad3f669500674dc4fce8c5
This commit is contained in:
Prakalp Srivastava 2020-02-03 13:06:53 -08:00 committed by TensorFlower Gardener
parent 2742aa7a3b
commit 7954fb8fd1
2 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir
View File

@ -407,6 +407,16 @@ func @bad_num_replicas() {
return
}
// -----
// Test that functions without TPUReplicateMetadata op are skipped without
// error
// CHECK-LABEL: func @missing_metadata_op
func @missing_metadata_op() {
// expected-warning@+1 {{TPUReplicateMetadata for associated '_tpu_replicate' attribute 'replicate' is missing}}
%0 = "tf.opA"() {_tpu_replicate = "replicate"} : () -> tensor<i1>
return
}
// -----

10
tensorflow/compiler/mlir/tensorflow/transforms/tpu_cluster_formation.cc
View File

@ -414,10 +414,12 @@ LogicalResult FormClustersInBlock(Block* block,
auto cluster_metadata = metadata_map.find(cluster.getFirst());
// No TPUReplicateMetadata for a `_tpu_replicate` attribute.
if (cluster_metadata == metadata_map.end())
return cluster_ops.front()->emitError()
<< "TPUReplicateMetadata for associated '" << kTPUReplicateAttr
<< "' attribute '" << cluster.getFirst() << "' is missing";
if (cluster_metadata == metadata_map.end()) {
cluster_ops.front()->emitWarning()
<< "TPUReplicateMetadata for associated '" << kTPUReplicateAttr
<< "' attribute '" << cluster.getFirst() << "' is missing";
continue;
}
llvm::SmallSetVector<Operation*, 8> preceding_users =
CollectClusterPrecedingUsers(block, cluster_ops);