[Mlir-commits] [mlir] [mlir] Avoid rewriting unreachable blocks in the greedy driver (PR #225814)

Mehdi Amini llvmlistbot at llvm.org
Fri Sep 25 05:02:52 PDT 2026


https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/225814

>From e9d3cd1dad96e6b888e9e10baa9d1d1b56b1e2a4 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Fri, 4 Sep 2026 00:01:22 -0700
Subject: [PATCH] [mlir] Avoid rewriting unreachable blocks in the greedy
 driver

A rewrite can disconnect a block after the iteration's initial CFG sweep.
Track reachability through rewriter notifications and skip worklist
operations in unreachable blocks until the next iteration removes them.

Cache reachability per region. Compare the final successors of changed
blocks after each rewrite, preserving the cache for successor-equivalent
changes and block merges that keep surviving blocks reachable. Recompute for
other CFG changes. This avoids repeated full CFG scans during branch
canonicalization.

Single-block regions return before cache lookup or CFG traversal. Before the
cache bookkeeping simplification, in the same Release build, canonicalizing
32,000 sibling single-block regions took 22.8 ms versus 21.9 ms with the
parent driver. With 16,000 regions containing folds, the times were 35.7 ms
and 34.3 ms, respectively. Nesting to 256 levels showed no measurable
difference.

Before that simplification, on a synthetic CFG chain whose conditional
branches have identical successors, median CanonicalizerPass times with
region simplification disabled were 0.3, 0.7, 1.3, and 2.5 ms in the parent
commit, versus 0.6, 1.2, 2.5, and 5.0 ms with this commit, for 500, 1,000,
2,000, and 4,000 blocks respectively. Both used the same Release build.

Cover inserted and redirected blocks, blocks connected before rewrite
completion, nested regions, and explicit op-list behavior. Exercise updates
to populated caches and observe whether affected operations are rewritten.
Invalidate the source cache on cross-region block moves so a subsequent
erasure cannot leave a dangling pointer. Share the traversal, use the
successor map as the reachable set, and use set membership for successors.

Validate connected blocks within a single greedy iteration. Avoid snapshots
of uncached blocks and redundant erased-block validation. Document that
skipped operations are not automatically retried if a later rewrite makes
their block reachable.

Assisted-by: Codex
---
 .../Transforms/GreedyPatternRewriteDriver.h   |   6 +
 .../Utils/GreedyPatternRewriteDriver.cpp      | 190 +++++++++++++++++-
 mlir/test/Transforms/canonicalize-dce.mlir    |  41 ++++
 .../Transforms/greedy-cfg-reachability.mlir   |  77 +++++++
 .../test-strict-pattern-driver.mlir           |  20 +-
 mlir/test/lib/Dialect/Test/TestPatterns.cpp   |  65 +++++-
 6 files changed, 395 insertions(+), 4 deletions(-)
 create mode 100644 mlir/test/Transforms/greedy-cfg-reachability.mlir

diff --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index 7fef17d806ef3b..7627f510f36c11 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.,
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba06..9cb20e9a4c2ed2 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -27,6 +27,7 @@
 #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"
@@ -353,7 +354,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 +380,26 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
   llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
 
 private:
+  struct ReachabilityState {
+    // Remember the root of the cached traversal. Inserting or moving a block
+    // before it changes reachability without changing any existing successor.
+    Block *entry = nullptr;
+    // Keys are the reachable blocks; values are their cached 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;
+
+  void invalidateRegion(Region *region);
+  void markBlockChanged(Region *region, Block *block);
+  void markBlockErased(Block *block);
+  bool isReachable(Operation *op);
+
   /// Look over the provided operands for any defining operations that should
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
@@ -441,6 +463,9 @@ GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
 }
 
 bool GreedyPatternRewriteDriver::processWorklist() {
+  reachableBlocks.clear();
+  llvm::scope_exit clearReachability([&] { reachableBlocks.clear(); });
+
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
@@ -465,6 +490,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 +667,139 @@ bool GreedyPatternRewriteDriver::processWorklist() {
   return changed;
 }
 
+void GreedyPatternRewriteDriver::invalidateRegion(Region *region) {
+  if (region)
+    reachableBlocks.erase(region);
+}
+
+void GreedyPatternRewriteDriver::markBlockChanged(Region *region,
+                                                  Block *block) {
+  if (region) {
+    auto it = reachableBlocks.find(region);
+    if (it != reachableBlocks.end())
+      it->second.changedBlocks.insert(block);
+  }
+}
+
+void GreedyPatternRewriteDriver::markBlockErased(Block *block) {
+  Region *region = block->getParent();
+  if (!region)
+    return;
+  auto it = reachableBlocks.find(region);
+  if (it != reachableBlocks.end()) {
+    if (it->second.entry == block) {
+      reachableBlocks.erase(it);
+      return;
+    }
+    it->second.changedBlocks.insert(block);
+    it->second.erasedBlocks.insert(block);
+  }
+}
+
+static SmallVector<Block *, 2> getUniqueSuccessors(Block *block) {
+  SmallVector<Block *, 2> successors;
+  llvm::SmallPtrSet<Block *, 4> seen;
+  for (Block *successor :
+       llvm::make_range(block->succ_begin(), block->succ_end()))
+    if (seen.insert(successor).second)
+      successors.push_back(successor);
+  return successors;
+}
+
+bool GreedyPatternRewriteDriver::isReachable(Operation *op) {
+  // A detached operation or block has no region entry to compare against.
+  Block *block = op->getBlock();
+  if (!block)
+    return true;
+  Region *region = block->getParent();
+  if (!region || region->hasOneBlock())
+    return true;
+
+  auto [it, inserted] = reachableBlocks.try_emplace(region);
+  ReachabilityState &state = it->second;
+  // Extend the cache from these roots, stopping at already visited blocks.
+  auto visit = [&](SmallVector<Block *> worklist) {
+    while (!worklist.empty()) {
+      Block *reachable = worklist.pop_back_val();
+      auto [it, inserted] = state.successors.try_emplace(reachable);
+      if (!inserted)
+        continue;
+      it->second = getUniqueSuccessors(reachable);
+      worklist.append(it->second);
+    }
+  };
+  auto recompute = [&] {
+    state = ReachabilityState();
+    state.entry = &region->front();
+    visit({state.entry});
+  };
+  // A changed region entry invalidates the traversal rooted at state.entry,
+  // even if no existing block changed its successors.
+  if (inserted || state.entry != &region->front()) {
+    recompute();
+  } else if (!state.changedBlocks.empty()) {
+    // Check the final CFG, after all notifications from the rewrite. Replacing
+    // a terminator without changing its successors needs no traversal. A
+    // block merge contracts an erased successor into its predecessor: if the
+    // predecessor now reaches all of the erased block's old successors, every
+    // path through that block still reaches the same surviving blocks.
+    DenseMap<Block *, SmallVector<Block *, 2>> currentSuccessors;
+    for (Block *changed : state.changedBlocks) {
+      // Cross-region moves invalidate the source cache immediately.
+      // Newly reachable blocks will have their successors read by visit().
+      if (state.successors.contains(changed) &&
+          !state.erasedBlocks.contains(changed))
+        currentSuccessors[changed] = getUniqueSuccessors(changed);
+    }
+
+    // On any old path from the entry to an erased block, the first erased
+    // block has a surviving predecessor. Checking that predecessor's lost edge
+    // either forces a rescan or preserves all of the erased block's successors.
+    // No separate check of erased blocks is needed.
+    bool mustRecompute = false;
+    for (auto &entry : currentSuccessors) {
+      Block *changed = entry.first;
+      auto &current = entry.second;
+      auto oldChanged = state.successors.find(changed);
+      llvm::SmallPtrSet<Block *, 4> currentSet(current.begin(), current.end());
+      for (Block *oldSuccessor : oldChanged->second) {
+        if (currentSet.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 = state.successors.find(oldSuccessor);
+        if (!state.erasedBlocks.contains(oldSuccessor) ||
+            old == state.successors.end() ||
+            !llvm::all_of(old->second, [&](Block *successor) {
+              return currentSet.contains(successor);
+            })) {
+          mustRecompute = true;
+          break;
+        }
+      }
+      if (mustRecompute)
+        break;
+    }
+    if (mustRecompute) {
+      recompute();
+    } else {
+      // Additions can only make more blocks reachable. Visit those blocks
+      // without rescanning the already reachable part of the region.
+      for (Block *erased : state.erasedBlocks)
+        state.successors.erase(erased);
+      SmallVector<Block *> worklist;
+      for (auto &[changed, current] : currentSuccessors) {
+        state.successors[changed] = current;
+        worklist.append(current);
+      }
+      visit(std::move(worklist));
+      state.changedBlocks.clear();
+      state.erasedBlocks.clear();
+    }
+  }
+  return state.successors.contains(block);
+}
+
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
   assert(op && "expected valid op");
   // Gather potential ancestors while looking for a "scope" parent region.
@@ -665,17 +827,34 @@ 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);
+  if (auto it = reachableBlocks.find(region);
+      it != reachableBlocks.end() && it->second.erasedBlocks.contains(block))
+    invalidateRegion(region);
+  markBlockChanged(region, 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 (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors()) {
+    markBlockChanged(op->getParentRegion(), op->getBlock());
+    if (previous.isSet())
+      markBlockChanged(previous.getBlock()->getParent(), previous.getBlock());
+  }
+  for (Region &region : op->getRegions())
+    invalidateRegion(&region);
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
@@ -688,6 +867,10 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(
 }
 
 void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
+  if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+    markBlockChanged(op->getParentRegion(), op->getBlock());
+  for (Region &region : op->getRegions())
+    invalidateRegion(&region);
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
                        << ")\n";
@@ -731,6 +914,11 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
 }
 
 void GreedyPatternRewriteDriver::notifyOperationErased(Operation *op) {
+  if (op->mightHaveTrait<OpTrait::IsTerminator>() || op->getNumSuccessors())
+    markBlockChanged(op->getParentRegion(), op->getBlock());
+  for (Region &region : op->getRegions()) {
+    reachableBlocks.erase(&region);
+  }
   LLVM_DEBUG({
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir
index 84631947970ded..e73e2b134b1941 100644
--- a/mlir/test/Transforms/canonicalize-dce.mlir
+++ b/mlir/test/Transforms/canonicalize-dce.mlir
@@ -187,3 +187,44 @@ func.func @f() {
   %0 = "test.test_effects_result"() : () -> i32
   return
 }
+
+// -----
+
+// Test case: Terminate when a block becomes unreachable while patterns run.
+
+// CHECK:      func @f
+// CHECK-NEXT:   return
+
+func.func @f(%cond: i1, %init: i32) {
+  %true = arith.constant true
+  %c1 = arith.constant 1 : i32
+  cf.cond_br %true, ^exit, ^header(%init : i32)
+^header(%iv: i32):
+  cf.cond_br %cond, ^exit, ^latch
+^latch:
+  %next = arith.addi %iv, %c1 : i32
+  cf.br ^header(%next : i32)
+^exit:
+  return
+}
+
+// -----
+
+// The same CFG change can occur in a region nested below the driver root.
+// CHECK-LABEL: func.func @nested_unreachable
+// CHECK-NOT: arith.addi
+// CHECK: return
+func.func @nested_unreachable(%cond: i1, %init: i32) {
+  scf.execute_region {
+    %true = arith.constant true
+    cf.cond_br %true, ^exit, ^header(%init : i32)
+  ^header(%iv: i32):
+    cf.cond_br %cond, ^exit, ^latch
+  ^latch:
+    %next = arith.addi %iv, %iv : i32
+    cf.br ^header(%next : i32)
+  ^exit:
+    scf.yield
+  }
+  return
+}
diff --git a/mlir/test/Transforms/greedy-cfg-reachability.mlir b/mlir/test/Transforms/greedy-cfg-reachability.mlir
new file mode 100644
index 00000000000000..ae5c9f27b2ab35
--- /dev/null
+++ b/mlir/test/Transforms/greedy-cfg-reachability.mlir
@@ -0,0 +1,77 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -test-greedy-patterns="top-down max-iterations=1" -verify-diagnostics | FileCheck %s
+
+// A rewrite can create an unreachable block after the iteration's initial
+// unreachable-block sweep. Do not process its self-referential operation.
+// CHECK-LABEL: func.func @insert_unreachable
+// CHECK-NOT: arith.addi
+// CHECK: return
+func.func @insert_unreachable(%arg: i32) {
+  "test.greedy_create_block"(%arg) {mode = "unreachable"} : (i32) -> ()
+  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) {
+  // 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() {
+  "test.greedy_create_block"() {mode = "entry"} : () -> ()
+  "test.greedy_create_block"() {mode = "observe"} : () -> ()
+  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() {
+  "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() {
+  "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
+}
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index c87444cba8e1a2..8a4735b3f2c894 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -75,8 +75,8 @@ func.func @test_replace_with_erase_op() {
 //       CHECK-AN: return
 //       CHECK-AN: ^[[BB1:[^:]*]]:
 //       CHECK-AN: "test.implicit_change_op"()[^[[BB1]]]
-func.func @test_trigger_rewrite_through_block() {
-  return
+func.func @test_trigger_rewrite_through_block(%cond: i1) {
+  cf.cond_br %cond, ^bb1, ^bb3
 ^bb1:
   // Uses bb1. ChangeBlockOp replaces that and all other usages of bb1 with bb2.
   "test.change_block_op"() [^bb1, ^bb2] : () -> ()
@@ -91,6 +91,22 @@ func.func @test_trigger_rewrite_through_block() {
 
 // -----
 
+// Explicitly selected operations in unreachable blocks are skipped.
+// CHECK-AN-LABEL: func @test_skip_unreachable
+// CHECK-AN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = false
+// CHECK-AN: "test.erase_op"()
+// CHECK-EN-LABEL: func @test_skip_unreachable
+// CHECK-EN-SAME: pattern_driver_all_erased = false, pattern_driver_changed = false
+// CHECK-EN: "test.erase_op"()
+func.func @test_skip_unreachable() {
+  return
+^dead:
+  "test.erase_op"() : () -> ()
+  return
+}
+
+// -----
+
 // CHECK-AN: notifyOperationErased: test.foo_b
 // CHECK-AN: notifyOperationErased: test.foo_a
 // CHECK-AN: notifyOperationErased: test.graph_region
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 06517910e270e8..78db40ca6d227d 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,68 @@ class CreateAndEraseOpAndBlock : public RewritePattern {
   }
 };
 
+/// Exercise reachability changes while a greedy rewrite is processing its
+/// worklist. A new block is either disconnected, connected within the same
+/// rewrite, or inserted before the old entry block; another mode redirects a
+/// branch in-place.
+class CreateBlockDuringGreedyRewrite : public RewritePattern {
+public:
+  CreateBlockDuringGreedyRewrite(MLIRContext *context)
+      : RewritePattern("test.greedy_create_block", /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    StringRef mode = cast<StringAttr>(op->getAttr("mode")).getValue();
+    Region *region = op->getParentRegion();
+    Location loc = op->getLoc();
+    if (mode == "observe") {
+      op->emitRemark("processed reachable block");
+    } else if (mode == "move-erase") {
+      auto branch = cast<cf::BranchOp>(op->getBlock()->getTerminator());
+      Block *moved = branch.getDest();
+      rewriter.modifyOpInPlace(branch,
+                               [&] { branch.setDest(&region->back()); });
+      // The source region retains two 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 == "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(&region->back()); });
+    } else {
+      Block *newBlock = rewriter.createBlock(region);
+      if (mode == "unreachable") {
+        Value input = op->getOperand(0);
+        auto add = arith::AddIOp::create(rewriter, loc, input, input);
+        rewriter.modifyOpInPlace(add, [&] {
+          add->setOperand(0, add.getResult());
+          add->setOperand(1, add.getResult());
+        });
+        cf::BranchOp::create(rewriter, loc, newBlock);
+      } else {
+        assert(mode == "reachable");
+        OperationState state(loc, "test.greedy_create_block");
+        state.addAttribute("mode", rewriter.getStringAttr("observe"));
+        rewriter.create(state);
+        func::ReturnOp::create(rewriter, loc);
+        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 +549,7 @@ struct TestGreedyPatternDriver
     patterns.add<FoldingPattern, TestNamedPatternRule,
                  FolderInsertBeforePreviouslyFoldedConstantPattern,
                  FolderCommutativeOp2WithConstant, HoistEligibleOps,
-                 MakeOpEligible>(&getContext());
+                 MakeOpEligible, CreateBlockDuringGreedyRewrite>(&getContext());
 
     // Additional patterns for testing the GreedyPatternRewriteDriver.
     patterns.insert<IncrementIntAttribute<3>>(&getContext());



More information about the Mlir-commits mailing list