[Mlir-commits] [mlir] Revert "[MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks" (PR #189010)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/189010

Reverts llvm/llvm-project#188956

Hit "merge" by accident on the wrong tab, juggling too may PRs in parallel...

>From 0d7c0dd55af8b7a9d9df886ecb475f610abb92f9 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Mar 2026 15:29:26 +0100
Subject: [PATCH] =?UTF-8?q?Revert=20"[MLIR]=20Fix=20ErasedOpsListener=20fa?=
 =?UTF-8?q?lse=20positives=20for=20newly=20created=20ops/bl=E2=80=A6"?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This reverts commit 06b057d2546ed64018b44df5ebe0e3ecd1c07b26.
---
 .../Linalg/Transforms/DropUnitDims.cpp        |  6 ++--
 .../Utils/WalkPatternRewriteDriver.cpp        | 32 ++-----------------
 mlir/test/lib/Dialect/Test/TestPatterns.cpp   |  3 +-
 3 files changed, 5 insertions(+), 36 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 4bb4ad61cf8da..c3dca148b7f94 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -578,10 +578,8 @@ linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
     b.cloneRegionBefore(genericOp.getRegion(), replacementOp.getRegion(),
                         replacementOp.getRegion().begin());
     // 5a. Replace `linalg.index` operations that refer to the dropped unit
-    //     dimensions. Use a fresh IRRewriter to avoid inheriting any listener
-    //     from the builder (e.g., WalkPatternRewriter's erasure listener),
-    //     since the ops being erased here are newly cloned, not the matched op.
-    IRRewriter rewriter(b.getContext());
+    //     dimensions.
+    IRRewriter rewriter(b);
     replaceUnitDimIndexOps(replacementOp, droppedDims, rewriter);
 
     return replacementOp;
diff --git a/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
index 40fcb351ee079..1382550e0f7e6 100644
--- a/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
@@ -19,7 +19,6 @@
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Rewrite/PatternApplicator.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/DebugLog.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -62,38 +61,16 @@ struct WalkAndApplyPatternsAction final
 // ops/blocks. Because we use walk-based pattern application, erasing the
 // op/block from the *next* iteration (e.g., a user of the visited op) is not
 // valid. Note that this is only used with expensive pattern API checks.
-//
-// Ops and blocks that were *created* during the current pattern application are
-// exempt: they were not in the walk schedule before the pattern ran, so erasing
-// them cannot disrupt the walk.
 struct ErasedOpsListener final : RewriterBase::ForwardingListener {
   using RewriterBase::ForwardingListener::ForwardingListener;
 
-  void notifyOperationInserted(Operation *op,
-                               OpBuilder::InsertPoint previous) override {
-    if (visitedOp)
-      newlyCreatedOps.insert(op);
-    ForwardingListener::notifyOperationInserted(op, previous);
-  }
-
-  void notifyBlockInserted(Block *block, Region *previous,
-                           Region::iterator previousIt) override {
-    if (visitedOp)
-      newlyCreatedBlocks.insert(block);
-    ForwardingListener::notifyBlockInserted(block, previous, previousIt);
-  }
-
   void notifyOperationErased(Operation *op) override {
-    if (!newlyCreatedOps.contains(op))
-      checkErasure(op);
-    newlyCreatedOps.erase(op);
+    checkErasure(op);
     ForwardingListener::notifyOperationErased(op);
   }
 
   void notifyBlockErased(Block *block) override {
-    if (!newlyCreatedBlocks.contains(block))
-      checkErasure(block->getParentOp());
-    newlyCreatedBlocks.erase(block);
+    checkErasure(block->getParentOp());
     ForwardingListener::notifyBlockErased(block);
   }
 
@@ -109,9 +86,6 @@ struct ErasedOpsListener final : RewriterBase::ForwardingListener {
   }
 
   Operation *visitedOp = nullptr;
-  // Ops and blocks inserted since visitedOp was last set; may be freely erased.
-  DenseSet<Operation *> newlyCreatedOps;
-  DenseSet<Block *> newlyCreatedBlocks;
 };
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 } // namespace
@@ -230,8 +204,6 @@ void walkAndApplyPatterns(Operation *op,
                    << OpWithFlags(op, OpPrintingFlags().skipRegions());
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
             erasedListener.visitedOp = op;
-            erasedListener.newlyCreatedOps.clear();
-            erasedListener.newlyCreatedBlocks.clear();
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
             if (succeeded(applicator.matchAndRewrite(op, rewriter)))
               LDBG() << "\tOp matched and rewritten";
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 7e8ed94630b61..6c44ace831e96 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -383,8 +383,7 @@ struct CloneRegionBeforeOp : public RewritePattern {
       return failure();
     for (Region &r : op->getRegions())
       rewriter.cloneRegionBefore(r, op->getBlock());
-    rewriter.modifyOpInPlace(
-        op, [&]() { op->setAttr("was_cloned", rewriter.getUnitAttr()); });
+    op->setAttr("was_cloned", rewriter.getUnitAttr());
     return success();
   }
 };