[mlir][IR] Clean up mergeBlockBefore and mergeBlocks

[springerm]

Summary

• RewriterBase::mergeBlocks is simplified: it is implemented in terms of mergeBlockBefore.
• The signature of mergeBlockBefore is consistent with other API (such as inlineRegionBefore): an overload for a Block::iterator is added.
• Additional safety checks are added to mergeBlockBefore: detect cases where the resulting IR could be invalid (no more dropAllUses) or partly unreachable (likely a case of incorrect API usage).

Depends On: D144549

[springerm]

Note: mergeBlockBefore could be renamed to inlineBlockBefore to be consistent with the inlineRegionBefore name, but not sure if it's worth the churn.

[springerm]

rebase

[rriddle]

Feel free to rename the API, I don't think the API churn in this case would be that bad.

[springerm]

update

mlir/include/mlir/IR/PatternMatch.h

public:
   }
 
   /// This method erases an operation that is known to have no uses.
   virtual void eraseOp(Operation *op);
 
   /// This method erases all operations in a block.
   virtual void eraseBlock(Block *block);
 
-  /// 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,
+  /// Inline the operations of block 'source' into block 'dest' before the given
+  /// position. The source block will be deleted and must have no uses.
+  /// 'argValues' is used to replace the block arguments of 'source'.
+  ///
+  /// If the source block is inserted at the end of the dest block, the dest
+  /// block must have no successors. Similarly, if the source block is inserted
+  /// somewhere in the middle (or beginning) of the dest block, the source block
+  /// must have no successors. Otherwise, the resulting IR would have
+  /// unreachable operations.
+  virtual void inlineBlockBefore(Block *source, Block *dest,
+                                 Block::iterator before,
                                  ValueRange argValues = std::nullopt);
 
-  // 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 = std::nullopt);
+  /// Inline the operations of block 'source' before the operation 'op'. The
+  /// source block will be deleted and must have no uses. 'argValues' is used to
+  /// replace the block arguments of 'source'
+  ///
+  /// The source block must have no successors. Otherwise, the resulting IR
+  /// would have unreachable operations.
+  void inlineBlockBefore(Block *source, Operation *op,
+                         ValueRange argValues = std::nullopt);
+
+  /// Inline the operations of block 'source' into the end of block 'dest'. The
+  /// source block will be deleted and must have no uses. 'argValues' is used to
+  /// replace the block arguments of 'source'
+  ///
+  /// The dest block must have no successors. Otherwise, the resulting IR would
+  /// have unreachable operation.
+  void mergeBlocks(Block *source, Block *dest,
 
   /// 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`.

mlir/include/mlir/Transforms/DialectConversion.h

public:
   void eraseBlock(Block *block) override;
 
   /// PatternRewriter hook creating a new block.
   void notifyBlockCreated(Block *block) override;
 
   /// PatternRewriter hook for splitting a block into two parts.
   Block *splitBlock(Block *block, Block::iterator before) override;
 
-  /// PatternRewriter hook for merging a block into another.
-  void mergeBlocks(Block *source, Block *dest, ValueRange argValues) override;
+  /// PatternRewriter hook for inlining the ops of a block into another block.
+  void inlineBlockBefore(Block *source, Block *dest, Block::iterator before,
+                         ValueRange argValues = std::nullopt) override;
+  using PatternRewriter::inlineBlockBefore;
 
   /// PatternRewriter hook for moving blocks out of a region.
   void inlineRegionBefore(Region &region, Region &parent,
                           Region::iterator before) override;
   using PatternRewriter::inlineRegionBefore;
 
   /// PatternRewriter hook for cloning blocks of one region into another. The
   /// given region to clone *must* not have been modified as part of conversion

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

for (auto &op : *parallelOp.getBody()) {
     auto reduce = dyn_cast<ReduceOp>(op);
     if (!reduce)
       continue;
 
     Block &reduceBlock = reduce.getReductionOperator().front();
     Value arg = iterArgs[yieldOperands.size()];
     yieldOperands.push_back(reduceBlock.getTerminator()->getOperand(0));
     rewriter.eraseOp(reduceBlock.getTerminator());
-    rewriter.mergeBlockBefore(&reduceBlock, &op, {arg, reduce.getOperand()});
+    rewriter.inlineBlockBefore(&reduceBlock, &op, {arg, reduce.getOperand()});
     rewriter.eraseOp(reduce);
   }
 
   // Then merge the loop body without the terminator.
   rewriter.eraseOp(parallelOp.getBody()->getTerminator());
   Block *newBody = rewriter.getInsertionBlock();
   if (newBody->empty())
     rewriter.mergeBlocks(parallelOp.getBody(), newBody, ivs);
   else
-    rewriter.mergeBlockBefore(parallelOp.getBody(), newBody->getTerminator(),
+    rewriter.inlineBlockBefore(parallelOp.getBody(), newBody->getTerminator(),
                                ivs);
 
   // Finally, create the terminator if required (for loops with no results, it
   // has been already created in loop construction).
   if (!yieldOperands.empty()) {
     rewriter.setInsertionPointToEnd(rewriter.getInsertionBlock());
     rewriter.create<scf::YieldOp>(loc, yieldOperands);
   }
 

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

LogicalResult matchAndRewrite(AffineIfOp op,
     } else if (isTriviallyTrue(affineIfConditions)) {
       blockToMove = op.getThenBlock();
     } else {
       return failure();
     }
     Operation *blockToMoveTerminator = blockToMove->getTerminator();
     // Promote the "blockToMove" block to the parent operation block between the
     // prologue and epilogue of "op".
-    rewriter.mergeBlockBefore(blockToMove, op);
+    rewriter.inlineBlockBefore(blockToMove, op);
     // Replace the "op" operation with the operands of the
     // "blockToMoveTerminator" operation. Note that "blockToMoveTerminator" is
     // the affine.yield operation present in the "blockToMove" block. It has no
     // operands when affine.if is not returning results and therefore, in that
     // case, replaceOp just erases "op". When affine.if is not returning
     // results, the affine.yield operation can be omitted. It gets inserted
     // implicitly.
     rewriter.replaceOp(op, blockToMoveTerminator->getOperands());

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

 simplifyBrToBlockWithSinglePred(BranchOp op, PatternRewriter &rewriter) {
   // Check that the successor block has a single predecessor.
   Block *succ = op.getDest();
   Block *opParent = op->getBlock();
   if (succ == opParent || !llvm::hasSingleElement(succ->getPredecessors()))
     return failure();
 
   // Merge the successor into the current block and erase the branch.
-  rewriter.mergeBlocks(succ, opParent, op.getOperands());
+  SmallVector<Value> brOperands(op.getOperands());
   rewriter.eraseOp(op);
+  rewriter.mergeBlocks(succ, opParent, brOperands);
   return success();
 }
 
 ///   br ^bb1
 /// ^bb1
 ///   br ^bbN(...)
 ///
 ///  -> br ^bbN(...)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (hasPotentialAlloca) {
         return failure();
       if (!lastNonTerminatorInRegion(op))
         return failure();
     }
 
     Block *block = &op.getRegion().front();
     Operation *terminator = block->getTerminator();
     ValueRange results = terminator->getOperands();
-    rewriter.mergeBlockBefore(block, op);
+    rewriter.inlineBlockBefore(block, op);
     rewriter.replaceOp(op, results);
     rewriter.eraseOp(terminator);
     return success();
   }
 };
 
 /// Move allocations into an allocation scope, if it is legal to
 /// move them (e.g. their operands are available at the location

mlir/lib/Dialect/SCF/IR/SCF.cpp

 /// Replaces the given op with the contents of the given single-block region,
 /// using the operands of the block terminator to replace operation results.
 static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
                                 Region &region, ValueRange blockArgs = {}) {
   assert(llvm::hasSingleElement(region) && "expected single-region block");
   Block *block = &region.front();
   Operation *terminator = block->getTerminator();
   ValueRange results = terminator->getOperands();
-  rewriter.mergeBlockBefore(block, op, blockArgs);
+  rewriter.inlineBlockBefore(block, op, blockArgs);
   rewriter.replaceOp(op, results);
   rewriter.eraseOp(terminator);
 }
 
 ///
 /// (ssa-id `=`)? `execute_region` `->` function-result-type `{`
 ///    block+
 /// `}`

 // iterators (inputs) are yielded.
 // 3) The iter arguments have no use and the corresponding (operation) results
 // have no use.
 //
 // These arguments must be defined outside of
 // the ForOp region and can just be forwarded after simplifying the op inits,
 // yields and returns.
 //
-// The implementation uses `mergeBlockBefore` to steal the content of the
+// The implementation uses `inlineBlockBefore` to steal the content of the
 // original ForOp and avoid cloning.
 struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
   using OpRewritePattern<scf::ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(scf::ForOp forOp,
                                 PatternRewriter &rewriter) const final {
     bool canonicalize = false;
     Block &block = forOp.getRegion().front();
     auto yieldOp = cast<scf::YieldOp>(block.getTerminator());
 
     // An internal flat vector of block transfer
     // arguments `newBlockTransferArgs` keeps the 1-1 mapping of original to
     // transformed block argument mappings. This plays the role of a
     // IRMapping for the particular use case of calling into
-    // `mergeBlockBefore`.
+    // `inlineBlockBefore`.
     SmallVector<bool, 4> keepMask;
     keepMask.reserve(yieldOp.getNumOperands());
     SmallVector<Value, 4> newBlockTransferArgs, newIterArgs, newYieldValues,
         newResultValues;
     newBlockTransferArgs.reserve(1 + forOp.getNumIterOperands());
     newBlockTransferArgs.push_back(Value()); // iv placeholder with null value
     newIterArgs.reserve(forOp.getNumIterOperands());
     newYieldValues.reserve(yieldOp.getNumOperands());
LogicalResult matchAndRewrite(scf::ForOp forOp,
     assert(oldBlock.getNumArguments() == newBlockTransferArgs.size() &&
            "unexpected argument size mismatch");
 
     // No results case: the scf::ForOp builder already created a zero
     // result terminator. Merge before this terminator and just get rid of the
     // original terminator that has been merged in.
     if (newIterArgs.empty()) {
       auto newYieldOp = cast<scf::YieldOp>(newBlock.getTerminator());
-      rewriter.mergeBlockBefore(&oldBlock, newYieldOp, newBlockTransferArgs);
+      rewriter.inlineBlockBefore(&oldBlock, newYieldOp, newBlockTransferArgs);
       rewriter.eraseOp(newBlock.getTerminator()->getPrevNode());
       rewriter.replaceOp(forOp, newResultValues);
       return success();
     }
 
     // No terminator case: merge and rewrite the merged terminator.
     auto cloneFilteredTerminator = [&](scf::YieldOp mergedTerminator) {
       OpBuilder::InsertionGuard g(rewriter);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

foreachInSparseConstant(
         SmallVector<Value> args;
         args.append(coords.begin(), coords.end());
         args.push_back(v);
         args.append(reduc);
         // Clones the foreach op to get a copy of the loop body.
         auto cloned = cast<ForeachOp>(rewriter.clone(*op.getOperation()));
         assert(args.size() == cloned.getBody()->getNumArguments());
         Operation *yield = cloned.getBody()->getTerminator();
-        rewriter.mergeBlockBefore(cloned.getBody(), op, args);
+        rewriter.inlineBlockBefore(cloned.getBody(), op, args);
         // clean up
         rewriter.eraseOp(cloned);
         reduc = yield->getOperands();
         rewriter.eraseOp(yield);
       });
 
   rewriter.replaceOp(op, reduc);
   return success();
LogicalResult matchAndRewrite(ForeachOp op,
 
     // Inline body.
     if (!reducValue.empty()) {
       rewriter.mergeBlocks(srcBlock, rewriter.getBlock(), args);
     } else {
       // This is annoying, since scf.for inserts a implicit yield op when
       // there is no reduction variable upon creation, in this case we need to
       // merge the block *before* the yield op.
-      rewriter.mergeBlockBefore(srcBlock, &*rewriter.getInsertionPoint(), args);
+      rewriter.inlineBlockBefore(srcBlock, &*rewriter.getInsertionPoint(),
+                                 args);
     }
 
     for (Dimension d = 0; d < dimRank; d++) {
       // Link the reduction chain. Note that loop emitter update the reducValue
       // in place.
       loopEmitter.exitCurrentLoop(rewriter, loc, reducValue);
       loopEmitter.exitCurrentLoopSeq();
     }

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

static Value insertYieldOp(RewriterBase &rewriter, Location loc, Region &region,
   // Make a clone of overlap region.
   Region tmpRegion;
   IRMapping mapper;
   region.cloneInto(&tmpRegion, tmpRegion.begin(), mapper);
   Block &clonedBlock = tmpRegion.front();
   YieldOp clonedYield = cast<YieldOp>(clonedBlock.getTerminator());
   // Merge cloned block and return yield value.
   Operation *placeholder = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-  rewriter.mergeBlockBefore(&tmpRegion.front(), placeholder, vals);
+  rewriter.inlineBlockBefore(&tmpRegion.front(), placeholder, vals);
   Value val = clonedYield.getResult();
   rewriter.eraseOp(clonedYield);
   rewriter.eraseOp(placeholder);
   return val;
 }
 
 static Value buildUnaryPresent(RewriterBase &rewriter, Location loc,
                                Operation *op, Value v0) {

mlir/lib/IR/PatternMatch.cpp

 }
 
 void RewriterBase::finalizeRootUpdate(Operation *op) {
   // Notify the listener that the operation was modified.
   if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
     rewriteListener->notifyOperationModified(op);
 }
 
-/// 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 RewriterBase::mergeBlocks(Block *source, Block *dest,
-                               ValueRange 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))
-    replaceAllUsesWith(std::get<0>(it), 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();
-}
-
 /// Find uses of `from` and replace them with `to` if the `functor` returns
 /// true. It also marks every modified uses and notifies the rewriter that an
 /// in-place operation modification is about to happen.
 void RewriterBase::replaceUsesWithIf(Value from, Value to,
                                      function_ref<bool(OpOperand &)> functor) {
   for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
     if (functor(operand))
       updateRootInPlace(operand.getOwner(), [&]() { operand.set(to); });
   }
 }
 
-// Merge the operations of block 'source' before the operation 'op'. Source
-// block should not have existing predecessors or successors.
-void RewriterBase::mergeBlockBefore(Block *source, Operation *op,
+void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
+                                     Block::iterator before,
                                      ValueRange argValues) {
+  assert(argValues.size() == source->getNumArguments() &&
+         "incorrect # of argument replacement values");
+
+  // The source block will be deleted, so it should not have any users (i.e.,
+  // there should be no predecessors).
   assert(source->hasNoPredecessors() &&
          "expected 'source' to have no predecessors");
+
+  if (dest->end() != before) {
+    // The source block will be inserted in the middle of the dest block, so
+    // the source block should have no successors. Otherwise, the remainder of
+    // the dest block would be unreachable.
     assert(source->hasNoSuccessors() &&
            "expected 'source' to have no successors");
+  } else {
+    // The source block will be inserted at the end of the dest block, so the
+    // dest block should have no successors. Otherwise, the inserted operations
+    // will be unreachable.
+    assert(dest->hasNoSuccessors() && "expected 'dest' to have no successors");
+  }
 
-  // Split the block containing 'op' into two, one containing 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());
+  // Replace all of the successor arguments with the provided values.
+  for (auto it : llvm::zip(source->getArguments(), argValues))
+    replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
 
-  // Merge the source block at the end of the prologue.
-  mergeBlocks(source, prologue, argValues);
+  // Move operations from the source block to the dest block and erase the
+  // source block.
+  dest->getOperations().splice(before, source->getOperations());
+  source->erase();
+}
+
+void RewriterBase::inlineBlockBefore(Block *source, Operation *op,
+                                     ValueRange argValues) {
+  inlineBlockBefore(source, op->getBlock(), op->getIterator(), argValues);
+}
 
-  // Merge the epilogue at the end the prologue.
-  mergeBlocks(epilogue, prologue);
+void RewriterBase::mergeBlocks(Block *source, Block *dest,
+                               ValueRange argValues) {
+  inlineBlockBefore(source, dest, dest->end(), argValues);
 }
 
 /// Split the operations starting at "before" (inclusive) out of the given
 /// block into a new block, and return it.
 Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {
   return block->splitBlock(before);
 }
 

mlir/lib/Transforms/Utils/DialectConversion.cpp

 //===----------------------------------------------------------------------===//
 // BlockAction
 
 /// The kind of the block action performed during the rewrite.  Actions can be
 /// undone if the conversion fails.
 enum class BlockActionKind {
   Create,
   Erase,
-  Merge,
+  Inline,
   Move,
   Split,
   TypeConversion
 };
 
 /// Original position of the given block in its parent region. During undo
 /// actions, the block needs to be placed after `insertAfterBlock`.
 struct BlockPosition {
   Region *region;
   Block *insertAfterBlock;
 };
 
-/// Information needed to undo the merge actions.
-/// - the source block, and
-/// - the Operation that was the last operation in the dest block before the
-///   merge (could be null if the dest block was empty).
-struct MergeInfo {
+/// Information needed to undo inlining actions.
+/// - the source block
+/// - the first inlined operation (could be null if the source block was empty)
+/// - the last inlined operation (could be null if the source block was empty)
+struct InlineInfo {
   Block *sourceBlock;
-  Operation *destBlockLastInst;
+  Operation *firstInlinedInst;
+  Operation *lastInlinedInst;
 };
 
 /// The storage class for an undoable block action (one of BlockActionKind),
 /// contains the information necessary to undo this action.
 struct BlockAction {
   static BlockAction getCreate(Block *block) {
     return {BlockActionKind::Create, block, {}};
   }
   static BlockAction getErase(Block *block, BlockPosition originalPosition) {
     return {BlockActionKind::Erase, block, {originalPosition}};
   }
-  static BlockAction getMerge(Block *block, Block *sourceBlock) {
-    BlockAction action{BlockActionKind::Merge, block, {}};
-    action.mergeInfo = {sourceBlock, block->empty() ? nullptr : &block->back()};
+  static BlockAction getInline(Block *block, Block *srcBlock,
+                               Block::iterator before) {
+    BlockAction action{BlockActionKind::Inline, block, {}};
+    action.inlineInfo = {srcBlock,
+                         srcBlock->empty() ? nullptr : &srcBlock->front(),
+                         srcBlock->empty() ? nullptr : &srcBlock->back()};
     return action;
   }
   static BlockAction getMove(Block *block, BlockPosition originalPosition) {
     return {BlockActionKind::Move, block, {originalPosition}};
   }
   static BlockAction getSplit(Block *block, Block *originalBlock) {
     BlockAction action{BlockActionKind::Split, block, {}};
     action.originalBlock = originalBlock;
     return action;
   }
   static BlockAction getTypeConversion(Block *block) {
     return BlockAction{BlockActionKind::TypeConversion, block, {}};
   }
 
   // The action kind.
   BlockActionKind kind;
 
   // A pointer to the block that was created by the action.
   Block *block;
 
   union {
-    // In use if kind == BlockActionKind::Move or BlockActionKind::Erase, and
+    // In use if kind == BlockActionKind::Inline or BlockActionKind::Erase, and
     // contains a pointer to the region that originally contained the block as
     // well as the position of the block in that region.
     BlockPosition originalPosition;
     // In use if kind == BlockActionKind::Split and contains a pointer to the
     // block that was split into two parts.
     Block *originalBlock;
-    // In use if kind == BlockActionKind::Merge, and contains the information
-    // needed to undo the merge.
-    MergeInfo mergeInfo;
+    // In use if kind == BlockActionKind::Inline, and contains the information
+    // needed to undo the inlining.
+    InlineInfo inlineInfo;
   };
 };
 
 //===----------------------------------------------------------------------===//
 // UnresolvedMaterialization
 
 /// This class represents an unresolved materialization, i.e. a materialization
 /// that was inserted during conversion that needs to be legalized at the end of
struct ConversionPatternRewriterImpl {
   void notifyBlockIsBeingErased(Block *block);
 
   /// Notifies that a block was created.
   void notifyCreatedBlock(Block *block);
 
   /// Notifies that a block was split.
   void notifySplitBlock(Block *block, Block *continuation);
 
-  /// Notifies that `block` is being merged with `srcBlock`.
-  void notifyBlocksBeingMerged(Block *block, Block *srcBlock);
+  /// Notifies that a block is being inlined into another block.
+  void notifyBlockBeingInlined(Block *block, Block *srcBlock,
+                               Block::iterator before);
 
   /// Notifies that the blocks of a region are about to be moved.
   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
                                         Region::iterator before);
 
   /// Notifies that the blocks of a region were cloned into another.
   void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
                                    Location origRegionLoc);
case BlockActionKind::Erase: {
       auto &blockList = action.originalPosition.region->getBlocks();
       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
       blockList.insert((insertAfterBlock
                             ? std::next(Region::iterator(insertAfterBlock))
                             : blockList.begin()),
                        action.block);
       break;
     }
-    // Split the block at the position which was originally the end of the
-    // destination block (owned by action), and put the instructions back into
-    // the block used before the merge.
-    case BlockActionKind::Merge: {
-      Block *sourceBlock = action.mergeInfo.sourceBlock;
-      Block::iterator splitPoint =
-          (action.mergeInfo.destBlockLastInst
-               ? ++Block::iterator(action.mergeInfo.destBlockLastInst)
-               : action.block->begin());
-      sourceBlock->getOperations().splice(sourceBlock->begin(),
-                                          action.block->getOperations(),
-                                          splitPoint, action.block->end());
+    // Put the instructions from the destination block (owned by the action)
+    // back into the source block.
+    case BlockActionKind::Inline: {
+      Block *sourceBlock = action.inlineInfo.sourceBlock;
+      if (action.inlineInfo.firstInlinedInst) {
+        assert(action.inlineInfo.lastInlinedInst && "expected operation");
+        sourceBlock->getOperations().splice(
+            sourceBlock->begin(), action.block->getOperations(),
+            Block::iterator(action.inlineInfo.firstInlinedInst),
+            ++Block::iterator(action.inlineInfo.lastInlinedInst));
+      }
       break;
     }
     // Move the block back to its original position.
     case BlockActionKind::Move: {
       Region *originalRegion = action.originalPosition.region;
       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
       originalRegion->getBlocks().splice(
           (insertAfterBlock ? std::next(Region::iterator(insertAfterBlock))
void ConversionPatternRewriterImpl::notifyCreatedBlock(Block *block) {
   blockActions.push_back(BlockAction::getCreate(block));
 }
 
 void ConversionPatternRewriterImpl::notifySplitBlock(Block *block,
                                                      Block *continuation) {
   blockActions.push_back(BlockAction::getSplit(continuation, block));
 }
 
-void ConversionPatternRewriterImpl::notifyBlocksBeingMerged(Block *block,
-                                                            Block *srcBlock) {
-  blockActions.push_back(BlockAction::getMerge(block, srcBlock));
+void ConversionPatternRewriterImpl::notifyBlockBeingInlined(
+    Block *block, Block *srcBlock, Block::iterator before) {
+  blockActions.push_back(BlockAction::getInline(block, srcBlock, before));
 }
 
 void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
     Region &region, Region &parent, Region::iterator before) {
   if (region.empty())
     return;
   Block *laterBlock = &region.back();
   for (auto &earlierBlock : llvm::drop_begin(llvm::reverse(region), 1)) {

 
 Block *ConversionPatternRewriter::splitBlock(Block *block,
                                              Block::iterator before) {
   auto *continuation = PatternRewriter::splitBlock(block, before);
   impl->notifySplitBlock(block, continuation);
   return continuation;
 }
 
-void ConversionPatternRewriter::mergeBlocks(Block *source, Block *dest,
+void ConversionPatternRewriter::inlineBlockBefore(Block *source, Block *dest,
+                                                  Block::iterator before,
                                                   ValueRange argValues) {
-  impl->notifyBlocksBeingMerged(dest, source);
-  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");
+#ifndef NDEBUG
+  auto opIgnored = [&](Operation *op) { return impl->isOpIgnored(op); };
+#endif // NDEBUG
+  // The source block will be deleted, so it should not have any users (i.e.,
+  // there should be no predecessors).
+  assert(llvm::all_of(source->getUsers(), opIgnored) &&
+         "expected 'source' to have no predecessors");
+
+  impl->notifyBlockBeingInlined(dest, source, before);
   for (auto it : llvm::zip(source->getArguments(), argValues))
     replaceUsesOfBlockArgument(std::get<0>(it), std::get<1>(it));
-  dest->getOperations().splice(dest->end(), source->getOperations());
+  dest->getOperations().splice(before, source->getOperations());
   eraseBlock(source);
 }
 
 void ConversionPatternRewriter::inlineRegionBefore(Region &region,
                                                    Region &parent,
                                                    Region::iterator before) {
   impl->notifyRegionIsBeingInlinedBefore(region, parent, before);
   PatternRewriter::inlineRegionBefore(region, parent, before);

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

matchAndRewrite(SingleBlockImplicitTerminatorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     SingleBlockImplicitTerminatorOp parentOp =
         op->getParentOfType<SingleBlockImplicitTerminatorOp>();
     if (!parentOp)
       return failure();
     Block &innerBlock = op.getRegion().front();
     TerminatorOp innerTerminator =
         cast<TerminatorOp>(innerBlock.getTerminator());
-    rewriter.mergeBlockBefore(&innerBlock, op);
+    rewriter.inlineBlockBefore(&innerBlock, op);
     rewriter.eraseOp(innerTerminator);
     rewriter.eraseOp(op);
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }
 };
 
 struct TestMergeBlocksPatternDriver