[Mlir-commits] [mlir] 06b057d - [MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks (#188956)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-27T11:46:46+01:00
New Revision: 06b057d2546ed64018b44df5ebe0e3ecd1c07b26

URL: https://github.com/llvm/llvm-project/commit/06b057d2546ed64018b44df5ebe0e3ecd1c07b26
DIFF: https://github.com/llvm/llvm-project/commit/06b057d2546ed64018b44df5ebe0e3ecd1c07b26.diff

LOG: [MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks (#188956)

WalkPatternRewriteDriver's ErasedOpsListener incorrectly flagged
erasures of ops/blocks that were created during the current pattern
application. Since those ops were never in the walk schedule, erasing
them is safe.

Track newly inserted ops and blocks per visited op; skip the erasure
check for them. Also fix two related issues:
- DropUnitDims: use a fresh IRRewriter (not inheriting the walk pattern
rewriter's listener) for replaceUnitDimIndexOps on cloned ops.
- TestPatterns CloneRegionBeforeOp: wrap op->setAttr() in
modifyOpInPlace to properly notify the rewriter of the in-place change.

Assisted-by: Claude Code

Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

Co-authored-by: Claude Sonnet 4.6 <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
    mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
    mlir/test/lib/Dialect/Test/TestPatterns.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index c3dca148b7f94..4bb4ad61cf8da 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -578,8 +578,10 @@ 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.
-    IRRewriter rewriter(b);
+    //     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());
     replaceUnitDimIndexOps(replacementOp, droppedDims, rewriter);
 
     return replacementOp;

diff  --git a/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
index 1382550e0f7e6..40fcb351ee079 100644
--- a/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
@@ -19,6 +19,7 @@
 #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"
@@ -61,16 +62,38 @@ 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 {
-    checkErasure(op);
+    if (!newlyCreatedOps.contains(op))
+      checkErasure(op);
+    newlyCreatedOps.erase(op);
     ForwardingListener::notifyOperationErased(op);
   }
 
   void notifyBlockErased(Block *block) override {
-    checkErasure(block->getParentOp());
+    if (!newlyCreatedBlocks.contains(block))
+      checkErasure(block->getParentOp());
+    newlyCreatedBlocks.erase(block);
     ForwardingListener::notifyBlockErased(block);
   }
 
@@ -86,6 +109,9 @@ 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
@@ -204,6 +230,8 @@ 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 6c44ace831e96..7e8ed94630b61 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -383,7 +383,8 @@ struct CloneRegionBeforeOp : public RewritePattern {
       return failure();
     for (Region &r : op->getRegions())
       rewriter.cloneRegionBefore(r, op->getBlock());
-    op->setAttr("was_cloned", rewriter.getUnitAttr());
+    rewriter.modifyOpInPlace(
+        op, [&]() { op->setAttr("was_cloned", rewriter.getUnitAttr()); });
     return success();
   }
 };