[MLIR] Add PatternRewriter::mergeBlockBefore() to merge a block in the middle of another block.

[jurahul]

Summary

• This utility to merge a block into anywhere into another one can help inline single block regions into other blocks.
• Modified patterns test to use the new function.

[rriddle]

Thanks!

mlir/include/mlir/IR/Block.h

pred_iterator pred_begin() {
     return pred_iterator((BlockOperand *)getFirstUse());
   }
   pred_iterator pred_end() { return pred_iterator(nullptr); }
   iterator_range<pred_iterator> getPredecessors() {
     return {pred_begin(), pred_end()};
   }
 
   /// Return true if this block has no predecessors.
-  bool hasNoPredecessors();
+  bool hasNoPredecessors() { return pred_begin() == pred_end(); }
+
+  /// Returns true if this blocks has no successors.
+  bool hasNoSuccessors() { return succ_begin() == succ_end(); }
 
   /// If this block has exactly one predecessor, return it.  Otherwise, return
   /// null.
   ///
   /// Note that if a block has duplicate predecessors from a single block (e.g.
   /// if you have a conditional branch with the same block as the true/false
   /// destinations) is not considered to be a single predecessor.
   Block *getSinglePredecessor();

mlir/include/mlir/IR/PatternMatch.h

public:
 
   /// Merge the operations of block 'source' into the end of block 'dest'.
   /// 'source's predecessors must either be empty or only contain 'dest`.
   /// 'argValues' is used to replace the block arguments of 'source' after
   /// merging.
   virtual void mergeBlocks(Block *source, Block *dest,
                            ValueRange argValues = llvm::None);
 
+  // Merge the operations of block 'source' before the operation 'op'. Source
+  // block should not have existing predecessors or successors.
+  void mergeBlockBefore(Block *source, Operation *op,
+                        ValueRange argValues = llvm::None);
+
   /// Split the operations starting at "before" (inclusive) out of the given
   /// block into a new block, and return it.
   virtual Block *splitBlock(Block *block, Block::iterator before);
 
   /// This method is used to notify the rewriter that an in-place operation
   /// modification is about to happen. A call to this function *must* be
   /// followed by a call to either `finalizeRootUpdate` or `cancelRootUpdate`.
   /// This is a minor efficiency win (it avoids creating a new operation and

mlir/lib/IR/Block.cpp

 
 /// Get the terminator operation of this block. This function asserts that
 /// the block has a valid terminator operation.
 Operation *Block::getTerminator() {
   assert(!empty() && !back().isKnownNonTerminator());
   return &back();
 }
 
-/// Return true if this block has no predecessors.
-bool Block::hasNoPredecessors() { return pred_begin() == pred_end(); }
-
 // Indexed successor access.
 unsigned Block::getNumSuccessors() {
   return empty() ? 0 : back().getNumSuccessors();
 }
 
 Block *Block::getSuccessor(unsigned i) {
   assert(i < getNumSuccessors());
   return getTerminator()->getSuccessor(i);

mlir/lib/IR/PatternMatch.cpp

void PatternRewriter::mergeBlocks(Block *source, Block *dest,
 
   // 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();
 }
 
+// Merge the operations of block 'source' before the operation 'op'. Source
+// block should not have existing predecessors or successors.
+void PatternRewriter::mergeBlockBefore(Block *source, Operation *op,
+                                       ValueRange argValues) {
+  assert(source->hasNoPredecessors() &&
+         "expected 'source' to have no predecessors");
+  assert(source->hasNoSuccessors() &&
+         "expected 'source' to have no successors");
+
+  // Split the block containing 'op' into two, one containg all operations
+  // before 'op' (prologue) and another (epilogue) containing 'op' and all
+  // operations after it.
+  Block *prologue = op->getBlock();
+  Block *epilogue = splitBlock(prologue, op->getIterator());
+
+  // Merge the source block at the end of the prologue.
+  mergeBlocks(source, prologue, argValues);
+
+  // Merge the epilogue at the end the prologue.
+  mergeBlocks(epilogue, prologue);
+}
+
 /// 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

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestMergeSingleBlockOps
 
   LogicalResult
   matchAndRewrite(SingleBlockImplicitTerminatorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     SingleBlockImplicitTerminatorOp parentOp =
         op.getParentOfType<SingleBlockImplicitTerminatorOp>();
     if (!parentOp)
       return failure();
-    Block &parentBlock = parentOp.region().front();
     Block &innerBlock = op.region().front();
     TerminatorOp innerTerminator =
         cast<TerminatorOp>(innerBlock.getTerminator());
-    Block *parentPrologue =
-        rewriter.splitBlock(&parentBlock, Block::iterator(op));
+    rewriter.mergeBlockBefore(&innerBlock, op);
     rewriter.eraseOp(innerTerminator);
-    rewriter.mergeBlocks(&innerBlock, &parentBlock, {});
     rewriter.eraseOp(op);
-    rewriter.mergeBlocks(parentPrologue, &parentBlock, {});
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }
 };
 
 struct TestMergeBlocksPatternDriver
     : public PassWrapper<TestMergeBlocksPatternDriver,
                          OperationPass<ModuleOp>> {