[Mlir-commits] [mlir] 7002c4f - [MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks (#192291)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 19 01:49:03 PDT 2026
Author: Mehdi Amini
Date: 2026-08-19T10:48:58+02:00
New Revision: 7002c4f33be0f2ce8b3fa4da7b84cf6dbde0f9d1
URL: https://github.com/llvm/llvm-project/commit/7002c4f33be0f2ce8b3fa4da7b84cf6dbde0f9d1
DIFF: https://github.com/llvm/llvm-project/commit/7002c4f33be0f2ce8b3fa4da7b84cf6dbde0f9d1.diff
LOG: [MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks (#192291)
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 the TestPatterns CloneRegionBeforeOp pattern
to wrap op->setAttr() in modifyOpInPlace so the rewriter observes the change.
Add a focused walk-driver regression that creates and erases an
operation and a block during one pattern application, and checks that
listener notifications are still forwarded.
Fix some failures present with
MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.
Assisted-by: Codex
Co-authored-by: Claude Sonnet 4.6 <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
mlir/test/Conversion/SCFToOpenMP/reductions.mlir
mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir
mlir/test/Conversion/SCFToOpenMP/vector-reduction.mlir
mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
mlir/test/IR/test-walk-pattern-rewrite-driver.mlir
mlir/test/Transforms/test-strict-pattern-driver.mlir
mlir/test/lib/Dialect/Test/TestPatterns.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/WalkPatternRewriteDriver.cpp
index 1382550e0f7e6..0aec6e7e7d920 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 invalidate the current walk iterator.
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/Conversion/SCFToOpenMP/reductions.mlir b/mlir/test/Conversion/SCFToOpenMP/reductions.mlir
index 879291d58b615..8e2343df48877 100644
--- a/mlir/test/Conversion/SCFToOpenMP/reductions.mlir
+++ b/mlir/test/Conversion/SCFToOpenMP/reductions.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt -convert-scf-to-openmp -split-input-file %s | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// CHECK: omp.declare_reduction @[[$REDF:.*]] : f32
// CHECK: init
diff --git a/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir b/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir
index e41b23c36c047..d362bb6092419 100644
--- a/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir
+++ b/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt -convert-scf-to-openmp='num-threads=4' %s | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// CHECK-LABEL: @parallel
func.func @parallel(%arg0: index, %arg1: index, %arg2: index,
%arg3: index, %arg4: index, %arg5: index) {
diff --git a/mlir/test/Conversion/SCFToOpenMP/vector-reduction.mlir b/mlir/test/Conversion/SCFToOpenMP/vector-reduction.mlir
index b146aec0e287d..baad2e9937c63 100644
--- a/mlir/test/Conversion/SCFToOpenMP/vector-reduction.mlir
+++ b/mlir/test/Conversion/SCFToOpenMP/vector-reduction.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt %s --convert-scf-to-openmp | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// CHECK-LABEL: omp.declare_reduction @__scf_reduction : vector<2xi1>
// CHECK: init {
// CHECK: %[[INIT:.*]] = llvm.mlir.constant(dense<true> : vector<2xi1>) : vector<2xi1>
diff --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
index dd7ce878cae90..841d0e5f56512 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -1,8 +1,6 @@
// RUN: mlir-opt %s -linalg-fold-unit-extent-dims -split-input-file | FileCheck %s
// RUN: mlir-opt %s -linalg-fold-unit-extent-dims="use-rank-reducing-slices" -cse -split-input-file | FileCheck %s --check-prefix=CHECK-SLICES
-// XFAIL: mlir-expensive-checks
-
#accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>,
affine_map<(i, j, k, l, m) -> ()>,
diff --git a/mlir/test/IR/test-walk-pattern-rewrite-driver.mlir b/mlir/test/IR/test-walk-pattern-rewrite-driver.mlir
index c3063416b0360..c16a40f1a1a61 100644
--- a/mlir/test/IR/test-walk-pattern-rewrite-driver.mlir
+++ b/mlir/test/IR/test-walk-pattern-rewrite-driver.mlir
@@ -1,6 +1,13 @@
// RUN: mlir-opt %s --test-walk-pattern-rewrite-driver="dump-notifications=true" \
// RUN: --allow-unregistered-dialect --split-input-file | FileCheck %s
+// Check that newly created operations and blocks may be erased without
+// invalidating the walk, while still forwarding listener notifications.
+// CHECK: notifyOperationInserted: test.transient_op, was unlinked
+// CHECK-NEXT: notifyOperationErased: test.transient_op
+// CHECK-NEXT: notifyBlockInserted into func.func: was unlinked
+// CHECK-NEXT: notifyBlockErased
+
// The following op is updated in-place and will not be added back to the worklist.
// CHECK-LABEL: func.func @inplace_update()
// CHECK: "test.any_attr_of_i32_str"() <{attr = 1 : i32}> : () -> ()
@@ -120,6 +127,12 @@ func.func @erase_nested_block() -> i32 {
return %a : i32
}
+// CHECK-LABEL: func.func @create_and_erase_op_and_block
+// CHECK: "test.create_and_erase_op_and_block"() {was_rewritten}
+func.func @create_and_erase_op_and_block() {
+ "test.create_and_erase_op_and_block"() : () -> ()
+ return
+}
// CHECK-LABEL: func.func @unreachable_replace_with_new_op
// CHECK: "test.new_op"
@@ -138,4 +151,3 @@ func.func @unreachable_replace_with_new_op() {
%c = "test.replace_with_new_op"() : () -> (i32)
return
}
-
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index b4b620e1519b2..c87444cba8e1a 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -1,5 +1,3 @@
-// XFAIL: mlir-expensive-checks
-
// RUN: mlir-opt \
// RUN: -test-strict-pattern-driver="strictness=AnyOp" \
// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHECK-AN
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 552a1a473c9fd..9121155f70245 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -384,7 +384,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();
}
};
@@ -437,6 +438,33 @@ class EraseFirstBlock : public RewritePattern {
}
};
+/// Creates and immediately erases an operation and a block.
+class CreateAndEraseOpAndBlock : public RewritePattern {
+public:
+ CreateAndEraseOpAndBlock(MLIRContext *context)
+ : RewritePattern("test.create_and_erase_op_and_block", /*benefit=*/1,
+ context) {}
+
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ if (op->hasAttr("was_rewritten"))
+ return failure();
+
+ Operation *newOp = rewriter.create(
+ op->getLoc(),
+ OperationName("test.transient_op", op->getContext()).getIdentifier(),
+ ValueRange(), TypeRange());
+ rewriter.eraseOp(newOp);
+
+ Block *newBlock = rewriter.createBlock(op->getParentRegion());
+ rewriter.eraseBlock(newBlock);
+
+ rewriter.modifyOpInPlace(
+ op, [&]() { op->setAttr("was_rewritten", rewriter.getUnitAttr()); });
+ return success();
+ }
+};
+
struct TestGreedyPatternDriver
: public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)
@@ -709,8 +737,8 @@ struct TestWalkPatternDriver final
// Patterns for testing the WalkPatternRewriteDriver.
patterns.add<IncrementIntAttribute<3>, MoveBeforeParentOp,
- MoveAfterParentOp, CloneOp, ReplaceWithNewOp, EraseFirstBlock>(
- &getContext());
+ MoveAfterParentOp, CloneOp, ReplaceWithNewOp, EraseFirstBlock,
+ CreateAndEraseOpAndBlock>(&getContext());
DumpNotifications dumpListener;
walkAndApplyPatterns(getOperation(), std::move(patterns),
More information about the Mlir-commits
mailing list