Add an overload to 'PatternRewriter::inlineRegionBefore' that accepts a parent region for the insertion position. This allows for inlining the given region into the end of another region.

PiperOrigin-RevId: 254367375

mlir/include/mlir/IR/PatternMatch.h

@@ -303,10 +303,12 @@ public:
   }
 
   /// Move the blocks that belong to "region" before the given position in
-  /// another region.  The two regions must be different.  The caller is in
-  /// charge to update create the operation transferring the control flow to the
-  /// region and pass it the correct block arguments.
-  virtual void inlineRegionBefore(Region &region, Region::iterator before);
+  /// another region "parent".  The two regions must be different.  The caller
+  /// is in charge to update create the operation transferring the control flow
+  /// to the region and pass it the correct block arguments.
+  virtual void inlineRegionBefore(Region &region, Region &parent,
+                                  Region::iterator before);
+  void inlineRegionBefore(Region &region, Block *before);
 
   /// This method performs the final replacement for a pattern, where the
   /// results of the operation are updated to use the specified list of SSA

mlir/lib/IR/PatternMatch.cpp

@@ -119,9 +119,12 @@ void PatternRewriter::replaceOpWithResultsOfAnotherOp(
 /// another region.  The two regions must be different.  The caller is in
 /// charge to update create the operation transferring the control flow to the
 /// region and pass it the correct block arguments.
-void PatternRewriter::inlineRegionBefore(Region &region,
+void PatternRewriter::inlineRegionBefore(Region &region, Region &parent,
                                          Region::iterator before) {
-  before->getParent()->getBlocks().splice(before, region.getBlocks());
+  parent.getBlocks().splice(before, region.getBlocks());
+}
+void PatternRewriter::inlineRegionBefore(Region &region, Block *before) {
+  inlineRegionBefore(region, *before->getParent(), before->getIterator());
 }
 
 /// This method is used as the final notification hook for patterns that end

mlir/lib/Transforms/DialectConversion.cpp

@@ -418,7 +418,8 @@ struct DialectConversionRewriter final : public PatternRewriter {
   }
 
   /// PatternRewriter hook for moving blocks out of a region.
-  void inlineRegionBefore(Region &region, Region::iterator before) override {
+  void inlineRegionBefore(Region &region, Region &parent,
+                          Region::iterator before) override {
     for (auto &pair : llvm::enumerate(region)) {
       Block &block = pair.value();
       unsigned position = pair.index();
@@ -428,7 +429,7 @@ struct DialectConversionRewriter final : public PatternRewriter {
       action.originalPosition = {&region, position};
       blockActions.push_back(action);
     }
-    PatternRewriter::inlineRegionBefore(region, before);
+    PatternRewriter::inlineRegionBefore(region, parent, before);
   }
 
   /// PatternRewriter hook for creating a new operation.

mlir/lib/Transforms/LowerAffine.cpp

@@ -394,7 +394,7 @@ public:
     auto *firstBodyBlock =
         rewriter.splitBlock(conditionBlock, conditionBlock->begin());
     auto *lastBodyBlock = &forOp.getRegion().back();
-    rewriter.inlineRegionBefore(forOp.getRegion(), Region::iterator(endBlock));
+    rewriter.inlineRegionBefore(forOp.getRegion(), endBlock);
     auto *iv = conditionBlock->getArgument(0);
 
     // Append the induction variable stepping logic to the last body block and
@@ -518,8 +518,7 @@ public:
     auto *thenBlock = &ifOp.getThenBlocks().front();
     rewriter.setInsertionPointToEnd(&ifOp.getThenBlocks().back());
     rewriter.create<BranchOp>(loc, continueBlock);
-    rewriter.inlineRegionBefore(ifOp.getThenBlocks(),
-                                Region::iterator(continueBlock));
+    rewriter.inlineRegionBefore(ifOp.getThenBlocks(), continueBlock);
 
     // Move blocks from the "else" region (if present) to the region containing
     // 'affine.if', place it before the continuation block and branch to it.  It
@@ -529,8 +528,7 @@ public:
       elseBlock = &ifOp.getElseBlocks().front();
       rewriter.setInsertionPointToEnd(&ifOp.getElseBlocks().back());
       rewriter.create<BranchOp>(loc, continueBlock);
-      rewriter.inlineRegionBefore(ifOp.getElseBlocks(),
-                                  Region::iterator(continueBlock));
+      rewriter.inlineRegionBefore(ifOp.getElseBlocks(), continueBlock);
     }
 
     // Now we just have to handle the condition logic.

mlir/test/TestDialect/TestPatterns.cpp

@@ -60,7 +60,8 @@ struct TestRegionRewriteBlockMovement : public ConversionPattern {
                                      PatternRewriter &rewriter) const final {
     // Inline this region into the parent region.
     auto &parentRegion = *op->getContainingRegion();
-    rewriter.inlineRegionBefore(op->getRegion(0), parentRegion.end());
+    rewriter.inlineRegionBefore(op->getRegion(0), parentRegion,
+                                parentRegion.end());
 
     // Drop this operation.
     rewriter.replaceOp(op, llvm::None);