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

PiperOrigin-RevId: 297274143
Change-Id: Ia33df38d23875a1a4021c52e70fca17b8aeb9c96

tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir

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

tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc

@@ -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
+      // For all other ops, we assume it mutates all resources it uses, so
-      // to tell us which resource, so that it can be set as potentially written
+      // this errs on the side of being conservative. We should improve
-      // to.
+      // this by using either a property or a trait that clearly
-      if (DoesOpUseResources(op)) {
+      // identifies ops with resource mutating behavior.
-        op->emitError() << "unknown op that uses resources";
+      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