[Mlir-commits] [mlir] 2c6c3a3 - [mlir] Avoid rewriting unreachable blocks in the greedy driver (#225814)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 08:51:12 PDT 2026
Author: Mehdi Amini
Date: 2026-10-23T17:51:07+02:00
New Revision: 2c6c3a3c3a5969b86d4afeb446ab8e16baa13bac
URL: https://github.com/llvm/llvm-project/commit/2c6c3a3c3a5969b86d4afeb446ab8e16baa13bac
DIFF: https://github.com/llvm/llvm-project/commit/2c6c3a3c3a5969b86d4afeb446ab8e16baa13bac.diff
LOG: [mlir] Avoid rewriting unreachable blocks in the greedy driver (#225814)
A rewrite can disconnect a block after the iteration's initial CFG sweep.
Operations queued in that block can then develop self-referential SSA uses,
which are permitted in unreachable code but can cause folds or patterns to
crash or repeatedly re-enqueue operations. Worklist processing may never
return, preventing the next unreachable-block sweep from running. The
constant-branch-folding reproducer in #221152 demonstrates this failure.
Track reachability through rewriter notifications and skip worklist
operations in unreachable blocks until the next iteration removes them.
This extends #153957, which fixed the unreachable `arith.addi`
canonicalization hang in #153732 by erasing unreachable blocks before
applying patterns and deferred handling blocks that become unreachable
during rewriting. #154038 similarly protects the walk-based rewrite driver
by skipping unreachable blocks instead of erasing them. The reachable
graph-region self-cycle crash in #194824 and #205064 is a separate
problem; #207185 proposes a fold-specific fix for it.
A forwarding `ReachabilityListener` sits between the rewriter and the
worklist driver. Notifications only record what changed: terminator
changes and blocks that leave a region are recorded per block. Queries happen
between rewrites, once per popped operation, and bring the cache up to
date. A region is traversed from its entry block when it has no cached
state or when its entry block changed. For recorded terminator changes,
the cached successors of each block are compared with its current
successors.
Added edges extend the reachable set. A lost edge whose target is gone
or has no predecessors left, and whose target's cached successors are all
still successors of the changed block, only drops the target; this is the
shape left by folding a constant conditional branch and by merging a block
into its predecessor. Any other lost edge marks the blocks behind its target
uncertain, and only a query for one of them traverses the region.
In-place folds are reported to the listener by the driver. Nested regions check
their containing blocks up to the scope. Caches are cleared at worklist-pass
boundaries. With `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`,
every query is checked against a fresh traversal.
Tests cover inserted and redirected blocks, blocks connected before
rewrite completion, nested regions, block moves and merges, entry-block
insertion and erasure, edges lost by two blocks in one rewrite, a lost edge into a
loop, a block reconnected after being disconnected, in-place folds that
retarget a branch, and explicit op lists. Observer diagnostics check
whether affected operations are processed within a single greedy iteration.
Skipped operations are not automatically retried if a later rewrite makes their
block reachable; this is documented and tested in each strictness mode.
Compile-time impact of `mlir-opt -canonicalize` on synthetic inputs,
same Release build with assertions, min / median wall time in seconds over 3
runs. Outputs are identical with and without this patch. The inputs
behind a loop header cannot merge into the entry block, so every query goes
through the cache; invalidating the cache on block insertion and
erasure, or traversing after every lost edge, made them quadratic (up to 50x
slower).
| Input | Baseline | This patch |
|---|---|---|
| 1 block, 80k foldable arith ops | 0.135 / 0.139 | 0.144 / 0.145 |
| 20k foldable ops in 30-deep nested `scf.if` | 0.079 / 0.080 | 0.078 /
0.079 |
| 20k-block chain of foldable ops behind a loop header, no CFG change |
0.410 / 0.414 | 0.416 / 0.420 |
| 20k functions, each a 4-block CFG with one branch | 0.565 / 0.575 |
0.623 / 0.633 |
| 20k-block chain of constant `cf.cond_br` | 0.151 / 0.158 | 0.158 /
0.159 |
| 20k-block chain of diamonds whose pass-through blocks fold | 0.145 /
0.151 | 0.164 / 0.165 |
| Same diamonds behind a loop header | 0.145 / 0.146 | 0.172 / 0.177 |
| 20k blocks behind a loop header, identical-successor `cf.cond_br` fold
and merge | 0.232 / 0.237 | 0.253 / 0.256 |
| 20k constant `cf.cond_br` into a shared dead block, behind a loop
header | 1.571 / 1.578 | 1.622 / 1.626 |
Assisted-by: Codex
Assisted-by: Claude Code
Added:
mlir/test/Transforms/canonicalize-cfg-reachability.mlir
mlir/test/Transforms/greedy-cfg-reachability.mlir
Modified:
mlir/docs/PatternRewriter.md
mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
mlir/test/Transforms/canonicalize-dce.mlir
mlir/test/Transforms/test-strict-pattern-driver.mlir
mlir/test/lib/Dialect/Test/TestEnumDefs.td
mlir/test/lib/Dialect/Test/TestOpDefs.cpp
mlir/test/lib/Dialect/Test/TestOps.td
mlir/test/lib/Dialect/Test/TestPatterns.cpp
Removed:
################################################################################
diff --git a/mlir/docs/PatternRewriter.md b/mlir/docs/PatternRewriter.md
index 39605c7acb7fe0..d70b049f917fbf 100644
--- a/mlir/docs/PatternRewriter.md
+++ b/mlir/docs/PatternRewriter.md
@@ -381,6 +381,19 @@ from being added to the worklist throughout the rewrite process:
* `GreedyRewriteStrictness::ExistingOps`: Only pre-existing ops (with which
the worklist was initialized) are added to the worklist.
+Both drivers skip operations in blocks that are unreachable from their region's
+entry block, including operations nested within such blocks. Unreachable code
+can contain SSA cycles that cause folds or patterns to crash or loop indefinitely.
+Reachability is checked between rewrites, so a block connected before a rewrite
+returns can be processed. If an operation is skipped, reconnecting its block in a
+later rewrite does not automatically add it back to the worklist. It may remain
+unprocessed even if the driver returns success, regardless of strictness.
+
+The region-based driver erases unreachable blocks at the start of worklist
+processing in each iteration. The op-based driver leaves unreachable blocks and
+their operations in place, even if those operations were explicitly provided to
+the driver.
+
Note: This driver listens for IR changes via the callbacks provided by
`RewriterBase`. It is important that patterns announce all IR changes to the
rewriter and do not bypass the rewriter API by modifying ops directly.
diff --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index 7fef17d806ef3b..a9e4e1a788322d 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -255,6 +255,17 @@ 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 listed in `ops`, and are not re-enqueued if a later
+/// rewrite reconnects their block. Such operations may remain unprocessed even
+/// when this function returns success. This protects folds and patterns from
+/// the SSA cycles that are legal in unreachable code. Reachability is checked
+/// between rewrites, so a block connected before a rewrite returns is eligible
+/// for processing.
+///
+/// Unlike `applyPatternsGreedily`, this function does not erase unreachable
+/// blocks. They are left in place, along with their operations.
+///
/// 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 +282,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 adc571061b7f77..4bcde3c5bf5550 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -27,6 +27,8 @@
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/ScopedPrinter.h"
#include "llvm/Support/raw_ostream.h"
@@ -56,8 +58,8 @@ namespace {
/// This struct stores finger prints of ops to determine whether the IR has
/// changed or not.
struct ExpensiveChecks : public RewriterBase::ForwardingListener {
- ExpensiveChecks(RewriterBase::Listener *driver, Operation *topLevel)
- : RewriterBase::ForwardingListener(driver), topLevel(topLevel) {}
+ ExpensiveChecks(RewriterBase::Listener *listener, Operation *topLevel)
+ : RewriterBase::ForwardingListener(listener), topLevel(topLevel) {}
/// Compute finger prints of the given op and its nested ops.
void computeFingerPrints(Operation *topLevel) {
@@ -133,6 +135,18 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
/// Invalidate the finger print of the given op, i.e., remove it from the map.
void invalidateFingerPrint(Operation *op) { fingerprints.erase(op); }
+ void notifyBlockInserted(Block *block, Region *previous,
+ Region::iterator previousIt) override {
+ RewriterBase::ForwardingListener::notifyBlockInserted(block, previous,
+ previousIt);
+
+ // Inserting or moving a block changes the region structure of both the
+ // new and previous owners. Either region may be detached.
+ invalidateFingerPrint(block->getParentOp());
+ if (previous)
+ invalidateFingerPrint(previous->getParentOp());
+ }
+
void notifyBlockErased(Block *block) override {
RewriterBase::ForwardingListener::notifyBlockErased(block);
@@ -315,6 +329,95 @@ class RandomizedWorklist : public Worklist {
// GreedyPatternRewriteDriver
//===----------------------------------------------------------------------===//
+/// Tracks which blocks are reachable from their region's entry block while
+/// patterns rewrite the CFG, so that the driver can skip operations in blocks
+/// that a rewrite disconnected. Notifications only record what changed; the
+/// queries, which the driver issues between rewrites, bring the cache up to
+/// date. `Block::isReachable` is not used because it traverses the region on
+/// every call, while the driver queries once per processed operation. All
+/// notifications are forwarded to the wrapped listener.
+class ReachabilityListener : public RewriterBase::ForwardingListener {
+public:
+ ReachabilityListener(RewriterBase::Listener *listener, Region *scope)
+ : RewriterBase::ForwardingListener(listener), scope(scope) {}
+
+ /// Return whether `op` is reachable: its block is reachable from the entry
+ /// block of its region, and so is every containing block up to `scope`.
+ /// Must be called between rewrites, when the IR is in a consistent state.
+ bool isReachable(Operation *op);
+
+ /// Drop all cached reachability information.
+ void clear() { regionStates.clear(); }
+
+ void notifyBlockInserted(Block *block, Region *previous,
+ Region::iterator previousIt) override;
+ void notifyBlockErased(Block *block) override;
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override;
+ void notifyOperationModified(Operation *op) override;
+ void notifyOperationErased(Operation *op) override;
+
+private:
+ /// The reachability cache of one region.
+ struct RegionState {
+ /// The entry block at the time of the traversal. The traversal starts
+ /// from the entry block, so a
diff erent entry block invalidates the cache.
+ Block *entry = nullptr;
+ /// The reachable blocks, mapped to their successors within the region as
+ /// of the last query. The snapshots let the next query tell lost edges
+ /// from added ones.
+ DenseMap<Block *, SmallVector<Block *, 2>> reachable;
+ /// Blocks of `reachable` whose reachability could not be decided
+ /// incrementally. A query for one of them traverses the region. Marking
+ /// and traversing both cost O(blocks), so a lost edge into a loop whose
+ /// target keeps other predecessors costs O(blocks) per rewrite if one of
+ /// the marked blocks is queried before the next change. Bottom-up
+ /// processing has usually popped the operations of these blocks already;
+ /// top-down processing pops them next.
+ llvm::SmallDenseSet<Block *, 4> uncertain;
+ /// Blocks whose terminator changed since the last query, in notification
+ /// order, so that the cache evolves deterministically.
+ llvm::SmallSetVector<Block *, 4> changedBlocks;
+ /// Blocks erased from or moved out of the region since the last query.
+ /// Their pointers are only compared, never dereferenced.
+ llvm::SmallDenseSet<Block *, 4> removedBlocks;
+
+ /// Rebuild the cache from the entry block of `region`.
+ void traverse(Region *region);
+ /// Add `roots` and every block reachable from them to `reachable`. The
+ /// added blocks are uncertain if `isUncertain` is set.
+ void addReachable(SmallVector<Block *> roots, bool isUncertain);
+ /// Mark `block` and every block behind it in the snapshots uncertain.
+ void markUncertain(Block *block);
+ /// Apply the pending changes, if any. On failure, the state is
+ /// inconsistent and the region must be traversed again.
+ LogicalResult applyChanges();
+ };
+
+ /// Record that the terminator of `block` changed. Only tracked regions
+ /// record changes, so no pointer to a block of an untracked region is kept.
+ void markChanged(Block *block) {
+ if (!block)
+ return;
+ auto it = regionStates.find(block->getParent());
+ if (it != regionStates.end())
+ it->second.changedBlocks.insert(block);
+ }
+ /// Record that `block` left `region`.
+ void markRemoved(Region *region, Block *block) {
+ auto it = regionStates.find(region);
+ if (it != regionStates.end())
+ it->second.removedBlocks.insert(block);
+ }
+ bool isBlockReachable(Block *block);
+
+ /// The tracked regions. A region without an entry is traversed by the next
+ /// query that needs it.
+ DenseMap<Region *, RegionState> regionStates;
+
+ Region *const scope;
+};
+
/// This is a worklist-driven driver for the PatternMatcher, which repeatedly
/// applies the locally optimal patterns.
///
@@ -353,7 +456,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
@@ -377,6 +481,10 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
/// `config.strictMode` is GreedyRewriteStrictness::AnyOp.
llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
+ /// Tracks block reachability while the worklist is processed. The rewriter
+ /// notifies this listener, which forwards to the driver.
+ ReachabilityListener reachabilityListener;
+
private:
/// Look over the provided operands for any defining operations that should
/// be re-added to the worklist. This function should be called when an
@@ -414,14 +522,222 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
};
} // namespace
+/// Return whether `op` may define successors of its block.
+static bool mayDefineSuccessors(Operation *op) {
+ return op->mightHaveTrait<OpTrait::IsTerminator>() ||
+ op->getNumSuccessors() != 0;
+}
+
+/// Collect the distinct successors of `block` that are in the same region.
+/// A successor in another region is invalid IR that cannot make
+/// a block of this region reachable.
+static SmallPtrSet<Block *, 2> getRegionSuccessors(Block *block) {
+ SmallPtrSet<Block *, 2> successors;
+ for (Block *successor : block->getSuccessors())
+ if (successor->getParent() == block->getParent())
+ successors.insert(successor);
+ return successors;
+}
+
+void ReachabilityListener::RegionState::traverse(Region *region) {
+ // Clear rather than reassign, to keep the allocations for the next
+ // traversal.
+ reachable.clear();
+ uncertain.clear();
+ changedBlocks.clear();
+ removedBlocks.clear();
+ entry = ®ion->front();
+ addReachable({entry}, /*isUncertain=*/false);
+}
+
+void ReachabilityListener::RegionState::addReachable(SmallVector<Block *> roots,
+ bool isUncertain) {
+ while (!roots.empty()) {
+ auto [it, inserted] = reachable.try_emplace(roots.pop_back_val());
+ if (!inserted)
+ continue;
+ auto successors = getRegionSuccessors(it->first);
+ it->second.append(successors.begin(), successors.end());
+ if (isUncertain)
+ uncertain.insert(it->first);
+ roots.append(it->second);
+ }
+}
+
+void ReachabilityListener::RegionState::markUncertain(Block *block) {
+ SmallVector<Block *> worklist{block};
+ while (!worklist.empty()) {
+ Block *current = worklist.pop_back_val();
+ auto it = reachable.find(current);
+ if (it == reachable.end() || !uncertain.insert(current).second)
+ continue;
+ worklist.append(it->second);
+ }
+}
+
+LogicalResult ReachabilityListener::RegionState::applyChanges() {
+ if (changedBlocks.empty() && removedBlocks.empty())
+ return success();
+ // Blocks are dropped after the loop, so that the snapshot of a lost block
+ // remains available to every changed block that lost an edge to it.
+ SmallVector<Block *> unreachable;
+ for (Block *changed : changedBlocks) {
+ // A removed block is no longer part of the region, and a change in an
+ // uncached block cannot make any block reachable.
+ if (removedBlocks.contains(changed))
+ continue;
+ auto it = reachable.find(changed);
+ if (it == reachable.end())
+ continue;
+ // Keep even a disconnected block's snapshot up to date: a later rewrite
+ // may reconnect it without querying the block itself. Its added successors
+ // must remain uncertain until a traversal establishes their reachability.
+ if (changed != entry && changed->hasNoPredecessors())
+ markUncertain(changed);
+ // Keep the snapshot available while marking lost successors uncertain.
+ // A lost self-edge can make that traversal visit `changed` itself.
+ const auto &old = it->second;
+ auto current = getRegionSuccessors(changed);
+ for (Block *lost : old) {
+ if (current.contains(lost))
+ continue;
+ // The successors of a reachable block are reachable, so `lost` has a
+ // snapshot unless the cache is inconsistent. This is not expected to
+ // happen; the failure is a safety net that forces a traversal.
+ auto lostIt = reachable.find(lost);
+ if (lostIt == reachable.end())
+ return failure();
+ // `lost` and every block behind it may have become unreachable. When
+ // `lost` is gone or has no predecessors left, it is unreachable; if all
+ // of its cached successors are still successors of `changed`, nothing
+ // else changed. This is the shape left by folding a constant conditional
+ // branch and by merging a block into its predecessor. Any other lost
+ // edge leaves the blocks behind `lost` uncertain until a query for one
+ // of them traverses the region.
+ bool isGone = removedBlocks.contains(lost) || lost->hasNoPredecessors();
+ if (isGone && llvm::all_of(lostIt->second, [&](Block *successor) {
+ return current.contains(successor);
+ })) {
+ unreachable.push_back(lost);
+ continue;
+ }
+ markUncertain(lost);
+ }
+ // Added edges can only make more blocks reachable. Blocks reached from an
+ // uncertain block are uncertain as well.
+ SmallVector<Block *> roots;
+ for (Block *successor : current)
+ if (!reachable.contains(successor))
+ roots.push_back(successor);
+ // Nothing was inserted into `reachable` since `it` was looked up.
+ it->second.assign(current.begin(), current.end());
+ addReachable(std::move(roots), uncertain.contains(changed));
+ }
+ // Drop every block that left the region, so that no stale key can match a
+ // block inserted later at the same address.
+ unreachable.append(removedBlocks.begin(), removedBlocks.end());
+ for (Block *block : unreachable) {
+ reachable.erase(block);
+ uncertain.erase(block);
+ }
+ changedBlocks.clear();
+ removedBlocks.clear();
+ return success();
+}
+
+bool ReachabilityListener::isBlockReachable(Block *block) {
+ Region *region = block->getParent();
+ // A detached block is treated as reachable, as is the entry block.
+ if (!region || block == ®ion->front())
+ return true;
+ if (block->hasNoPredecessors())
+ return false;
+ auto [it, inserted] = regionStates.try_emplace(region);
+ RegionState &state = it->second;
+ bool isStale = inserted || state.entry != ®ion->front();
+ if (isStale || failed(state.applyChanges()) ||
+ state.uncertain.contains(block))
+ state.traverse(region);
+ bool result = state.reachable.contains(block);
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ assert(result == region->front().isReachable(block) &&
+ "reachability cache disagrees with a traversal");
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+ return result;
+}
+
+bool ReachabilityListener::isReachable(Operation *op) {
+ // A nested region can be locally reachable while its owning operation is in
+ // an unreachable block, so check every containing block up to the scope.
+ while (Block *block = op->getBlock()) {
+ if (!isBlockReachable(block))
+ return false;
+ Region *region = block->getParent();
+ if (!region || region == scope)
+ break;
+ op = region->getParentOp();
+ }
+ return true;
+}
+
+void ReachabilityListener::notifyBlockInserted(Block *block, Region *previous,
+ Region::iterator previousIt) {
+ Region *region = block->getParent();
+ // A block moved within its region keeps its reachability, unless it becomes
+ // the entry block, which the next query detects. A block that is new to the
+ // region is unreachable until a recorded terminator change adds an edge to
+ // it, but a stale entry from an earlier removal cannot be trusted.
+ if (previous != region) {
+ if (previous)
+ markRemoved(previous, block);
+ auto it = regionStates.find(region);
+ if (it != regionStates.end() && (it->second.removedBlocks.contains(block) ||
+ it->second.reachable.contains(block)))
+ regionStates.erase(it);
+ }
+ RewriterBase::ForwardingListener::notifyBlockInserted(block, previous,
+ previousIt);
+}
+
+void ReachabilityListener::notifyBlockErased(Block *block) {
+ markRemoved(block->getParent(), block);
+ RewriterBase::ForwardingListener::notifyBlockErased(block);
+}
+
+void ReachabilityListener::notifyOperationInserted(
+ Operation *op, OpBuilder::InsertPoint previous) {
+ if (mayDefineSuccessors(op)) {
+ markChanged(op->getBlock());
+ if (previous.isSet())
+ markChanged(previous.getBlock());
+ }
+ RewriterBase::ForwardingListener::notifyOperationInserted(op, previous);
+}
+
+void ReachabilityListener::notifyOperationModified(Operation *op) {
+ if (mayDefineSuccessors(op))
+ markChanged(op->getBlock());
+ RewriterBase::ForwardingListener::notifyOperationModified(op);
+}
+
+void ReachabilityListener::notifyOperationErased(Operation *op) {
+ if (mayDefineSuccessors(op))
+ markChanged(op->getBlock());
+ // Do not retain keys for regions that are about to be destroyed.
+ for (Region ®ion : op->getRegions())
+ regionStates.erase(®ion);
+ RewriterBase::ForwardingListener::notifyOperationErased(op);
+}
+
GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
const GreedyRewriteConfig &config)
- : rewriter(ctx), config(config), matcher(patterns)
+ : rewriter(ctx), config(config),
+ reachabilityListener(this, config.getScope()), matcher(patterns)
#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
// clang-format off
, expensiveChecks(
- /*driver=*/this,
+ /*listener=*/&reachabilityListener,
/*topLevel=*/config.getScope() ? config.getScope()->getParentOp()
: nullptr)
// clang-format on
@@ -432,15 +748,20 @@ GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
// Set up listener.
#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
- // Send IR notifications to the debug handler. This handler will then forward
- // all notifications to this GreedyPatternRewriteDriver.
+ // Chain notifications through the debug handler, reachability listener, and
+ // worklist driver, in that order.
rewriter.setListener(&expensiveChecks);
#else
- rewriter.setListener(this);
+ rewriter.setListener(&reachabilityListener);
#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
}
bool GreedyPatternRewriteDriver::processWorklist() {
+ // The cache is only maintained while this worklist pass runs. Between
+ // passes, the IR changes without notifications, e.g. in region
+ // simplification and in the unreachable-block sweep.
+ llvm::scope_exit clearReachability([&] { reachabilityListener.clear(); });
+
#ifndef NDEBUG
const char *logLineComment =
"//===-------------------------------------------===//\n";
@@ -465,6 +786,11 @@ bool GreedyPatternRewriteDriver::processWorklist() {
(numRewrites < config.getMaxNumRewrites() ||
config.getMaxNumRewrites() == GreedyRewriteConfig::kNoLimit)) {
auto *op = worklist.pop();
+ if (!reachabilityListener.isReachable(op)) {
+ LLVM_DEBUG(logger.startLine() << "Skipping unreachable operation : '"
+ << op->getName() << "'(" << op << ")\n");
+ continue;
+ }
LLVM_DEBUG({
logger.getOStream() << "\n";
@@ -503,8 +829,10 @@ bool GreedyPatternRewriteDriver::processWorklist() {
Operation *dumpRootOp = getDumpRootOp(op);
#endif // NDEBUG
if (foldResults.empty()) {
- // Op was modified in-place.
- notifyOperationModified(op);
+ // Folding modified the op in place without going through the
+ // rewriter.
+ // Notify both the reachability listener and the worklist driver.
+ reachabilityListener.notifyOperationModified(op);
changed = true;
LLVM_DEBUG(logSuccessfulFolding(dumpRootOp));
#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
@@ -854,8 +1182,10 @@ LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
worklist.clear();
// `OperationFolder` CSE's constant ops (and may move them into parents
- // regions to enable more aggressive CSE'ing).
- OperationFolder folder(ctx, this);
+ // regions to enable more aggressive CSE'ing). It must notify the
+ // reachability listener, which forwards to the driver, so that the
+ // reachability cache sees every CFG change.
+ OperationFolder folder(ctx, &reachabilityListener);
auto insertKnownConstant = [&](Operation *op) {
// Check for existing constants when populating the worklist. This avoids
// accidentally reversing the constant order during processing.
@@ -897,7 +1227,10 @@ LogicalResult RegionPatternRewriteDriver::simplify(bool *changed) && {
// %add = arith.addi %add, %add : i64
// are legal in unreachable code. Unfortunately many patterns would be
// unsafe to apply on such IR and can lead to crashes or infinite
- // loops.
+ // loops. A rewrite can disconnect more blocks while processWorklist()
+ // runs; processWorklist() skips their queued operations until the
+ // next sweep, since repeated rewrites on such IR can prevent the
+ // worklist pass from returning.
continueRewrites |=
succeeded(eraseUnreachableBlocks(rewriter, region));
diff --git a/mlir/test/Transforms/canonicalize-cfg-reachability.mlir b/mlir/test/Transforms/canonicalize-cfg-reachability.mlir
new file mode 100644
index 00000000000000..6048601d617c09
--- /dev/null
+++ b/mlir/test/Transforms/canonicalize-cfg-reachability.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-opt %s -split-input-file -pass-pipeline='builtin.module(func.func(canonicalize{max-iterations=1}))' | FileCheck %s
+
+// A single iteration shows which blocks the driver still processes after a
+// rewrite disconnects a block. These tests check that blocks which remain
+// reachable are not skipped; they pass without reachability tracking. The
+// regression tests for rewriting disconnected blocks are in
+// canonicalize-dce.mlir. The constant branches sit behind a loop header, so
+// that the reachability cache is populated before they fold.
+
+// Folding the constant branch drops the only edge into ^dead. The remaining
+// blocks must stay eligible: the addition in ^next folds in the same iteration.
+// CHECK-LABEL: func @lose_only_predecessor
+// CHECK-NOT: arith.addi
+// CHECK: cf.cond_br %arg1, ^{{.*}}(%[[A:.*]] : i32), ^[[EXIT:.*]]
+// CHECK: ^[[EXIT]]:
+// CHECK-NEXT: return %[[A]]
+func.func @lose_only_predecessor(%x: i32, %c: i1) -> i32 {
+ %true = arith.constant true
+ %zero = arith.constant 0 : i32
+ cf.br ^header(%x : i32)
+^header(%h: i32):
+ cf.cond_br %true, ^next(%h : i32), ^dead
+^dead:
+ cf.br ^next(%zero : i32)
+^next(%a: i32):
+ %r = arith.addi %a, %zero : i32
+ cf.cond_br %c, ^header(%r : i32), ^exit
+^exit:
+ return %r : i32
+}
+
+// -----
+
+// The dropped edge leads into a loop, so ^loop keeps a predecessor. The loop
+// is unreachable and must be erased.
+// CHECK-LABEL: func @lose_edge_into_loop
+// CHECK-NOT: arith.addi
+// CHECK: ^[[HEADER:[^:]*]]:
+// CHECK-NEXT: cf.cond_br %arg1, ^[[BODY:[^,]*]], ^[[EXIT:.*]]
+// CHECK: ^[[BODY]]:
+// CHECK-NEXT: cf.br ^[[HEADER]]
+// CHECK: ^[[EXIT]]:
+// CHECK-NEXT: return
+func.func @lose_edge_into_loop(%init: i32, %c: i1) {
+ %true = arith.constant true
+ %c1 = arith.constant 1 : i32
+ cf.br ^header
+^header:
+ cf.cond_br %c, ^body, ^exit
+^body:
+ cf.cond_br %true, ^header, ^loop(%init : i32)
+^loop(%iv: i32):
+ %next = arith.addi %iv, %c1 : i32
+ cf.br ^loop(%next : i32)
+^exit:
+ return
+}
diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir
index 84631947970ded..85fac234d4f3fc 100644
--- a/mlir/test/Transforms/canonicalize-dce.mlir
+++ b/mlir/test/Transforms/canonicalize-dce.mlir
@@ -187,3 +187,149 @@ func.func @f() {
%0 = "test.test_effects_result"() : () -> i32
return
}
+
+// -----
+
+// Test case: Terminate when a block becomes unreachable while patterns run.
+
+// CHECK-LABEL: func @unreachable_loop_terminates
+// CHECK-NEXT: return
+
+func.func @unreachable_loop_terminates(%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
+}
+
+// -----
+
+// A constant branch disconnects a loop with two latches and three block
+// arguments. Folding in the dead loop must not produce an infinite rewrite.
+// Reduced from a Fortran GOTO loop after inlining a constant argument.
+// CHECK-LABEL: func @unreachable_two_latch_loop
+// CHECK-NEXT: %[[C1:.*]] = arith.constant 1 : i32
+// CHECK-NEXT: return %[[C1]]
+func.func @unreachable_two_latch_loop(%p: i1, %n: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %true = arith.constant true
+ cf.cond_br %true, ^bb3(%c1 : i32), ^bb1(%c1, %c1, %c1 : i32, i32, i32)
+^bb1(%a: i32, %b: i32, %c: i32):
+ %j = arith.addi %a, %c1 : i32
+ cf.cond_br %p, ^bb2, ^bb4
+^bb2:
+ %b1 = arith.addi %b, %c1 : i32
+ %gt = arith.cmpi sgt, %b1, %n : i32
+ cf.cond_br %gt, ^bb3(%j : i32), ^bb1(%j, %b1, %c : i32, i32, i32)
+^bb4:
+ %c2 = arith.addi %c, %c1 : i32
+ %le = arith.cmpi sle, %c2, %c0 : i32
+ cf.cond_br %le, ^bb1(%j, %b, %c2 : i32, i32, i32), ^bb3(%j : i32)
+^bb3(%r: i32):
+ return %r : i32
+}
+
+// -----
+
+// The lost edge into ^dead cannot be decided incrementally: ^dead keeps ^pass
+// as a predecessor, which is only reachable from ^dead itself. The blocks
+// behind ^dead hold an SSA cycle through block arguments.
+// CHECK-LABEL: func.func @uncertain_chain
+// CHECK-NOT: arith.addi
+// CHECK: return
+func.func @uncertain_chain(%c: i1, %x: i32) -> i32 {
+ %true = arith.constant true
+ %zero = arith.constant 0 : i32
+ cf.br ^header(%x : i32)
+^header(%h: i32):
+ cf.cond_br %true, ^next(%h : i32), ^dead(%h : i32)
+^pass:
+ cf.br ^dead(%x : i32)
+^dead(%d: i32):
+ %s = arith.addi %d, %d : i32
+ cf.cond_br %c, ^dead2(%s : i32), ^pass
+^dead2(%e: i32):
+ %t = arith.addi %e, %s : i32
+ cf.br ^dead(%t : i32)
+^next(%a: i32):
+ %r = arith.addi %a, %zero : i32
+ cf.cond_br %c, ^header(%r : i32), ^exit
+^exit:
+ return %r : i32
+}
+
+// -----
+
+// Nested regions: a constant scf.if whose region holds a CFG is inlined, so
+// its cached state must be dropped, and a block that becomes unreachable holds
+// nested CFG and scf.if regions whose operations must be skipped.
+// CHECK-LABEL: func.func @nested_two_level
+// CHECK-NEXT: return %arg1
+func.func @nested_two_level(%c: i1, %x: i32) -> i32 {
+ %true = arith.constant true
+ %zero = arith.constant 0 : i32
+ cf.br ^header(%x : i32)
+^header(%h: i32):
+ %v = scf.if %true -> i32 {
+ %w = scf.execute_region -> i32 {
+ cf.br ^a
+ ^a:
+ cf.cond_br %true, ^b(%h : i32), ^c(%h : i32)
+ ^b(%bb: i32):
+ scf.yield %bb : i32
+ ^c(%cc: i32):
+ %q = arith.addi %cc, %cc : i32
+ cf.br ^c(%q : i32)
+ }
+ scf.yield %w : i32
+ } else {
+ scf.yield %zero : i32
+ }
+ cf.cond_br %true, ^exit(%v : i32), ^dead(%v : i32)
+^dead(%d: i32):
+ %r = scf.execute_region -> i32 {
+ cf.br ^l(%d : i32)
+ ^l(%iv: i32):
+ %n = arith.addi %iv, %d : i32
+ %m = scf.if %c -> i32 {
+ %mm = arith.addi %n, %iv : i32
+ scf.yield %mm : i32
+ } else {
+ scf.yield %n : i32
+ }
+ cf.cond_br %c, ^l(%m : i32), ^done(%m : i32)
+ ^done(%z: i32):
+ scf.yield %z : i32
+ }
+ cf.br ^dead(%r : i32)
+^exit(%o: i32):
+ return %o : i32
+}
+
+// -----
+
+// 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..c48a38d2bd54ba
--- /dev/null
+++ b/mlir/test/Transforms/greedy-cfg-reachability.mlir
@@ -0,0 +1,334 @@
+// RUN: mlir-opt %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_cfg_rewrite
+// CHECK: return
+func.func @insert_unreachable() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite append_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_cfg_rewrite
+// CHECK: return
+func.func @connect_before_return(%cond: i1) {
+ cf.br ^start
+^start:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite append_reachable %cond
+ 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_cfg_rewrite insert_entry
+ test.greedy_cfg_rewrite 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_cfg_rewrite move_entry
+ test.greedy_cfg_rewrite observe
+ cf.br ^exit
+^exit:
+ return
+}
+
+// -----
+
+// Moving the old entry away from the front also changes the traversal root.
+// The rewrite inserts a return-only block after the entry, then moves the old
+// entry to the end. No successor changes, but the observer becomes unreachable.
+// CHECK-LABEL: func.func @move_entry_away
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @move_entry_away() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite move_entry_away
+ test.greedy_cfg_rewrite observe
+ 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_cfg_rewrite
+// CHECK: return
+func.func @move_nested_unreachable() {
+ cf.br ^start
+^start:
+ scf.execute_region {
+ scf.execute_region {
+ test.greedy_cfg_rewrite move_parent
+ test.greedy_cfg_rewrite 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_cfg_rewrite
+// CHECK: return
+func.func @disconnect_nested() {
+ cf.br ^body
+^body:
+ scf.execute_region {
+ test.greedy_cfg_rewrite disconnect_parent
+ test.greedy_cfg_rewrite 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_cfg_rewrite
+// CHECK: return
+func.func @redirect() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite redirect
+ cf.br ^body
+^body:
+ test.greedy_cfg_rewrite 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_cfg_rewrite
+// CHECK: return
+func.func @move_then_erase() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite move_erase
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite 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_cfg_rewrite
+// CHECK: return
+func.func @merge() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite merge
+ cf.br ^body
+^body:
+ cf.br ^exit
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ return
+}
+
+// -----
+
+// Merging a block and dropping one of its successors leaves ^dead without
+// predecessors. Its observer must be skipped and ^exit must still be processed.
+// CHECK-LABEL: func.func @merge_disconnect
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: return
+func.func @merge_disconnect(%cond: i1) {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite merge_disconnect
+ cf.br ^body
+^body:
+ cf.cond_br %cond, ^exit, ^dead
+^dead:
+ test.greedy_cfg_rewrite observe
+ return
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ return
+}
+
+// -----
+
+// Erasing the entry block makes its successor the new entry. The cache is
+// populated before the rewrite, and the observer must run in this iteration.
+// CHECK-LABEL: func.func @erase_entry
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: return
+func.func @erase_entry() {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite erase_entry
+ cf.br ^exit
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ return
+}
+
+// -----
+
+// Two blocks lose their edge to ^x in one rewrite. ^x has no predecessors left
+// and ^s, its successor, is still a successor of ^p, but the entry lost ^x as
+// well. Deciding ^s from the edge of ^p alone would keep it cached as reachable
+// and process its observer.
+// CHECK-LABEL: func.func @double_redirect
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @double_redirect() {
+ cf.br ^x
+^x:
+ cf.br ^s
+^p:
+ test.greedy_cfg_rewrite double_redirect
+ cf.br ^x
+^s:
+ test.greedy_cfg_rewrite observe
+ cf.br ^p
+}
+
+// -----
+
+// ^c and ^l change their terminators in one rewrite: ^c branches past ^l to
+// ^s, and ^l, left without predecessors, branches to a new block holding an
+// observer. Applying the change of ^l before that of ^c would record the new
+// block as reachable; the observer in ^s must be processed, the new one not.
+// CHECK-LABEL: func.func @redirect_retarget
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: return
+func.func @redirect_retarget() {
+ cf.br ^c
+^c:
+ test.greedy_cfg_rewrite redirect_retarget
+ cf.br ^l
+^l:
+ cf.br ^s
+^s:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ return
+}
+
+// -----
+
+// An in-place fold that retargets a branch is reported by the driver, not by
+// the rewriter. The disconnected block must be skipped and the block that
+// remains reachable must be processed in this iteration.
+// CHECK-LABEL: func.func @fold_disconnects
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: return
+func.func @fold_disconnects() {
+ cf.br ^start
+^start:
+ "test.br"()[^mid] {fold_to_last} : () -> ()
+^mid:
+ test.greedy_cfg_rewrite observe
+ cf.br ^exit
+^exit:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ return
+}
+
+// -----
+
+// The lost edge leads into a block that keeps a predecessor: itself. The
+// blocks behind it cannot be decided incrementally, and the query for the
+// observer must traverse the region.
+// CHECK-LABEL: func.func @lose_edge_into_loop
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: return
+func.func @lose_edge_into_loop(%cond: i1) {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite redirect
+ cf.br ^loop
+^loop:
+ test.greedy_cfg_rewrite observe
+ cf.cond_br %cond, ^loop, ^exit
+^exit:
+ return
+}
+
+// -----
+
+// The first rewrite disconnects ^body, the second one adds an edge back to
+// it. The observer in ^body is not re-enqueued by the reconnection, but it
+// is still on the worklist and must be processed in this iteration.
+// CHECK-LABEL: func.func @reconnect
+// CHECK-NOT: test.greedy_cfg_rewrite
+// CHECK: cf.cond_br %{{.*}}, ^[[EXIT:.*]], ^[[BODY:.*]]
+// CHECK: ^[[BODY]]:
+// CHECK-NEXT: cf.br ^[[EXIT]]
+// CHECK: ^[[EXIT]]:
+// CHECK-NEXT: return
+func.func @reconnect(%cond: i1) {
+ cf.br ^start
+^start:
+ test.greedy_cfg_rewrite redirect
+ test.greedy_cfg_rewrite reconnect %cond
+ cf.br ^body
+^body:
+ // expected-remark @+1 {{processed reachable block}}
+ test.greedy_cfg_rewrite observe
+ 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..2333507795ebe3 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -10,6 +10,56 @@
// RUN: -test-strict-pattern-driver="strictness=ExistingOps" \
// RUN: --split-input-file %s | FileCheck %s --check-prefix=CHECK-EX
+// Removing a self-edge must propagate uncertainty through the old successor
+// snapshot before removing the entry edge into the loop. Put the tail before
+// the loop so the explicit worklist processes the rewrite before its observer.
+// ExistingOps excludes the new branch, whose query would otherwise rebuild the
+// cache before querying the tail and hide the missing uncertainty propagation.
+// CHECK-EX-LABEL: func.func @disconnect_loop_after_self_edge
+// CHECK-EX: return
+// CHECK-EX: test.greedy_cfg_rewrite observe
+// CHECK-EX-NOT: test.greedy_cfg_rewrite double_redirect
+func.func @disconnect_loop_after_self_edge(%cond: i1) {
+ cf.br ^loop
+^tail:
+ test.greedy_cfg_rewrite observe
+ cf.br ^loop
+^loop:
+ test.greedy_cfg_rewrite double_redirect
+ cf.cond_br %cond, ^loop, ^tail
+}
+
+// -----
+
+// A terminator changes in ^b, then the entry disconnects ^b before the next
+// cache update. Querying ^y applies both pending changes. Reconnecting ^b must
+// discover its new successor ^c, even when ^c is queried before ^b. ExistingOps
+// excludes new branches whose queries could otherwise repair the cache first.
+// Replacing the observer operands explicitly re-enqueues the skipped observers.
+// CHECK-EX-LABEL: func.func @reconnect_stale_successors
+// CHECK-EX-NOT: test.greedy_cfg_rewrite
+// CHECK-EX: {{^ *[}]}}
+func.func @reconnect_stale_successors(%cond: i1) {
+ %for_c = "test.replace_with_new_op"() : () -> i1
+ test.greedy_cfg_rewrite reconnect %cond
+ %for_y = "test.replace_with_new_op"() : () -> i1
+ test.greedy_cfg_rewrite redirect
+ cf.br ^b
+^b:
+ test.greedy_cfg_rewrite reconnect %cond
+ cf.br ^x
+^c:
+ test.greedy_cfg_rewrite observe %for_c {tag = "C"}
+ return
+^x:
+ return
+^y:
+ test.greedy_cfg_rewrite observe %for_y {tag = "Y"}
+ return
+}
+
+// -----
+
// CHECK-EN-LABEL: func @test_erase
// CHECK-EN-SAME: pattern_driver_all_erased = true, pattern_driver_changed = true}
// CHECK-EN: "test.arg0"
@@ -75,8 +125,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 +141,55 @@ 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"()
+// CHECK-EX-LABEL: func @test_skip_unreachable
+// CHECK-EX-SAME: pattern_driver_all_erased = false, pattern_driver_changed = false
+// CHECK-EX: "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_cfg_rewrite 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/TestEnumDefs.td b/mlir/test/lib/Dialect/Test/TestEnumDefs.td
index c98be29a048963..daf898ca7ee061 100644
--- a/mlir/test/lib/Dialect/Test/TestEnumDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestEnumDefs.td
@@ -106,6 +106,34 @@ def TestSimpleEnum64 : I64Enum<"SimpleEnum64", "", [
let cppNamespace = "::test";
}
+//===----------------------------------------------------------------------===//
+// Greedy CFG Rewrite Mode
+//===----------------------------------------------------------------------===//
+
+// The CFG change performed by `test.greedy_cfg_rewrite`. The cases are
+// documented on the op.
+def TestGreedyCfgRewriteMode
+ : I32Enum<"GreedyCfgRewriteMode", "a CFG change during a greedy rewrite", [
+ I32EnumCase<"Observe", 0, "observe">,
+ I32EnumCase<"AppendUnreachable", 1, "append_unreachable">,
+ I32EnumCase<"AppendReachable", 2, "append_reachable">,
+ I32EnumCase<"InsertEntry", 3, "insert_entry">,
+ I32EnumCase<"MoveEntry", 4, "move_entry">,
+ I32EnumCase<"MoveEntryAway", 5, "move_entry_away">,
+ I32EnumCase<"EraseEntry", 6, "erase_entry">,
+ I32EnumCase<"Redirect", 7, "redirect">,
+ I32EnumCase<"DoubleRedirect", 8, "double_redirect">,
+ I32EnumCase<"Reconnect", 9, "reconnect">,
+ I32EnumCase<"DisconnectParent", 10, "disconnect_parent">,
+ I32EnumCase<"MoveParent", 11, "move_parent">,
+ I32EnumCase<"MoveErase", 12, "move_erase">,
+ I32EnumCase<"Merge", 13, "merge">,
+ I32EnumCase<"MergeDisconnect", 14, "merge_disconnect">,
+ I32EnumCase<"RedirectRetarget", 15, "redirect_retarget">,
+ ]> {
+ let cppNamespace = "::test";
+}
+
//===----------------------------------------------------------------------===//
// Test Bit Enum
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 57cd56f70fac0f..ef1fc725b53cda 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -59,6 +59,15 @@ SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
return SuccessorOperands(getTargetOperandsMutable());
}
+LogicalResult TestBranchOp::fold(FoldAdaptor adaptor,
+ SmallVectorImpl<OpFoldResult> &results) {
+ if (!getFoldToLast())
+ return failure();
+ (*this)->setSuccessor(&(*this)->getParentRegion()->back(), 0);
+ removeFoldToLastAttr();
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// TestProducingBranchOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 0a64885d8e2921..fc096a9e61d710 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1133,10 +1133,55 @@ def UpdateAttr : Pat<(I32ElementsAttrOp $attr),
(I32ElementsAttrOp ConstantAttr<I32ElementsAttr, "0">),
[(IsNotScalar $attr)]>;
+def TestGreedyCfgRewriteOp : TEST_Op<"greedy_cfg_rewrite"> {
+ let description = [{
+ Exercises reachability changes while the greedy driver is processing its
+ worklist. The matching test pattern performs the CFG change selected by
+ `mode`, then erases the op:
+
+ observe: emit a remark; used to observe whether a block is
+ processed
+ append_unreachable: append a disconnected block holding an observer
+ append_reachable: append a block holding an observer and branch to it
+ on `cond`
+ insert_entry: insert a new entry block, disconnecting the old one
+ move_entry: move the last block to the front
+ move_entry_away: insert a block after the entry and move the entry to
+ the end
+ erase_entry: erase the entry block; its successor becomes the entry
+ redirect: retarget the block's branch to the last block
+ double_redirect: retarget the block's branch to the last block and
+ replace the entry's branch with a return; for a
+ conditional branch, keep only the false destination
+ reconnect: replace the block's branch with a conditional branch
+ on `cond` to its destination and to the next block
+ disconnect_parent: retarget the entry branch of the parent region
+ move_parent: move the grandparent operation into a new block
+ move_erase: retarget the block's branch, then move its old
+ destination to another region and erase it
+ merge: merge the branch destination into the block
+ merge_disconnect: merge, then fold the merged conditional branch
+ redirect_retarget: retarget the block's branch to its destination's
+ destination, and the destination's branch to an
+ appended block holding an observer
+
+ Modes that operate on the block's branch expect an unconditional branch
+ without operands as the terminator, except that `double_redirect` also
+ accepts a conditional branch without successor operands.
+ }];
+ let arguments = (ins EnumProp<TestGreedyCfgRewriteMode>:$mode,
+ Optional<I1>:$cond);
+ let assemblyFormat = "$mode ($cond^)? attr-dict";
+}
+
def TestBranchOp : TEST_Op<"br",
[DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {
- let arguments = (ins Variadic<AnyType>:$targetOperands);
+ let arguments = (ins Variadic<AnyType>:$targetOperands,
+ UnitAttr:$fold_to_last);
let successors = (successor AnySuccessor:$target);
+ // With `fold_to_last`, the fold retargets the branch in place to the last
+ // block of the region, to exercise CFG changes made by folding.
+ let hasFolder = 1;
}
def TestProducingBranchOp : TEST_Op<"producing_br",
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 7072977b14de26..0417d6662814e0 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -12,6 +12,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"
@@ -470,6 +471,159 @@ class CreateAndEraseOpAndBlock : public RewritePattern {
}
};
+/// Perform the CFG change selected by the `mode` of a `test.greedy_cfg_rewrite`
+/// op while a greedy rewrite is processing its worklist. The modes are
+/// documented on the op.
+class GreedyCfgRewritePattern
+ : public OpRewritePattern<TestGreedyCfgRewriteOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(TestGreedyCfgRewriteOp op,
+ PatternRewriter &rewriter) const override {
+ Region *region = op->getParentRegion();
+ Location loc = op.getLoc();
+ // The unconditional branch that terminates the block of `op`.
+ auto getBranch = [&] {
+ return cast<cf::BranchOp>(op->getBlock()->getTerminator());
+ };
+ switch (op.getMode()) {
+ case GreedyCfgRewriteMode::Observe:
+ op.emitRemark("processed reachable block");
+ break;
+ case GreedyCfgRewriteMode::AppendUnreachable:
+ case GreedyCfgRewriteMode::AppendReachable: {
+ Block *newBlock = rewriter.createBlock(region);
+ TestGreedyCfgRewriteOp::create(rewriter, loc,
+ GreedyCfgRewriteMode::Observe,
+ /*cond=*/Value());
+ func::ReturnOp::create(rewriter, loc);
+ if (op.getMode() == GreedyCfgRewriteMode::AppendReachable) {
+ cf::BranchOp branch = getBranch();
+ rewriter.setInsertionPoint(branch);
+ cf::CondBranchOp::create(rewriter, loc, op.getCond(), branch.getDest(),
+ ValueRange(), newBlock, ValueRange());
+ rewriter.eraseOp(branch);
+ }
+ break;
+ }
+ case GreedyCfgRewriteMode::InsertEntry:
+ rewriter.createBlock(region, region->begin());
+ func::ReturnOp::create(rewriter, loc);
+ break;
+ case GreedyCfgRewriteMode::MoveEntry: {
+ // The new entry must not have predecessors after the rewrite.
+ cf::BranchOp branch = getBranch();
+ rewriter.setInsertionPoint(branch);
+ rewriter.replaceOpWithNewOp<func::ReturnOp>(branch);
+ rewriter.moveBlockBefore(®ion->back(), region, region->begin());
+ break;
+ }
+ case GreedyCfgRewriteMode::MoveEntryAway: {
+ Block *entry = ®ion->front();
+ rewriter.createBlock(region, std::next(region->begin()));
+ func::ReturnOp::create(rewriter, loc);
+ rewriter.moveBlockBefore(entry, region, region->end());
+ break;
+ }
+ case GreedyCfgRewriteMode::EraseEntry:
+ rewriter.eraseBlock(®ion->front());
+ break;
+ case GreedyCfgRewriteMode::Redirect:
+ case GreedyCfgRewriteMode::DoubleRedirect: {
+ if (auto conditional =
+ dyn_cast<cf::CondBranchOp>(op->getBlock()->getTerminator())) {
+ assert(op.getMode() == GreedyCfgRewriteMode::DoubleRedirect);
+ rewriter.setInsertionPoint(conditional);
+ cf::BranchOp::create(rewriter, loc, conditional.getFalseDest());
+ rewriter.eraseOp(conditional);
+ } else {
+ cf::BranchOp branch = getBranch();
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(®ion->back()); });
+ }
+ if (op.getMode() == GreedyCfgRewriteMode::DoubleRedirect) {
+ Operation *entryBranch = region->front().getTerminator();
+ rewriter.setInsertionPoint(entryBranch);
+ func::ReturnOp::create(rewriter, loc);
+ rewriter.eraseOp(entryBranch);
+ }
+ break;
+ }
+ case GreedyCfgRewriteMode::Reconnect: {
+ cf::BranchOp branch = getBranch();
+ rewriter.setInsertionPoint(branch);
+ cf::CondBranchOp::create(
+ rewriter, loc, op.getCond(), branch.getDest(), ValueRange(),
+ &*std::next(op->getBlock()->getIterator()), ValueRange());
+ rewriter.eraseOp(branch);
+ break;
+ }
+ case GreedyCfgRewriteMode::DisconnectParent: {
+ Region *parentRegion = region->getParentOp()->getParentRegion();
+ auto branch = cast<cf::BranchOp>(parentRegion->front().getTerminator());
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(&parentRegion->back()); });
+ break;
+ }
+ case GreedyCfgRewriteMode::MoveParent: {
+ Operation *outer = region->getParentOp()->getParentOp();
+ Block *newBlock = rewriter.createBlock(outer->getParentRegion());
+ func::ReturnOp::create(rewriter, loc);
+ rewriter.moveOpBefore(outer, newBlock, newBlock->begin());
+ break;
+ }
+ case GreedyCfgRewriteMode::MoveErase: {
+ cf::BranchOp branch = getBranch();
+ Block *moved = branch.getDest();
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(®ion->back()); });
+ // The block is erased while it belongs to another region. The source
+ // region keeps its entry and several blocks, so its cached state must
+ // not retain a pointer to the erased block.
+ Region temporary;
+ rewriter.moveBlockBefore(moved, &temporary, temporary.end());
+ rewriter.eraseBlock(moved);
+ break;
+ }
+ case GreedyCfgRewriteMode::RedirectRetarget: {
+ // Two terminators change in one rewrite. The old destination has no
+ // predecessors left, so its retargeted branch must not make the new
+ // block reachable, whichever block the cache processes first.
+ cf::BranchOp branch = getBranch();
+ auto lostBranch = cast<cf::BranchOp>(branch.getDest()->getTerminator());
+ Block *newBlock = rewriter.createBlock(region);
+ TestGreedyCfgRewriteOp::create(rewriter, loc,
+ GreedyCfgRewriteMode::Observe,
+ /*cond=*/Value());
+ func::ReturnOp::create(rewriter, loc);
+ rewriter.modifyOpInPlace(branch,
+ [&] { branch.setDest(lostBranch.getDest()); });
+ rewriter.modifyOpInPlace(lostBranch,
+ [&] { lostBranch.setDest(newBlock); });
+ break;
+ }
+ case GreedyCfgRewriteMode::Merge:
+ case GreedyCfgRewriteMode::MergeDisconnect: {
+ Block *dest = op->getBlock();
+ cf::BranchOp branch = getBranch();
+ Block *source = branch.getDest();
+ rewriter.eraseOp(branch);
+ rewriter.mergeBlocks(source, dest);
+ if (op.getMode() == GreedyCfgRewriteMode::MergeDisconnect) {
+ auto conditional = cast<cf::CondBranchOp>(dest->getTerminator());
+ rewriter.setInsertionPoint(conditional);
+ cf::BranchOp::create(rewriter, loc, conditional.getTrueDest());
+ rewriter.eraseOp(conditional);
+ }
+ break;
+ }
+ }
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
struct TestGreedyPatternDriver
: public PassWrapper<TestGreedyPatternDriver, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGreedyPatternDriver)
@@ -488,7 +642,7 @@ struct TestGreedyPatternDriver
patterns.add<FoldingPattern, TestNamedPatternRule,
FolderInsertBeforePreviouslyFoldedConstantPattern,
FolderCommutativeOp2WithConstant, HoistEligibleOps,
- MakeOpEligible>(&getContext());
+ MakeOpEligible, GreedyCfgRewritePattern>(&getContext());
// Additional patterns for testing the GreedyPatternRewriteDriver.
patterns.insert<IncrementIntAttribute<3>>(&getContext());
@@ -544,6 +698,7 @@ struct TestStrictPatternDriver
patterns.add<
// clang-format off
ChangeBlockOp,
+ GreedyCfgRewritePattern,
CloneOp,
CloneRegionBeforeOp,
EraseOp,
@@ -563,7 +718,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_cfg_rewrite") {
ops.push_back(op);
}
});
@@ -586,8 +742,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