Make early increment range when modifying underlying user list.
PiperOrigin-RevId: 335700983
Change-Id: I248f4c8aad523e947ab8ecdc8340874531aba9de
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
index 3e6d4f3..3302ec5 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir
@@ -258,6 +258,19 @@
// -----
+// Tests removal of dead local variables.
+
+// CHECK-LABEL: func @main
+func @main(%arg0: tensor<2xf32>) {
+ // CHECK-NOT: tf.MlirLocalVarOp
+ // CHECK-NOT: tf.AssignVariableOp
+ %0 = "tf.MlirLocalVarOp"() : () -> tensor<!tf.resource<tensor<2xf32>>>
+ "tf.AssignVariableOp"(%0, %arg0) : (tensor<!tf.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
+ return
+}
+
+// -----
+
// Tests first read of one resource is used as a value to write to another
// resource.
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting_cleanup.cc b/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting_cleanup.cc
index fb232c7..b635096 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting_cleanup.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting_cleanup.cc
@@ -52,11 +52,11 @@
}
}
for (auto local_var : local_vars) {
- if (llvm::all_of(local_var.resource().getUsers(),
- [](const Operation *user) {
- return isa<TF::AssignVariableOp>(user);
- })) {
- for (auto user : local_var.resource().getUsers()) user->erase();
+ auto users = local_var.resource().getUsers();
+ if (llvm::all_of(users, [](const Operation *user) {
+ return isa<TF::AssignVariableOp>(user);
+ })) {
+ for (auto user : llvm::make_early_inc_range(users)) user->erase();
local_var.erase();
}
}
