When canonicalizing only erase the operation after calling the 'fold' hook if replace...

PiperOrigin-RevId: 232757964

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 45c57e2f307084dacba71404d57425f8509c7da1..c2adf84edd3ffa1b07573b1b0369216bcbe139ca 100644 (file)
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -267,9 +267,9 @@ void GreedyPatternRewriteDriver::simplifyFunction() {
             addToWorklist(operand.getOwner());
           res->replaceAllUsesWith(resultValues[i]);
         }
-      }
 
-      op->erase();
+        op->erase();
+      }
       continue;
     }