Fix ResourceOpLiftingPass to handle call ops that return tf.ReadVariableOp results.

It is possible for a tf.StatefulPartitionedCall to return the result of a tf.ReadVariableOp (with some potential forwarding through ops like tf.Identity). As return op operands are captured prior to replacing tf.ReadVariableOp results with function args, the new function return operands may not be correct. Instead, when replacing tf.ReadVariableOp results with function args, the operands of the new return are updated.

PiperOrigin-RevId: 320597502
Change-Id: I81f614e0b89670c978da376d5810ff82502f601f

Author: andyly

tensorflow/compiler/mlir/tensorflow/tests/resource_op_lifting.mlir

@@ -710,3 +710,29 @@ func @callee(%arg0: tensor<*x!tf.resource<tensor<f32>>>, %arg1: tensor<*x!tf.res
   %0 = "tf._Unknown_"() : () -> tensor<*x!tf.resource<tensor<f32>>>
   return %0 : tensor<*x!tf.resource<tensor<f32>>>
 }
+
+// -----
+
+// Tests call op where it's result is the result of a tf.ReadVariableOp.
+
+// CHECK-LABEL: func @call_with_forwarded_read_only_result
+// CHECK-SAME: (%[[RESOURCE_ARG0:.*]]: tensor<*x!tf.resource<tensor<f32>>>)
+func @call_with_forwarded_read_only_result(%arg0: tensor<*x!tf.resource<tensor<f32>>>) {
+  // CHECK: %[[READ:.*]] = "tf.ReadVariableOp"(%[[RESOURCE_ARG0]])
+  %0 = "tf_device.cluster"() ( {
+    // CHECK: %[[CALL:.*]] = "tf.StatefulPartitionedCall"(%[[READ]])
+    %1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<*x!tf.resource<tensor<f32>>>) -> tensor<f32>
+    // CHECK-NEXT: tf_device.return %[[CALL]]
+    tf_device.return %1 : tensor<f32>
+  }) {} : () -> tensor<f32>
+  return
+}
+
+func @callee(%arg0: tensor<*x!tf.resource<tensor<f32>>>) -> tensor<f32> {
+  %0 = "tf.ReadVariableOp"(%arg0) {device = ""} : (tensor<*x!tf.resource<tensor<f32>>>) -> tensor<f32>
+  %1 = "tf.Identity"(%0) {device = ""} : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK:      func @callee_resource_lifted(%[[A0:.*]]: tensor<f32>) -> tensor<f32>
+// CHECK-NEXT:   return %[[A0]]

tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc

@@ -880,23 +880,34 @@ LogicalResult HandlePartitionedCallOpCallee(
     result->arg_data_type_and_updated_output_index[entry.getFirst()] = {
         entry.getSecond(), -1};
   }
-  llvm::SmallVector<Value, 4> new_retvals;
-  for (auto val : callee.front().getTerminator()->getOperands()) {
-    // Remove resource type outputs.
-    if (getElementTypeOrSelf(val.getType()).isa<TF::ResourceType>()) continue;
-    new_retvals.push_back(val);
+  llvm::SmallVector<int64_t, 4> retval_indices_to_preserve;
+  for (auto& val : callee.front().getTerminator()->getOpOperands()) {
+    // Store indices of results that are not resources.
+    if (!getElementTypeOrSelf(val.get().getType()).isa<TF::ResourceType>())
+      retval_indices_to_preserve.push_back(val.getOperandNumber());
   }
+  int64_t num_retvals = retval_indices_to_preserve.size();
+  llvm::SmallVector<Value, 4> new_retvals;
   // Lift resources.
   LiftArgRetResourcesForFunction(
       callee, remaining_resource_data_types, [&](int64_t index, Value value) {
         result->arg_data_type_and_updated_output_index[index].second =
-            new_retvals.size();
+            num_retvals++;
         new_retvals.push_back(value);
       });
+
   auto old_return = callee.front().getTerminator();
+  llvm::SmallVector<Value, 4> old_and_new_retvals;
+  old_and_new_retvals.reserve(retval_indices_to_preserve.size() +
+                              new_retvals.size());
+  for (int64_t retval_index : retval_indices_to_preserve)
+    old_and_new_retvals.push_back(old_return->getOperand(retval_index));
+
+  old_and_new_retvals.append(new_retvals.begin(), new_retvals.end());
   // Replace old return with the new ones with update values.
   OpBuilder builder(old_return);
-  auto new_return = builder.create<ReturnOp>(old_return->getLoc(), new_retvals);
+  auto new_return =
+      builder.create<ReturnOp>(old_return->getLoc(), old_and_new_retvals);
   old_return->erase();
   callee.setType(FunctionType::get(
       callee.getType().getInputs(),