[mlir][DialectConversion] Add support for mergeBlocks in ConversionPatternRewriter.

[rriddle]

Thanks Mahesh!! Looks good, just had some comments on the encoding.

mlir/lib/Transforms/DialectConversion.cpp

627

Assuming all of the other operations are tracked properly, all you really need to do a merge is:

• The source block
• The destination block
• A pointer to the operation that was the previous end of the block destination block. If the destination block was empty, this is just nullptr.

When reverting all you would need to do is splice back (origEnd++, curEnd) to the source block.

644

Please avoid any use of std::distance or random access on these containers, they are intrusive linked lists. These computations all run O(N).

I just noticed that getErase is doing the same thing, it shouldn't.

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

559

Can you add a new driver instead? This one is overly bloated at this point.

[mravishankar]

mlir/lib/Transforms/DialectConversion.cpp

644

Removed the std::distance from here.

Re: use of std::distance in getErased, does that mean I change BlockPosition to be {Region *, Block*} where the block * is the block erased (it is actually same as sourceBlock, so we only need the Region* then?

[mravishankar]

Addressing comments

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

623

Do you actually need this if you are relying on the erase action to do the backtracking?

644

Yeah I would change BlockPosition to use a similar scheme, i.e. keep a pointer to the previous block(or null if the region only has a single block).

1003

destBlockLastInst can be null here.

1335

typo: ezpected

mlir/test/lib/Dialect/Test/TestOps.td

1167

Move the brace to the previous line.

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

849

You can initialize the SmallVector with the range directly, no need to use begin/end.

871

Same here.

[mravishankar]

Addressing comments

[mravishankar]

Addressing missed comments

mlir/lib/Transforms/DialectConversion.cpp

623

No, its actually not used anywhere now. Removed.

644

Added a follow up patch to do this.

1003

Handled it now. Thanks.

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

642

block->empty()?

644

nit: Drop the ()

[mravishankar]

Address minor comments

mlir/lib/Transforms/DialectConversion.cpp

struct OpReplacement {
 
   /// An optional type converter that can be used to materialize conversions
   /// between the new and old values if necessary.
   TypeConverter *converter = nullptr;
 };
 
 /// The kind of the block action performed during the rewrite.  Actions can be
 /// undone if the conversion fails.
-enum class BlockActionKind { Create, Erase, Move, Split, TypeConversion };
+enum class BlockActionKind {
+  Create,
+  Erase,
+  Merge,
+  Move,
+  Split,
+  TypeConversion
+};
 
 /// Original position of the given block in its parent region.  We cannot use
 /// a region iterator because it could have been invalidated by other region
 /// operations since the position was stored.
 struct BlockPosition {
   Region *region;
   Region::iterator::difference_type position;
 };
 
+/// 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 {
+  Block *sourceBlock;
+  Operation *destBlockLastInst;
+};
+
 /// 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 originalPos) {
     return {BlockActionKind::Erase, block, {originalPos}};
   }
+  static BlockAction getMerge(Block *block, Block *sourceBlock) {
+    BlockAction action{BlockActionKind::Merge, block, {}};
+    action.mergeInfo = {sourceBlock, block->empty() ? nullptr : &block->back()};
+    return action;
+  }
   static BlockAction getMove(Block *block, BlockPosition originalPos) {
     return {BlockActionKind::Move, block, {originalPos}};
   }
   static BlockAction getSplit(Block *block, Block *originalBlock) {
     BlockAction action{BlockActionKind::Split, block, {}};
     action.originalBlock = originalBlock;
     return action;
   }
struct BlockAction {
   union {
     // In use if kind == BlockActionKind::Move 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;
   };
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriterImpl
 //===----------------------------------------------------------------------===//
 namespace mlir {
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 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);
 
for (auto &action :
     // Put the block (owned by action) back into its original position.
     case BlockActionKind::Erase: {
       auto &blockList = action.originalPosition.region->getBlocks();
       blockList.insert(
           std::next(blockList.begin(), action.originalPosition.position),
           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());
+      break;
+    }
     // Move the block back to its original position.
     case BlockActionKind::Move: {
       Region *originalRegion = action.originalPosition.region;
       originalRegion->getBlocks().splice(
           std::next(originalRegion->begin(), action.originalPosition.position),
           action.block->getParent()->getBlocks(), action.block);
       break;
     }
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::notifyRegionIsBeingInlinedBefore(
     Region &region, Region &parent, Region::iterator before) {
   for (auto &pair : llvm::enumerate(region)) {
     Block &block = pair.value();
     Region::iterator::difference_type position = pair.index();
     blockActions.push_back(BlockAction::getMove(&block, {&region, position}));
   }
 }
Block *ConversionPatternRewriter::splitBlock(Block *block,
   auto *continuation = PatternRewriter::splitBlock(block, before);
   impl->notifySplitBlock(block, continuation);
   return continuation;
 }
 
 /// PatternRewriter hook for merging a block into another.
 void ConversionPatternRewriter::mergeBlocks(Block *source, Block *dest,
                                             ValueRange argValues) {
-  // TODO: This requires fixing the implementation of
-  // 'replaceUsesOfBlockArgument', which currently isn't undoable.
-  llvm_unreachable("block merging updates are currently not supported");
+  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");
+  for (auto it : llvm::zip(source->getArguments(), argValues))
+    replaceUsesOfBlockArgument(std::get<0>(it), std::get<1>(it));
+  dest->getOperations().splice(dest->end(), source->getOperations());
+  eraseBlock(source);
 }
 
 /// PatternRewriter hook for moving blocks out of a region.
 void ConversionPatternRewriter::inlineRegionBefore(Region &region,
                                                    Region &parent,
                                                    Region::iterator before) {
   impl->notifyRegionIsBeingInlinedBefore(region, parent, before);
   PatternRewriter::inlineRegionBefore(region, parent, before);

mlir/test/Transforms/test-merge-blocks.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -test-merge-blocks -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: @merge_blocks
+func @merge_blocks(%arg0: i32, %arg1 : i32) -> () {
+  //      CHECK: "test.merge_blocks"() ( {
+  // CHECK-NEXT:   "test.return"
+  // CHECK-NEXT: })
+  // CHECK-NEXT: "test.return"
+  %0:2 = "test.merge_blocks"() ({
+  ^bb0:
+     "test.br"(%arg0, %arg1)[^bb1] : (i32, i32) -> ()
+  ^bb1(%arg3 : i32, %arg4 : i32):
+     "test.return"(%arg3, %arg4) : (i32, i32) -> ()
+  }) : () -> (i32, i32)
+  "test.return"(%0#0, %0#1) : (i32, i32) -> ()
+}
+
+// -----
+
+// The op in this function is rewritten to itself (and thus remains
+// illegal) by a pattern that merges the second block with the first
+// after adding an operation into it.  Check that we can undo block
+// removal succesfully.
+// CHECK-LABEL: @undo_blocks_merge
+func @undo_blocks_merge(%arg0: i32) {
+  "test.undo_blocks_merge"() ({
+    // expected-remark@-1 {{op 'test.undo_blocks_merge' is not legalizable}}
+    // CHECK: "unregistered.return"(%{{.*}})[^[[BB:.*]]]
+    "unregistered.return"(%arg0)[^bb1] : (i32) -> ()
+    // expected-remark@-1 {{op 'unregistered.return' is not legalizable}}
+  // CHECK: ^[[BB]]
+  ^bb1(%arg1 : i32):
+    // CHECK: "unregistered.return"
+    "unregistered.return"(%arg1) : (i32) -> ()
+    // expected-remark@-1 {{op 'unregistered.return' is not legalizable}}
+  }) : () -> ()
+}
+
+// -----
+
+// CHECK-LABEL: @inline_regions()
+func @inline_regions() -> ()
+{
+  //      CHECK: test.SingleBlockImplicitTerminator
+  // CHECK-NEXT:   %[[T0:.*]] = "test.type_producer"
+  // CHECK-NEXT:   "test.type_consumer"(%[[T0]])
+  // CHECK-NEXT:   "test.finish"
+  "test.SingleBlockImplicitTerminator"() ({
+  ^bb0:
+    %0 = "test.type_producer"() : () -> i32
+    "test.SingleBlockImplicitTerminator"() ({
+    ^bb1:
+      "test.type_consumer"(%0) : (i32) -> ()
+      "test.finish"() : () -> ()
+    }) : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  "test.return"() : () -> ()
+}

mlir/test/lib/Dialect/Test/TestOps.td

def TestInvalidOp : TEST_Op<"invalid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 def TestTypeProducerOp : TEST_Op<"type_producer">,
   Results<(outs AnyType)>;
 def TestTypeConsumerOp : TEST_Op<"type_consumer">,
   Arguments<(ins AnyType)>;
 def TestValidOp : TEST_Op<"valid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 
+def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
+  let summary = "merge_blocks operation";
+  let description = [{
+    Test op with multiple blocks that are merged with Dialect Conversion"
+  }];
+
+  let regions = (region AnyRegion:$body);
+  let results = (outs Variadic<AnyType>:$result);
+}
+
 //===----------------------------------------------------------------------===//
 // Test parser.
 //===----------------------------------------------------------------------===//
 
 def WrappedKeywordOp : TEST_Op<"wrapped_keyword"> {
   let arguments = (ins StrAttr:$keyword);
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let printer = [{ return ::print(p, *this); }];

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

void runOnOperation() override {
     populateWithGenerated(&getContext(), &patterns);
     patterns.insert<
         TestRegionRewriteBlockMovement, TestRegionRewriteUndo, TestCreateBlock,
         TestCreateIllegalBlock, TestUndoBlockArgReplace, TestUndoBlockErase,
         TestPassthroughInvalidOp, TestSplitReturnType,
         TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
         TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
         TestNonRootReplacement, TestBoundedRecursiveRewrite,
         TestNestedOpCreationUndoRewrite>(&getContext());
     patterns.insert<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
     mlir::populateCallOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());
mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     if (failed(applyPartialConversion(getOperation(), target, patterns)))
       signalPassFailure();
   }
 };
 } // end anonymous namespace
 
+namespace {
+/// A rewriter pattern that tests that blocks can be merged.
+struct TestMergeBlock : public OpConversionPattern<TestMergeBlocksOp> {
+  using OpConversionPattern<TestMergeBlocksOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(TestMergeBlocksOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    Block &firstBlock = op.body().front();
+    Operation *branchOp = firstBlock.getTerminator();
+    Block *secondBlock = &*(std::next(op.body().begin()));
+    auto succOperands = branchOp->getOperands();
+    SmallVector<Value, 2> replacements(succOperands);
+    rewriter.eraseOp(branchOp);
+    rewriter.mergeBlocks(secondBlock, &firstBlock, replacements);
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
+/// A rewrite pattern to tests the undo mechanism of blocks being merged.
+struct TestUndoBlocksMerge : public ConversionPattern {
+  TestUndoBlocksMerge(MLIRContext *ctx)
+      : ConversionPattern("test.undo_blocks_merge", /*benefit=*/1, ctx) {}
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    Block &firstBlock = op->getRegion(0).front();
+    Operation *branchOp = firstBlock.getTerminator();
+    Block *secondBlock = &*(std::next(op->getRegion(0).begin()));
+    rewriter.setInsertionPointToStart(secondBlock);
+    rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
+    auto succOperands = branchOp->getOperands();
+    SmallVector<Value, 2> replacements(succOperands);
+    rewriter.eraseOp(branchOp);
+    rewriter.mergeBlocks(secondBlock, &firstBlock, replacements);
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
+/// A rewrite mechanism to inline the body of the op into its parent, when both
+/// ops can have a single block.
+struct TestMergeSingleBlockOps
+    : public OpConversionPattern<SingleBlockImplicitTerminatorOp> {
+  using OpConversionPattern<
+      SingleBlockImplicitTerminatorOp>::OpConversionPattern;
+
+  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.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>> {
+  void runOnOperation() override {
+    mlir::OwningRewritePatternList patterns;
+    MLIRContext *context = &getContext();
+    patterns
+        .insert<TestMergeBlock, TestUndoBlocksMerge, TestMergeSingleBlockOps>(
+            context);
+    ConversionTarget target(*context);
+    target.addLegalOp<FuncOp, ModuleOp, ModuleTerminatorOp, TerminatorOp,
+                      TestBranchOp, TestTypeConsumerOp, TestTypeProducerOp,
+                      TestReturnOp>();
+    target.addIllegalOp<ILLegalOpF>();
+
+    /// Expect the op to have a single block after legalization.
+    target.addDynamicallyLegalOp<TestMergeBlocksOp>(
+        [&](TestMergeBlocksOp op) -> bool {
+          return llvm::hasSingleElement(op.body());
+        });
+
+    /// Only allow `test.br` within test.merge_blocks op.
+    target.addDynamicallyLegalOp<TestBranchOp>([&](TestBranchOp op) -> bool {
+      return op.getParentOfType<TestMergeBlocksOp>();
+    });
+
+    /// Expect that all nested test.SingleBlockImplicitTerminator ops are
+    /// inlined.
+    target.addDynamicallyLegalOp<SingleBlockImplicitTerminatorOp>(
+        [&](SingleBlockImplicitTerminatorOp op) -> bool {
+          return !op.getParentOfType<SingleBlockImplicitTerminatorOp>();
+        });
+
+    DenseSet<Operation *> unlegalizedOps;
+    (void)applyPartialConversion(getOperation(), target, patterns,
+                                 &unlegalizedOps);
+    for (auto *op : unlegalizedOps)
+      op->emitRemark() << "op '" << op->getName() << "' is not legalizable";
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void registerPatternsTestPass() {
   PassRegistration<TestReturnTypeDriver>("test-return-type",
                                          "Run return type functions");
void registerPatternsTestPass() {
 
   PassRegistration<TestUnknownRootOpDriver>(
       "test-legalize-unknown-root-patterns",
       "Test public remapped value mechanism in ConversionPatternRewriter");
 
   PassRegistration<TestTypeConversionDriver>(
       "test-legalize-type-conversion",
       "Test various type conversion functionalities in DialectConversion");
+
+  PassRegistration<TestMergeBlocksPatternDriver>{
+      "test-merge-blocks",
+      "Test Merging operation in ConversionPatternRewriter"};
 }
 } // namespace mlir