Add a PatternRewriter hook to merge blocks, and use it to support for folding branches.

A pattern rewriter hook, mergeBlock, is added that allows for merging the operations of one block into the end of another. This is used to support a canonicalization pattern for branch operations that folds the branch when the successor has a single predecessor(the branch block).

Example:
  ^bb0:
    %c0_i32 = constant 0 : i32
    br ^bb1(%c0_i32 : i32)
  ^bb1(%x : i32):
    return %x : i32

becomes:
  ^bb0:
    %c0_i32 = constant 0 : i32
    return %c0_i32 : i32
PiperOrigin-RevId: 278677825

1 files changed, 29 insertions, 0 deletions

diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 1f9b9b060aa..b8ecab97552 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -109,6 +109,35 @@ void PatternRewriter::eraseOp(Operation *op) {
   op->erase();
 }
 
+/// Merge the operations of block 'source' into the end of block 'dest'.
+/// 'source's predecessors must be empty or only contain 'dest`.
+/// 'argValues' is used to replace the block arguments of 'source' after
+/// merging.
+void PatternRewriter::mergeBlocks(Block *source, Block *dest,
+                                  ArrayRef<Value *> argValues) {
+  assert(llvm::all_of(source->getPredecessors(),
+                      [dest](Block *succ) { return succ == dest; }) &&
+         "expected 'source' to have no predecessors or only 'dest'");
+  assert(argValues.size() == source->getNumArguments() &&
+         "incorrect # of argument replacement values");
+
+  // Replace all of the successor arguments with the provided values.
+  for (auto it : llvm::zip(source->getArguments(), argValues))
+    std::get<0>(it)->replaceAllUsesWith(std::get<1>(it));
+
+  // Splice the operations of the 'source' block into the 'dest' block and erase
+  // it.
+  dest->getOperations().splice(dest->end(), source->getOperations());
+  source->dropAllUses();
+  source->erase();
+}
+
+/// Split the operations starting at "before" (inclusive) out of the given
+/// block into a new block, and return it.
+Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
+  return block->splitBlock(before);
+}
+
 /// op and newOp are known to have the same number of results, replace the
 /// uses of op with uses of newOp
 void PatternRewriter::replaceOpWithResultsOfAnotherOp(