[Mlir-commits] [mlir] b0750e2 - Fix rollback of first block erasure in a region.

[Tres Popp]

Author: Tres Popp
Date: 2020-11-19T21:24:10+01:00
New Revision: b0750e2df6e4d5251c67f74265edbc753c76afc2

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

LOG: Fix rollback of first block erasure in a region.

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

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/DialectConversion.cpp
    mlir/test/Transforms/test-legalizer-full.mlir
    mlir/test/lib/Dialect/Test/TestPatterns.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 1720c5dbde62..c4231035c38a 100644
--- 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
index fc6fda2994e1..2cf23611486c 100644
--- 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 at +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
index 40ae9ad52636..437a72176cf6 100644
--- 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);