[Mlir-commits] [mlir] [mlir] Avoid rewriting unreachable blocks in the greedy driver (PR #225814)
Mehdi Amini
llvmlistbot at llvm.org
Fri Sep 25 03:09:22 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/225814
>From 9f2f84294f8da9f66f3ab4700d921d5d60386c6e Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Fri, 4 Sep 2026 00:01:22 -0700
Subject: [PATCH] [mlir] Avoid rewriting unreachable blocks in the greedy
driver
A rewrite can disconnect a block after the iteration's initial CFG sweep.
Track reachability through rewriter notifications and skip worklist operations
in unreachable blocks until the next iteration removes them.
Cache reachability per region. Compare the final successors of changed blocks
after each rewrite, preserving the cache for successor-equivalent changes and
block merges that keep surviving blocks reachable. Recompute for other CFG
changes. This avoids repeated full CFG scans during branch canonicalization.
Single-block regions return before cache lookup or CFG traversal. In the same
Release build, canonicalizing 32,000 sibling single-block regions took 22.8 ms
versus 21.9 ms with the parent driver. With 16,000 regions containing folds,
the times were 35.7 ms and 34.3 ms, respectively. Nesting to 256 levels
showed no measurable difference.
Cover inserted and redirected blocks, blocks connected before rewrite
completion, nested regions, and explicit op-list behavior.
Assisted-by: Codex
---
.../Transforms/GreedyPatternRewriteDriver.h | 3 +
.../Utils/GreedyPatternRewriteDriver.cpp | 193 +++++++++++++++++-
mlir/test/Transforms/canonicalize-dce.mlir | 41 ++++
.../Transforms/greedy-cfg-reachability.mlir | 50 +++++
.../test-strict-pattern-driver.mlir | 20 +-
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 51 ++++-
6 files changed, 354 insertions(+), 4 deletions(-)
create mode 100644 mlir/test/Transforms/greedy-cfg-reachability.mlir
diff --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index 7fef17d806ef3b..5674728513f69d 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -255,6 +255,9 @@ applyPatternsGreedily(Operation *op, const FrozenRewritePatternSet &patterns,
/// Also performs simple dead-code elimination before attempting to match any of
/// the provided patterns.
///
+/// Operations in blocks that are unreachable from their region's entry block
+/// are skipped, even if they are explicitly listed in `ops`.
+///
/// Newly created ops and other pre-existing ops that use results of rewritten
/// ops or supply operands to such ops are also processed, unless such ops are
/// excluded via `config.strictMode`. Any other ops remain unmodified (i.e.,
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba06..cf825ec7a02285 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -353,7 +353,8 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
/// Process ops until the worklist is empty or `config.maxNumRewrites` is
- /// reached. Return `true` if any IR was changed.
+ /// reached. Skip operations in unreachable blocks. Return `true` if any IR
+ /// was changed.
bool processWorklist();
/// The pattern rewriter that is used for making IR modifications and is
@@ -378,6 +379,24 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
private:
+ struct ReachabilityState {
+ Block *entry = nullptr;
+ llvm::SmallDenseSet<Block *, 16> blocks;
+ DenseMap<Block *, SmallVector<Block *, 2>> successors;
+ llvm::SmallDenseSet<Block *, 4> changedBlocks;
+ llvm::SmallDenseSet<Block *, 4> erasedBlocks;
+ bool forceRecompute = false;
+ };
+
+ /// Listener notifications record changed blocks. Reachability is updated
+ /// when the next operation is processed, after the rewrite has completed.
+ DenseMap<Region *, ReachabilityState> reachableBlocks;
+
+ void invalidateRegion(Region *region);
+ void markBlockChanged(Region *region, Block *block);
+ void markBlockErased(Block *block);
+ bool isReachable(Operation *op);
+
/// Look over the provided operands for any defining operations that should
/// be re-added to the worklist. This function should be called when an
/// operation is modified or removed, as it may trigger further
@@ -441,6 +460,9 @@ GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
}
bool GreedyPatternRewriteDriver::processWorklist() {
+ reachableBlocks.clear();
+ llvm::scope_exit clearReachability([&] { reachableBlocks.clear(); });
+
#ifndef NDEBUG
const char *logLineComment =
"//===-------------------------------------------===//\n";
@@ -465,6 +487,10 @@ bool GreedyPatternRewriteDriver::processWorklist() {
(numRewrites < config.getMaxNumRewrites() ||
config.getMaxNumRewrites() == GreedyRewriteConfig::kNoLimit)) {
auto *op = worklist.pop();
+ if (!isReachable(op)) {
+ LLVM_DEBUG(logger.startLine() << "Skipping unreachable operation\n");
+ continue;
+ }
LLVM_DEBUG({
logger.getOStream() << "\n";
@@ -638,6 +664,149 @@ bool GreedyPatternRewriteDriver::processWorklist() {
return changed;
}
+void GreedyPatternRewriteDriver::invalidateRegion(Region *region) {
+ if (region)
+ reachableBlocks.erase(region);
+}
+
+void GreedyPatternRewriteDriver::markBlockChanged(Region *region,
+ Block *block) {
+ if (region) {
+ auto it = reachableBlocks.find(region);
+ if (it != reachableBlocks.end())
+ it->second.changedBlocks.insert(block);
+ }
+}
+
+void GreedyPatternRewriteDriver::markBlockErased(Block *block) {
+ Region *region = block->getParent();
+ if (!region)
+ return;
+ auto it = reachableBlocks.find(region);
+ if (it != reachableBlocks.end()) {
+ if (it->second.entry == block)
+ it->second.forceRecompute = true;
+ it->second.changedBlocks.insert(block);
+ it->second.erasedBlocks.insert(block);
+ }
+}
+
+static SmallVector<Block *, 2> getUniqueSuccessors(Block *block) {
+ SmallVector<Block *, 2> successors;
+ for (Block *successor :
+ llvm::make_range(block->succ_begin(), block->succ_end()))
+ if (!llvm::is_contained(successors, successor))
+ successors.push_back(successor);
+ return successors;
+}
+
+bool GreedyPatternRewriteDriver::isReachable(Operation *op) {
+ // A detached operation or block has no region entry to compare against.
+ Block *block = op->getBlock();
+ if (!block)
+ return true;
+ Region *region = block->getParent();
+ if (!region || region->hasOneBlock())
+ return true;
+
+ auto [it, inserted] = reachableBlocks.try_emplace(region);
+ ReachabilityState &state = it->second;
+ auto recompute = [&] {
+ state = ReachabilityState();
+ state.entry = ®ion->front();
+ SmallVector<Block *> worklist{state.entry};
+ while (!worklist.empty()) {
+ Block *reachable = worklist.pop_back_val();
+ if (!state.blocks.insert(reachable).second)
+ continue;
+ auto &successors = state.successors[reachable];
+ successors = getUniqueSuccessors(reachable);
+ worklist.append(successors);
+ }
+ };
+ if (inserted || state.forceRecompute || state.entry != ®ion->front()) {
+ recompute();
+ } else if (!state.changedBlocks.empty()) {
+ // Check the final CFG, after all notifications from the rewrite. Replacing
+ // a terminator without changing its successors needs no traversal. A
+ // block merge contracts an erased successor into its predecessor: if the
+ // predecessor now reaches all of the erased block's old successors, every
+ // path through that block still reaches the same surviving blocks.
+ DenseMap<Block *, SmallVector<Block *, 2>> currentSuccessors;
+ for (Block *changed : state.changedBlocks) {
+ if (!state.erasedBlocks.contains(changed) &&
+ changed->getParent() == region)
+ currentSuccessors[changed] = getUniqueSuccessors(changed);
+ }
+
+ llvm::SmallDenseSet<Block *, 4> contractedBlocks;
+ bool mustRecompute = false;
+ for (auto &entry : currentSuccessors) {
+ Block *changed = entry.first;
+ auto ¤t = entry.second;
+ if (!state.blocks.contains(changed))
+ continue;
+ for (Block *oldSuccessor : state.successors.lookup(changed)) {
+ if (llvm::is_contained(current, oldSuccessor))
+ continue;
+ auto old = state.successors.find(oldSuccessor);
+ if (!state.erasedBlocks.contains(oldSuccessor) ||
+ old == state.successors.end() ||
+ !llvm::all_of(old->second, [&](Block *successor) {
+ return llvm::is_contained(current, successor);
+ })) {
+ mustRecompute = true;
+ break;
+ }
+ contractedBlocks.insert(oldSuccessor);
+ }
+ if (mustRecompute)
+ break;
+ }
+ if (!mustRecompute) {
+ for (Block *changed : state.changedBlocks) {
+ if (currentSuccessors.contains(changed) ||
+ !state.blocks.contains(changed))
+ continue;
+ if (!state.erasedBlocks.contains(changed) ||
+ (!contractedBlocks.contains(changed) &&
+ !state.successors.lookup(changed).empty())) {
+ mustRecompute = true;
+ break;
+ }
+ }
+ }
+ if (mustRecompute) {
+ recompute();
+ } else {
+ // Additions can only make more blocks reachable. Visit those blocks
+ // without rescanning the already reachable part of the region.
+ for (Block *erased : state.erasedBlocks) {
+ state.blocks.erase(erased);
+ state.successors.erase(erased);
+ }
+ SmallVector<Block *> worklist;
+ for (auto &[changed, current] : currentSuccessors) {
+ if (!state.blocks.contains(changed))
+ continue;
+ state.successors[changed] = current;
+ worklist.append(current);
+ }
+ while (!worklist.empty()) {
+ Block *reachable = worklist.pop_back_val();
+ if (!state.blocks.insert(reachable).second)
+ continue;
+ auto &successors = state.successors[reachable];
+ successors = getUniqueSuccessors(reachable);
+ worklist.append(successors);
+ }
+ state.changedBlocks.clear();
+ state.erasedBlocks.clear();
+ }
+ }
+ return state.blocks.contains(block);
+}
+
void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
assert(op && "expected valid op");
// Gather potential ancestors while looking for a "scope" parent region.
@@ -665,17 +834,30 @@ void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
void GreedyPatternRewriteDriver::notifyBlockInserted(
Block *block, Region *previous, Region::iterator previousIt) {
+ markBlockChanged(previous, block);
+ markBlockChanged(block->getParent(), block);
+ if (auto it = reachableBlocks.find(block->getParent());
+ it != reachableBlocks.end() && it->second.erasedBlocks.contains(block))
+ it->second.forceRecompute = true;
if (RewriterBase::Listener *listener = config.getListener())
listener->notifyBlockInserted(block, previous, previousIt);
}
void GreedyPatternRewriteDriver::notifyBlockErased(Block *block) {
+ markBlockErased(block);
if (RewriterBase::Listener *listener = config.getListener())
listener->notifyBlockErased(block);
}
void GreedyPatternRewriteDriver::notifyOperationInserted(
Operation *op, OpBuilder::InsertPoint previous) {
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors()) {
+ markBlockChanged(op->getParentRegion(), op->getBlock());
+ if (previous.isSet())
+ markBlockChanged(previous.getBlock()->getParent(), previous.getBlock());
+ }
+ for (Region ®ion : op->getRegions())
+ invalidateRegion(®ion);
LLVM_DEBUG({
logger.startLine() << "** Insert : '" << op->getName() << "'(" << op
<< ")\n";
@@ -688,6 +870,10 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(
}
void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+ markBlockChanged(op->getParentRegion(), op->getBlock());
+ for (Region ®ion : op->getRegions())
+ invalidateRegion(®ion);
LLVM_DEBUG({
logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
<< ")\n";
@@ -731,6 +917,11 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
}
void GreedyPatternRewriteDriver::notifyOperationErased(Operation *op) {
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+ markBlockChanged(op->getParentRegion(), op->getBlock());
+ for (Region ®ion : op->getRegions()) {
+ reachableBlocks.erase(®ion);
+ }
LLVM_DEBUG({
logger.startLine() << "** Erase : '" << op->getName() << "'(" << op
<< ")\n";
diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir
index 84631947970ded..e73e2b134b1941 100644
--- a/mlir/test/Transforms/canonicalize-dce.mlir
+++ b/mlir/test/Transforms/canonicalize-dce.mlir
@@ -187,3 +187,44 @@ func.func @f() {
%0 = "test.test_effects_result"() : () -> i32
return
}
+
+// -----
+
+// Test case: Terminate when a block becomes unreachable while patterns run.
+
+// CHECK: func @f
+// CHECK-NEXT: return
+
+func.func @f(%cond: i1, %init: i32) {
+ %true = arith.constant true
+ %c1 = arith.constant 1 : i32
+ cf.cond_br %true, ^exit, ^header(%init : i32)
+^header(%iv: i32):
+ cf.cond_br %cond, ^exit, ^latch
+^latch:
+ %next = arith.addi %iv, %c1 : i32
+ cf.br ^header(%next : i32)
+^exit:
+ return
+}
+
+// -----
+
+// The same CFG change can occur in a region nested below the driver root.
+// CHECK-LABEL: func.func @nested_unreachable
+// CHECK-NOT: arith.addi
+// CHECK: return
+func.func @nested_unreachable(%cond: i1, %init: i32) {
+ scf.execute_region {
+ %true = arith.constant true
+ cf.cond_br %true, ^exit, ^header(%init : i32)
+ ^header(%iv: i32):
+ cf.cond_br %cond, ^exit, ^latch
+ ^latch:
+ %next = arith.addi %iv, %iv : i32
+ cf.br ^header(%next : i32)
+ ^exit:
+ scf.yield
+ }
+ return
+}
diff --git a/mlir/test/Transforms/greedy-cfg-reachability.mlir b/mlir/test/Transforms/greedy-cfg-reachability.mlir
new file mode 100644
index 00000000000000..b7fe5ddc46284f
--- /dev/null
+++ b/mlir/test/Transforms/greedy-cfg-reachability.mlir
@@ -0,0 +1,50 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -test-greedy-patterns | FileCheck %s
+
+// A rewrite can create an unreachable block after the iteration's initial
+// unreachable-block sweep. Do not process its self-referential operation.
+// CHECK-LABEL: func.func @insert_unreachable
+// CHECK-NOT: arith.addi
+// CHECK: return
+func.func @insert_unreachable(%arg: i32) {
+ "test.greedy_create_block"(%arg) {mode = "unreachable"} : (i32) -> ()
+ return
+}
+
+// -----
+
+// A block connected before the rewrite returns remains reachable.
+// CHECK-LABEL: func.func @connect_before_return
+// CHECK: "test.keep"()
+// CHECK: return
+func.func @connect_before_return() {
+ "test.greedy_create_block"() {mode = "reachable"} : () -> ()
+ return
+}
+
+// -----
+
+// Inserting a new entry can disconnect the old entry without editing its
+// terminator.
+// CHECK-LABEL: func.func @insert_entry
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @insert_entry() {
+ "test.greedy_create_block"() {mode = "entry"} : () -> ()
+ return
+}
+
+// -----
+
+// An in-place successor change can disconnect a block.
+// CHECK-LABEL: func.func @redirect
+// CHECK-NOT: "test.keep"
+// CHECK: return
+func.func @redirect() {
+ "test.greedy_create_block"() {mode = "redirect"} : () -> ()
+ cf.br ^body
+^body:
+ "test.keep"() : () -> ()
+ cf.br ^exit
+^exit:
+ return
+}
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index c87444cba8e1a2..8a4735b3f2c894 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -75,8 +75,8 @@ func.func @test_replace_with_erase_op() {
// CHECK-AN: return
// CHECK-AN: ^[[BB1:[^:]*]]:
// CHECK-AN: "test.implicit_change_op"()[^[[BB1]]]
-func.func @test_trigger_rewrite_through_block() {
- return
+func.func @test_trigger_rewrite_through_block(%cond: i1) {
+ cf.cond_br %cond, ^bb1, ^bb3
^bb1:
// Uses bb1. ChangeBlockOp replaces that and all other usages of bb1 with bb2.
"test.change_block_op"() [^bb1, ^bb2] : () -> ()
@@ -91,6 +91,22 @@ func.func @test_trigger_rewrite_through_block() {
// -----
+// Explicitly selected operations in unreachable blocks are skipped.
+// CHECK-AN-LABEL: func @test_skip_unreachable
+// CHECK-AN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = false
+// CHECK-AN: "test.erase_op"()
+// CHECK-EN-LABEL: func @test_skip_unreachable
+// CHECK-EN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = false
+// CHECK-EN: "test.erase_op"()
+func.func @test_skip_unreachable() {
+ return
+^dead:
+ "test.erase_op"() : () -> ()
+ return
+}
+
+// -----
+
// CHECK-AN: notifyOperationErased: test.foo_b
// CHECK-AN: notifyOperationErased: test.foo_a
// CHECK-AN: notifyOperationErased: test.graph_region
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 06517910e270e8..6acc658507483d 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -11,6 +11,7 @@
#include "TestTypes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/CommonFolders.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/ControlFlow/Transforms/StructuralTypeConversions.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
@@ -468,6 +469,54 @@ class CreateAndEraseOpAndBlock : public RewritePattern {
}
};
+/// Exercise reachability changes while a greedy rewrite is processing its
+/// worklist. A new block is either disconnected, connected within the same
+/// rewrite, or inserted before the old entry block; another mode redirects a
+/// branch in-place.
+class CreateBlockDuringGreedyRewrite : public RewritePattern {
+public:
+ CreateBlockDuringGreedyRewrite(MLIRContext *context)
+ : RewritePattern("test.greedy_create_block", /*benefit=*/1, context) {}
+
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ StringRef mode = cast<StringAttr>(op->getAttr("mode")).getValue();
+ Region *region = op->getParentRegion();
+ Location loc = op->getLoc();
+ if (mode == "entry") {
+ rewriter.createBlock(region, region->begin());
+ func::ReturnOp::create(rewriter, loc);
+ } else if (mode == "redirect") {
+ auto branch = cast<cf::BranchOp>(op->getBlock()->getTerminator());
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(®ion->back()); });
+ } else {
+ Block *newBlock = rewriter.createBlock(region);
+ if (mode == "unreachable") {
+ Value input = op->getOperand(0);
+ auto add = arith::AddIOp::create(rewriter, loc, input, input);
+ rewriter.modifyOpInPlace(add, [&] {
+ add->setOperand(0, add.getResult());
+ add->setOperand(1, add.getResult());
+ });
+ cf::BranchOp::create(rewriter, loc, newBlock);
+ } else {
+ assert(mode == "reachable");
+ rewriter.create(
+ loc, OperationName("test.keep", op->getContext()).getIdentifier(),
+ ValueRange(), TypeRange());
+ func::ReturnOp::create(rewriter, loc);
+ Operation *oldTerminator = op->getBlock()->getTerminator();
+ rewriter.setInsertionPoint(oldTerminator);
+ cf::BranchOp::create(rewriter, loc, newBlock);
+ rewriter.eraseOp(oldTerminator);
+ }
+ }
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
struct TestGreedyPatternDriver
: public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)
@@ -486,7 +535,7 @@ struct TestGreedyPatternDriver
patterns.add<FoldingPattern, TestNamedPatternRule,
FolderInsertBeforePreviouslyFoldedConstantPattern,
FolderCommutativeOp2WithConstant, HoistEligibleOps,
- MakeOpEligible>(&getContext());
+ MakeOpEligible, CreateBlockDuringGreedyRewrite>(&getContext());
// Additional patterns for testing the GreedyPatternRewriteDriver.
patterns.insert<IncrementIntAttribute<3>>(&getContext());
More information about the Mlir-commits
mailing list