Fix rollback of first block erasure in a region.

Differential Revision: https://reviews.llvm.org/D91788
llvm · Nov 19, 2020 · b0750e2 · b0750e2
1 parent b4ac05d
commit b0750e2
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -988,7 +988,7 @@ void ConversionPatternRewriterImpl::undoBlockActions(
       Block *insertAfterBlock = action.originalPosition.insertAfterBlock;
       blockList.insert((insertAfterBlock
                             ? std::next(Region::iterator(insertAfterBlock))
-                            : blockList.end()),
+                            : blockList.begin()),
                        action.block);
       break;
     }

diff --git a/mlir/test/Transforms/test-legalizer-full.mlir b/mlir/test/Transforms/test-legalizer-full.mlir
@@ -84,3 +84,18 @@ func @test_undo_region_inline() {
 
   "test.return"() : () -> ()
 }
+
+// -----
+
+// Test that multiple block erases can be properly undone.
+func @test_undo_block_erase() {
+   // expected-error@+1 {{failed to legalize operation 'test.region'}}
+  "test.region"() ({
+    ^bb1(%i0: i64):
+       br ^bb2(%i0 : i64)
+    ^bb2(%i1: i64):
+      "test.invalid"(%i1) : (i64) -> ()
+  }) {legalizer.should_clone, legalizer.erase_old_blocks} : () -> ()
+
+  "test.return"() : () -> ()
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -203,12 +203,16 @@ struct TestRegionRewriteBlockMovement : public ConversionPattern {
                   ConversionPatternRewriter &rewriter) const final {
     // Inline this region into the parent region.
     auto &parentRegion = *op->getParentRegion();
+    auto &opRegion = op->getRegion(0);
     if (op->getAttr("legalizer.should_clone"))
-      rewriter.cloneRegionBefore(op->getRegion(0), parentRegion,
-                                 parentRegion.end());
+      rewriter.cloneRegionBefore(opRegion, parentRegion, parentRegion.end());
     else
-      rewriter.inlineRegionBefore(op->getRegion(0), parentRegion,
-                                  parentRegion.end());
+      rewriter.inlineRegionBefore(opRegion, parentRegion, parentRegion.end());
+
+    if (op->getAttr("legalizer.erase_old_blocks")) {
+      while (!opRegion.empty())
+        rewriter.eraseBlock(&opRegion.front());
+    }
 
     // Drop this operation.
     rewriter.eraseOp(op);