[Mlir-commits] [mlir] 9d80f17 - [mlir][bufferize] Fix tracking of erased ops

[Matthias Springer]

Author: Matthias Springer
Date: 2023-03-27T08:58:53+02:00
New Revision: 9d80f1742294aeedfbe547349a5422c363b5fb1a

URL: https://github.com/llvm/llvm-project/commit/9d80f1742294aeedfbe547349a5422c363b5fb1a
DIFF: https://github.com/llvm/llvm-project/commit/9d80f1742294aeedfbe547349a5422c363b5fb1a.diff

LOG: [mlir][bufferize] Fix tracking of erased ops

This is a workaround until D144193 has landed.

Differential Revision: https://reviews.llvm.org/D146868

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 0b10bafb9f163..e5e125f731032 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -358,9 +358,12 @@ class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {
 
 protected:
   void notifyOperationRemoved(Operation *op) override {
-    erasedOps.insert(op);
-    // Erase if present.
-    toMemrefOps.erase(op);
+    // TODO: Walk can be removed when D144193 has landed.
+    op->walk([&](Operation *op) {
+      erasedOps.insert(op);
+      // Erase if present.
+      toMemrefOps.erase(op);
+    });
   }
 
   void notifyOperationInserted(Operation *op) override {