[Mlir-commits] [mlir] 9c7b0c4 - [MLIR] Add PatternRewriter::mergeBlockBefore() to merge a block in the middle of another block.

[Rahul Joshi]

Author: Rahul Joshi
Date: 2020-08-19T16:24:59-07:00
New Revision: 9c7b0c4aa5ebe5ec9fef9ca18bef0c3a916b0fca

URL: https://github.com/llvm/llvm-project/commit/9c7b0c4aa5ebe5ec9fef9ca18bef0c3a916b0fca
DIFF: https://github.com/llvm/llvm-project/commit/9c7b0c4aa5ebe5ec9fef9ca18bef0c3a916b0fca.diff

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

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

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/Block.h
    mlir/include/mlir/IR/PatternMatch.h
    mlir/lib/IR/Block.cpp
    mlir/lib/IR/PatternMatch.cpp
    mlir/test/lib/Dialect/Test/TestPatterns.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 859fa1713ffd..ca2523050b24 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -207,7 +207,10 @@ class Block : public IRObjectWithUseList<BlockOperand>,
   }
 
   /// 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.

diff  --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index f1c7c39a3e73..46dd96408ba7 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -326,6 +326,11 @@ class PatternRewriter : public OpBuilder, public OpBuilder::Listener {
   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);

diff  --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 80131325adff..71f368c49776 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -201,9 +201,6 @@ Operation *Block::getTerminator() {
   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();

diff  --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index e05d234c4ad0..a26bc63ed89d 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -128,6 +128,28 @@ void PatternRewriter::mergeBlocks(Block *source, Block *dest,
   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) {

diff  --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index be5d799a0253..cee87984cd88 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -893,16 +893,12 @@ struct TestMergeSingleBlockOps
         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();
   }