Remove no-op array updates.

PiperOrigin-RevId: 684530855
hzeller · Oct 10, 2024 · 83d5b9e · 83d5b9e
1 parent c1c2ebc
commit 83d5b9e
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/xls/passes/array_simplification_pass.cc b/xls/passes/array_simplification_pass.cc
@@ -608,6 +608,30 @@ absl::StatusOr<SimplifyResult> SimplifyArrayUpdate(
     return SimplifyResult::Changed({array_update});
   }
 
+  // An array update with the update value being an 'array index' on the to
+  // update array at the same index as the update can be replaced with the
+  // original array.
+  //
+  // arr := ARRAY...
+  // idx := INDEX...
+  // array-update(arr, array-index(arr, idx) idx)) -> arr
+  if (array_update->update_value()->Is<ArrayIndex>() &&
+      array_update->update_value()->As<ArrayIndex>()->array() ==
+          array_update->array_to_update()) {
+    ArrayIndex* val = array_update->update_value()->As<ArrayIndex>();
+    bool indices_are_equal =
+        val->indices().size() == array_update->indices().size();
+    for (int64_t i = 0; indices_are_equal && i < val->indices().size(); ++i) {
+      indices_are_equal = query_engine.NodesKnownUnsignedEquals(
+          val->indices()[i], array_update->indices()[i]);
+    }
+    if (indices_are_equal) {
+      XLS_RETURN_IF_ERROR(
+          array_update->ReplaceUsesWith(array_update->array_to_update()));
+      return SimplifyResult::Changed({});
+    }
+  }
+
   // Try to simplify a kArray operation followed by an ArrayUpdate operation
   // into a single kArray operation which uses the updated value. For example:
   //

diff --git a/xls/passes/array_simplification_pass_test.cc b/xls/passes/array_simplification_pass_test.cc
@@ -1413,5 +1413,21 @@ TEST_F(ArraySimplificationPassTest, RemovalOfUpdate) {
   EXPECT_THAT(f->nodes(), Each(Not(m::Array())));
 }
 
+TEST_F(ArraySimplificationPassTest, NoOpArrayUpdate) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue arr = fb.Param("array", p->GetArrayType(10, p->GetBitsType(32)));
+  BValue idx = fb.Param("upd_idx", p->GetBitsType(32));
+  BValue val_at_idx = fb.ArrayIndex(arr, {idx});
+  BValue update_arr = fb.ArrayUpdate(arr, val_at_idx, {idx});
+  fb.ArrayIndex(update_arr, {fb.Param("second_idx", p->GetBitsType(32))});
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.Build());
+  ScopedVerifyEquivalence sve(f);
+  ScopedRecordIr sri(p.get());
+  ASSERT_THAT(Run(f), IsOkAndHolds(true));
+  EXPECT_THAT(f->return_value(),
+              m::ArrayIndex(m::Param("array"), {m::Param("second_idx")}));
+}
+
 }  // namespace
 }  // namespace xls