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

Differential Revision: https://reviews.llvm.org/D84795

1 files changed, 116 insertions, 0 deletions

diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 5bc947fc8c91..f6607a5f5524 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -834,6 +834,118 @@ struct TestTypeConversionDriver
 };
 } // 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
 //===----------------------------------------------------------------------===//
@@ -866,5 +978,9 @@ void registerPatternsTestPass() {
   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