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

Set resource as mutable when unknown op is encountered in the optimize global tensors pass

Browse Source 
PiperOrigin-RevId: 297274143
Change-Id: Ia33df38d23875a1a4021c52e70fca17b8aeb9c96
This commit is contained in:
Ashwin Murthy 2020-02-25 21:08:14 -08:00 committed by TensorFlower Gardener
parent fa08cfd489
commit 3755911efc
2 changed files with 52 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir
View File

@ -156,3 +156,31 @@ module attributes {tf_saved_model.semantics} {
return %val : tensor<f32>
}
}
// -----
// CHECK-LABEL: module attributes {tf_saved_model.semantics}
module attributes {tf_saved_model.semantics} {
// Test case: Inter-procedural analysis with resource usage in an
// unknown op, we assume mutating behavior and propagate that.
// CHECK: "tf_saved_model.global_tensor"() {
// CHECK-SAME: is_mutable
// CHECK-SAME: } : () -> ()
"tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<f32>, value = dense<42.> : tensor<f32> } : () -> ()
func @exported_f(%arg0: tensor<!tf.resource<tensor<f32>>> {tf_saved_model.bound_input = @v}) -> (tensor<f32> {tf_saved_model.index_path = []})
attributes {tf_saved_model.exported_names = ["exported_f"]} {
%val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f} : (tensor<!tf.resource<tensor<f32>>>) -> (tensor<f32>)
return %val : tensor<f32>
}
// CHECK: func @f(%arg0: tensor<*x!tf.resource>) -> tensor<f32>
func @f(%arg0: tensor<*x!tf.resource>) -> tensor<f32> {
%c0 = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
"tf.AssignAddVariableOp"(%arg0, %c0) : (tensor<*x!tf.resource>, tensor<f32>) -> ()
return %c0 : tensor<f32>
}
}

44
tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
View File

@ -64,21 +64,6 @@ static bool IsResourceType(Type type) {
static bool IsResource(Value value) { return IsResourceType(value.getType()); }
static bool DoesOpUseResources(Operation* op) {
for (auto operand : op->getOperands()) {
if (IsResource(operand)) {
return true;
}
}
bool uses_resources = false;
visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
if (IsResource(operand->get())) {
uses_resources = true;
}
});
return uses_resources;
}
class ResourceAnalyzer {
public:
explicit ResourceAnalyzer(ModuleOp module) {
@ -145,17 +130,36 @@ class ResourceAnalyzer {
}
return;
}
// TODO(ashwinm): This should not error but refactor DoesOpUseResources
// to tell us which resource, so that it can be set as potentially written
// to.
if (DoesOpUseResources(op)) {
op->emitError() << "unknown op that uses resources";
// For all other ops, we assume it mutates all resources it uses, so
// this errs on the side of being conservative. We should improve
// this by using either a property or a trait that clearly
// identifies ops with resource mutating behavior.
if (PropagatePotentiallyWrittenWithinUnhandledOp(op)) {
return;
}
});
return success();
}
// If an op is not one of the handled ones, we assume all resource usages
// within its purview are mutating in nature.
bool PropagatePotentiallyWrittenWithinUnhandledOp(Operation* op) {
for (auto operand : op->getOperands()) {
if (IsResource(operand)) {
SetPotentiallyWritten(operand);
return true;
}
}
bool uses_resources = false;
visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
if (IsResource(operand->get())) {
SetPotentiallyWritten(operand->get());
uses_resources = true;
}
});
return uses_resources;
}
// Given a funcOp associated with the callee and operands from the
// corresponding callOp, propagate the potentially written decision to the
// callOp's operands, if the corresponding func's arguments are potentially