From eda3a5ec4d076dae7056b35452b8fd376273d9d3 Mon Sep 17 00:00:00 2001
From: Ian Wood <ianwood2024@u.northwestern.edu>
Date: Fri, 15 Nov 2024 18:13:57 -0800
Subject: [PATCH] Erase state of modified ops

Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
---
 .../Dialect/Util/Transforms/OptimizeIntArithmetic.cpp      | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
index d4b3a14b43f5..94b2ee237398 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
@@ -292,6 +292,13 @@ class DataFlowListener : public RewriterBase::Listener {
       s.eraseState(res);
   }
 
+  void notifyOperationModified(Operation *op) override {
+    s.eraseState(s.getProgramPointAfter(op));
+    for (Value res : op->getResults()) {
+      s.eraseState(res);
+    }
+  }
+
   DataFlowSolver &s;
 };