[Mlir-commits] [mlir] [mlir] Avoid rewriting unreachable blocks in the greedy driver (PR #225814)
Mehdi Amini
llvmlistbot at llvm.org
Sat Sep 26 05:26:27 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/225814
>From e091c08539fe78a2c02a88eff538af9068adcedc 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 and compare final successors after each
rewrite. Preserve the cache for successor-equivalent changes, added edges,
and block merges that keep surviving blocks reachable. Recompute for other
CFG changes, including entry changes and cross-region block moves.
Keep cache updates separate from reachability queries, returning failure
when a full traversal is required. Collect each successor set once for
validation and retain compact vector snapshots in the cache. Check
containing blocks when querying operations in nested regions, caching the
ancestor result until listener notifications invalidate it.
Cover inserted and redirected blocks, blocks connected before rewrite
completion, nested regions, block moves and merges, and explicit op lists.
Use observer diagnostics to check whether affected operations are processed
within a single greedy iteration. Document that skipped operations are not
automatically retried if a later rewrite makes their block reachable, and
test that behavior in each strictness mode.
Assisted-by: Codex
---
.../Transforms/GreedyPatternRewriteDriver.h | 12 +-
.../Utils/GreedyPatternRewriteDriver.cpp | 257 +++++++++++++++++-
mlir/test/Transforms/canonicalize-dce.mlir | 41 +++
.../Transforms/greedy-cfg-reachability.mlir | 190 +++++++++++++
.../test-strict-pattern-driver.mlir | 56 +++-
mlir/test/lib/Dialect/Test/TestPatterns.cpp | 93 ++++++-
6 files changed, 636 insertions(+), 13 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 7fef17d806ef3..1ae4565994f93 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -255,6 +255,12 @@ 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`. Making a skipped
+/// operation's block reachable later does not by itself re-enqueue that
+/// operation. Such operations may remain unprocessed even when this function
+/// returns success.
+///
/// 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.,
@@ -271,9 +277,9 @@ applyPatternsGreedily(Operation *op, const FrozenRewritePatternSet &patterns,
/// `applyPatternsGreedily` should be used.
///
/// Returns "success" if the iterative process converged (i.e., fixpoint was
-/// reached) and no more patterns can be matched. `changed` is set to "true" if
-/// the IR was modified at all. `allOpsErased` is set to "true" if all ops in
-/// `ops` were erased.
+/// reached) and no more patterns can be matched, except for operations skipped
+/// as described above. `changed` is set to "true" if the IR was modified at
+/// all. `allOpsErased` is set to "true" if all ops in `ops` were erased.
LogicalResult
applyOpPatternsGreedily(ArrayRef<Operation *> ops,
const FrozenRewritePatternSet &patterns,
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba0..ad87c4b1fdb0b 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -27,10 +27,14 @@
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/ScopedPrinter.h"
#include "llvm/Support/raw_ostream.h"
+#include <cstdint>
+#include <utility>
+
#ifdef MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
#include <random>
#endif // MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
@@ -353,7 +357,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 +383,37 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
private:
+ struct ReachabilityState {
+ /// Extend the successor cache from these roots, stopping at visited blocks.
+ void visit(SmallVector<Block *> worklist);
+
+ /// Update the cache after a rewrite. On failure, the caller must recompute
+ /// reachability because some successor snapshots may already be updated.
+ LogicalResult update();
+
+ // Keys are the reachable blocks; values are their cached unique successors.
+ DenseMap<Block *, SmallVector<Block *, 2>> successors;
+ llvm::SmallDenseSet<Block *, 4> changedBlocks;
+ // Erased block pointers must not be inspected after the rewrite.
+ llvm::SmallDenseSet<Block *, 4> erasedBlocks;
+ };
+
+ /// Listener notifications record changed blocks. Reachability is updated
+ /// when the next operation is processed, after the rewrite has completed.
+ DenseMap<Region *, ReachabilityState> reachableBlocks;
+
+ // Whether each region's owning operation is reachable in its containing
+ // region. Listener notifications invalidate these answers with a version
+ // change, avoiding repeated walks through unchanged containing regions.
+ DenseMap<Region *, std::pair<uint64_t, bool>> ancestorReachability;
+ uint64_t reachabilityVersion = 0;
+
+ void invalidateRegion(Region *region);
+ void markBlockChanged(Block *block);
+ void markBlockErased(Block *block);
+ bool isBlockReachable(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 +477,11 @@ GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
}
bool GreedyPatternRewriteDriver::processWorklist() {
+ llvm::scope_exit clearReachability([&] {
+ reachableBlocks.clear();
+ ancestorReachability.clear();
+ });
+
#ifndef NDEBUG
const char *logLineComment =
"//===-------------------------------------------===//\n";
@@ -465,6 +506,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 +683,183 @@ bool GreedyPatternRewriteDriver::processWorklist() {
return changed;
}
+void GreedyPatternRewriteDriver::invalidateRegion(Region *region) {
+ if (region) {
+ reachableBlocks.erase(region);
+ ++reachabilityVersion;
+ }
+}
+
+void GreedyPatternRewriteDriver::markBlockChanged(Block *block) {
+ if (!block)
+ return;
+ ++reachabilityVersion;
+ if (Region *region = block->getParent()) {
+ 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;
+ ++reachabilityVersion;
+ auto it = reachableBlocks.find(region);
+ if (it != reachableBlocks.end()) {
+ if (block == ®ion->front()) {
+ reachableBlocks.erase(it);
+ return;
+ }
+ it->second.changedBlocks.insert(block);
+ it->second.erasedBlocks.insert(block);
+ }
+}
+
+static llvm::SmallPtrSet<Block *, 4> getUniqueSuccessors(Block *block) {
+ llvm::SmallPtrSet<Block *, 4> successors;
+ // A pattern may temporarily leave an invalid cross-region block operand.
+ // Such an edge cannot make a block reachable in this region.
+ for (Block *successor : block->getSuccessors())
+ if (successor->getParent() == block->getParent())
+ successors.insert(successor);
+ return successors;
+}
+
+void GreedyPatternRewriteDriver::ReachabilityState::visit(
+ SmallVector<Block *> worklist) {
+ while (!worklist.empty()) {
+ Block *reachable = worklist.pop_back_val();
+ auto [it, inserted] = successors.try_emplace(reachable);
+ if (!inserted)
+ continue;
+ auto current = getUniqueSuccessors(reachable);
+ it->second.assign(current.begin(), current.end());
+ worklist.append(it->second);
+ }
+}
+
+LogicalResult GreedyPatternRewriteDriver::ReachabilityState::update() {
+ // 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.
+ // Keep erased blocks' old successors until validation finishes: they
+ // explain lost edges from surviving predecessors.
+ SmallVector<Block *> worklist;
+ for (Block *changed : changedBlocks) {
+ auto oldChanged = successors.find(changed);
+ // Cross-region moves invalidate the source cache immediately.
+ // Newly reachable blocks will have their successors read by visit().
+ if (oldChanged == successors.end() || erasedBlocks.contains(changed))
+ continue;
+ auto current = getUniqueSuccessors(changed);
+ for (Block *oldSuccessor : oldChanged->second) {
+ if (current.contains(oldSuccessor))
+ continue;
+ // A -> B -> C can become A -> C when B is merged into A. If the
+ // removed edge does not have this explanation, rescan the CFG.
+ auto old = successors.find(oldSuccessor);
+ if (!erasedBlocks.contains(oldSuccessor) || old == successors.end() ||
+ !llvm::all_of(old->second, [&](Block *successor) {
+ return current.contains(successor);
+ }))
+ return failure();
+ }
+ worklist.append(current.begin(), current.end());
+ oldChanged->second.assign(current.begin(), current.end());
+ }
+ // Additions can only make more blocks reachable. Visit those blocks
+ // without rescanning the already reachable part of the region.
+ for (Block *erased : erasedBlocks)
+ successors.erase(erased);
+ visit(std::move(worklist));
+ changedBlocks.clear();
+ erasedBlocks.clear();
+ return success();
+}
+
+bool GreedyPatternRewriteDriver::isBlockReachable(Block *block) {
+ Region *region = block->getParent();
+ // The entry block needs no cache lookup or CFG traversal.
+ if (!region || block == ®ion->front())
+ return true;
+ auto [it, inserted] = reachableBlocks.try_emplace(region);
+ ReachabilityState &state = it->second;
+ auto recompute = [&] {
+ state = ReachabilityState();
+ state.visit({®ion->front()});
+ };
+ if (inserted) {
+ recompute();
+ } else if (!state.changedBlocks.empty() && failed(state.update())) {
+ recompute();
+ }
+ return state.successors.contains(block);
+}
+
+bool GreedyPatternRewriteDriver::isReachable(Operation *op) {
+ Block *block = op->getBlock();
+ if (!block)
+ return true;
+ if (!isBlockReachable(block))
+ return false;
+ Region *region = block->getParent();
+ if (!region || region == config.getScope())
+ return true;
+
+ // Most regions are nested only a few levels below the scope. Their entry
+ // blocks need no cache entries or ancestor traversal beyond those levels.
+ Region *ancestor = region;
+ for (unsigned depth = 0; depth < 3; ++depth) {
+ Operation *parentOp = ancestor->getParentOp();
+ if (!parentOp)
+ return true;
+ Block *parentBlock = parentOp->getBlock();
+ if (!parentBlock)
+ return true;
+ Region *parentRegion = parentBlock->getParent();
+ if (!parentRegion || parentBlock != &parentRegion->front())
+ break;
+ if (parentRegion == config.getScope())
+ return true;
+ ancestor = parentRegion;
+ }
+
+ // A nested region can be locally reachable while its owning operation is in
+ // an unreachable block. Cache the parent-chain result until a notification
+ // changes a CFG or moves a region-owning operation.
+ SmallVector<Region *, 4> regionsToCache;
+ bool reachable = true;
+ while (region != config.getScope()) {
+ Operation *parentOp = region->getParentOp();
+ if (!parentOp)
+ break;
+ auto it = ancestorReachability.find(region);
+ if (it != ancestorReachability.end() &&
+ it->second.first == reachabilityVersion) {
+ reachable = it->second.second;
+ break;
+ }
+ regionsToCache.push_back(region);
+ block = parentOp->getBlock();
+ if (!block)
+ break;
+ if (!isBlockReachable(block)) {
+ reachable = false;
+ break;
+ }
+ region = block->getParent();
+ if (!region)
+ break;
+ }
+ for (Region *region : regionsToCache)
+ ancestorReachability[region] = {reachabilityVersion, reachable};
+ return reachable;
+}
+
void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
assert(op && "expected valid op");
// Gather potential ancestors while looking for a "scope" parent region.
@@ -665,17 +887,41 @@ void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
void GreedyPatternRewriteDriver::notifyBlockInserted(
Block *block, Region *previous, Region::iterator previousIt) {
+ Region *region = block->getParent();
+ // Do not retain pointers to blocks that can subsequently be erased in a
+ // different region without notifying the source region's cache.
+ if (previous != region)
+ invalidateRegion(previous);
+ // A newly allocated block may reuse an erased block's address. Discard its
+ // old successor snapshot and erased-block marker before tracking the new
+ // block.
+ if (auto it = reachableBlocks.find(region);
+ it != reachableBlocks.end() && it->second.erasedBlocks.contains(block))
+ invalidateRegion(region);
+ // The entry is the root of a cached traversal. Inserting or moving a block
+ // before it can change reachability without changing any successor.
+ if (region && block == ®ion->front())
+ invalidateRegion(region);
+ markBlockChanged(block);
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 (previous.isSet() && op->getNumRegions())
+ ++reachabilityVersion;
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors()) {
+ markBlockChanged(op->getBlock());
+ if (previous.isSet())
+ markBlockChanged(previous.getBlock());
+ }
LLVM_DEBUG({
logger.startLine() << "** Insert : '" << op->getName() << "'(" << op
<< ")\n";
@@ -688,6 +934,8 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(
}
void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+ markBlockChanged(op->getBlock());
LLVM_DEBUG({
logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
<< ")\n";
@@ -731,6 +979,13 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
}
void GreedyPatternRewriteDriver::notifyOperationErased(Operation *op) {
+ if (op->getNumRegions())
+ ++reachabilityVersion;
+ if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+ markBlockChanged(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 84631947970de..e73e2b134b194 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 0000000000000..211189c6ac341
--- /dev/null
+++ b/mlir/test/Transforms/greedy-cfg-reachability.mlir
@@ -0,0 +1,190 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -test-greedy-patterns="top-down max-iterations=1" -verify-diagnostics | FileCheck %s
+
+// Start each rewrite in a nonentry block to populate the reachability cache
+// before changing the CFG, even when entry-block queries bypass the cache.
+
+// A rewrite can create an unreachable block after the iteration's initial
+// unreachable-block sweep. Its observer must not be processed.
+// CHECK-LABEL: func.func @insert_unreachable
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @insert_unreachable() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "unreachable"} : () -> ()
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// A block connected before the rewrite returns is processed in this iteration,
+// without relying on a fresh reachability cache in the next iteration.
+// CHECK-LABEL: func.func @connect_before_return
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @connect_before_return(%cond: i1) {
+ cf.br ^start
+^start:
+ // expected-remark @+1 {{processed reachable block}}
+ "test.greedy_create_block"(%cond) {mode = "reachable"} : (i1) -> ()
+ cf.br ^exit
+^exit:
+ 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() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "entry"} : () -> ()
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// Moving an existing block to the front must also invalidate the old entry.
+// CHECK-LABEL: func.func @move_entry
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @move_entry() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "move-entry"} : () -> ()
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// The first nested op populates the ancestor cache, then moves its two-region
+// owner into an unreachable block. The observer must see the invalidation.
+// CHECK-LABEL: func.func @move_nested_unreachable
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @move_nested_unreachable() {
+ cf.br ^start
+^start:
+ scf.execute_region {
+ scf.execute_region {
+ "test.greedy_create_block"() {mode = "move-parent"} : () -> ()
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ scf.yield
+ }
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// A branch rewrite can disconnect an outer block while leaving its nested
+// region in place. The observer must not use the cached ancestor result.
+// CHECK-LABEL: func.func @disconnect_nested
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @disconnect_nested() {
+ cf.br ^body
+^body:
+ scf.execute_region {
+ "test.greedy_create_block"() {mode = "disconnect-parent"} : () -> ()
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ scf.yield
+ }
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// An in-place successor change can disconnect a block.
+// CHECK-LABEL: func.func @redirect
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @redirect() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "redirect"} : () -> ()
+ cf.br ^body
+^body:
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// Populate the source cache, then move and erase a nonentry block in one
+// rewrite. The next query must not inspect the erased block in that cache.
+// CHECK-LABEL: func.func @move_then_erase
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @move_then_erase() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "move-erase"} : () -> ()
+ // expected-remark @+1 {{processed reachable block}}
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ cf.br ^body
+^body:
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// Merging a block preserves reachability of its successors. The cache is
+// populated before the rewrite, and the observer must run in this iteration.
+// CHECK-LABEL: func.func @merge
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @merge() {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "merge"} : () -> ()
+ cf.br ^body
+^body:
+ cf.br ^exit
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ return
+}
+
+// -----
+
+// Merging a block and dropping one of its successors must force a rescan:
+// preserving reachability based only on the erased block would process ^dead.
+// CHECK-LABEL: func.func @merge_disconnect
+// CHECK-NOT: "test.greedy_create_block"
+// CHECK: return
+func.func @merge_disconnect(%cond: i1) {
+ cf.br ^start
+^start:
+ "test.greedy_create_block"() {mode = "merge-disconnect"} : () -> ()
+ cf.br ^body
+^body:
+ cf.cond_br %cond, ^exit, ^dead
+^dead:
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ return
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ "test.greedy_create_block"() {mode = "observe"} : () -> ()
+ return
+}
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index c87444cba8e1a..77feaa800cd9a 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -1,12 +1,12 @@
-// RUN: mlir-opt \
+// RUN: mlir-opt -allow-unregistered-dialect \
// RUN: -test-strict-pattern-driver="strictness=AnyOp" \
// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHECK-AN
-// RUN: mlir-opt \
+// RUN: mlir-opt -allow-unregistered-dialect \
// RUN: -test-strict-pattern-driver="strictness=ExistingAndNewOps" \
// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHECK-EN
-// RUN: mlir-opt \
+// RUN: mlir-opt -allow-unregistered-dialect \
// RUN: -test-strict-pattern-driver="strictness=ExistingOps" \
// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHECK-EX
@@ -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,52 @@ 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
+}
+
+// -----
+
+// The explicit worklist is processed in reverse walk order: skip the erase op
+// first, then reconnect its block. Reconnection alone must not re-enqueue it.
+// The pass checks that the driver returns success.
+// CHECK-AN-LABEL: func @test_reconnect_skipped
+// CHECK-AN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = true
+// CHECK-AN-NEXT: cf.br ^[[DEST:.*]]
+// CHECK-AN: ^[[DEST]]:
+// CHECK-AN-NEXT: "test.erase_op"()
+// CHECK-EN-LABEL: func @test_reconnect_skipped
+// CHECK-EN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = true
+// CHECK-EN-NEXT: cf.br ^[[DEST:.*]]
+// CHECK-EN: ^[[DEST]]:
+// CHECK-EN-NEXT: "test.erase_op"()
+// CHECK-EX-LABEL: func @test_reconnect_skipped
+// CHECK-EX-SAME: pattern_driver_all_erased = false, pattern_driver_changed = true
+// CHECK-EX-NEXT: cf.br ^[[DEST:.*]]
+// CHECK-EX: ^[[DEST]]:
+// CHECK-EX-NEXT: "test.erase_op"()
+func.func @test_reconnect_skipped() {
+ "test.greedy_create_block"() {mode = "redirect"} : () -> ()
+ cf.br ^exit
+^exit:
+ 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 06517910e270e..62a344934dc49 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,87 @@ 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. Other modes redirect a
+/// branch, merge blocks, move a nested operation, or move and erase a block.
+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 == "observe") {
+ op->emitRemark("processed reachable block");
+ } else if (mode == "disconnect-parent") {
+ Block *parentBlock = region->getParentOp()->getBlock();
+ Region *parentRegion = parentBlock->getParent();
+ auto branch = cast<cf::BranchOp>(parentRegion->front().getTerminator());
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(&parentRegion->back()); });
+ } else if (mode == "move-parent") {
+ Operation *outer = region->getParentOp()->getParentOp();
+ Region *outerRegion = outer->getParentRegion();
+ Block *newBlock = rewriter.createBlock(outerRegion);
+ func::ReturnOp::create(rewriter, loc);
+ rewriter.moveOpBefore(outer, newBlock, newBlock->begin());
+ } else if (mode == "move-entry") {
+ rewriter.moveBlockBefore(®ion->back(), region, region->begin());
+ } else if (mode == "move-erase") {
+ auto branch = cast<cf::BranchOp>(op->getBlock()->getTerminator());
+ Block *moved = branch.getDest();
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(®ion->back()); });
+ // The source region retains multiple blocks and its original entry.
+ // Erasing the moved block must not leave a dangling pointer in its cached
+ // state.
+ Region temporary;
+ rewriter.moveBlockBefore(moved, &temporary, temporary.end());
+ rewriter.eraseBlock(moved);
+ } else if (mode == "merge" || mode == "merge-disconnect") {
+ Block *dest = op->getBlock();
+ auto branch = cast<cf::BranchOp>(dest->getTerminator());
+ Block *source = branch.getDest();
+ rewriter.eraseOp(branch);
+ rewriter.mergeBlocks(source, dest);
+ if (mode == "merge-disconnect") {
+ auto conditional = cast<cf::CondBranchOp>(dest->getTerminator());
+ rewriter.setInsertionPoint(conditional);
+ cf::BranchOp::create(rewriter, loc, conditional.getTrueDest());
+ rewriter.eraseOp(conditional);
+ }
+ } else 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 {
+ assert(mode == "reachable" || mode == "unreachable");
+ Block *newBlock = rewriter.createBlock(region);
+ OperationState state(loc, "test.greedy_create_block");
+ state.addAttribute("mode", rewriter.getStringAttr("observe"));
+ rewriter.create(state);
+ func::ReturnOp::create(rewriter, loc);
+ if (mode == "reachable") {
+ Operation *oldTerminator = op->getBlock()->getTerminator();
+ rewriter.setInsertionPoint(oldTerminator);
+ Block *oldDest = cast<cf::BranchOp>(oldTerminator).getDest();
+ cf::CondBranchOp::create(rewriter, loc, op->getOperand(0), oldDest,
+ ValueRange(), newBlock, ValueRange());
+ rewriter.eraseOp(oldTerminator);
+ }
+ }
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
struct TestGreedyPatternDriver
: public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)
@@ -486,7 +568,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());
@@ -581,6 +663,7 @@ struct TestStrictPatternDriver
patterns.add<
// clang-format off
ChangeBlockOp,
+ CreateBlockDuringGreedyRewrite,
CloneOp,
CloneRegionBeforeOp,
EraseOp,
@@ -600,7 +683,8 @@ struct TestStrictPatternDriver
opName == "test.move_before_parent_op" ||
opName == "test.inline_blocks_into_parent" ||
opName == "test.split_block_here" || opName == "test.clone_me" ||
- opName == "test.clone_region_before") {
+ opName == "test.clone_region_before" ||
+ opName == "test.greedy_create_block") {
ops.push_back(op);
}
});
@@ -623,8 +707,9 @@ struct TestStrictPatternDriver
// operation will trigger the assertion while processing.
bool changed = false;
bool allErased = false;
- (void)applyOpPatternsGreedily(ArrayRef(ops), std::move(patterns), config,
- &changed, &allErased);
+ if (failed(applyOpPatternsGreedily(ArrayRef(ops), std::move(patterns),
+ config, &changed, &allErased)))
+ return signalPassFailure();
Builder b(ctx);
getOperation()->setDiscardableAttr("pattern_driver_changed",
b.getBoolAttr(changed));
More information about the Mlir-commits
mailing list