[mlir]use correct iterator when eraseOp

[HerrCai0907]

#66771 introduce llvm::post_order(&r.front()) which is equal to r.front().getSuccessor(...).
It will visit the succ block of current block. But actually here need to visit all block of region in reverse order.
Fixes: #77420.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Congcong Cai (HerrCai0907)

Changes

llvm::post_order(&r.front()) is equal to r.front().getSuccessor(...). It will visit the succ block of current block. But actually here need to visit all block in this region. Fixes: #77420.

---

Full diff: https://github.com/llvm/llvm-project/pull/83444.diff

3 Files Affected:

• (modified) mlir/lib/IR/PatternMatch.cpp (+3-3)
• (added) mlir/test/Transforms/Title: [mlir] --canonicalize causes the Segmentation Fault in eraseUnreachableBlocks. #77420.mlir (+21)
• (modified) mlir/test/Transforms/test-strict-pattern-driver.mlir (+7-7)

diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5ba5328f14b89e..56e3c44acf1913 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -229,14 +229,14 @@ void RewriterBase::eraseOp(Operation *op) {
       // until the region is empty. (The block graph could be disconnected.)
       while (!r.empty()) {
         SmallVector<Block *> erasedBlocks;
-        for (Block *b : llvm::post_order(&r.front())) {
+        for (Block &b : llvm::reverse(r.getBlocks())) {
           // Visit ops in reverse order.
           for (Operation &op :
-               llvm::make_early_inc_range(ReverseIterator::makeIterable(*b)))
+               llvm::make_early_inc_range(ReverseIterator::makeIterable(b)))
             eraseTree(&op);
           // Do not erase the block immediately. This is not supprted by the
           // post_order iterator.
-          erasedBlocks.push_back(b);
+          erasedBlocks.push_back(&b);
         }
         for (Block *b : erasedBlocks) {
           // Explicitly drop all uses in case there is a cycle in the block
diff --git a/mlir/test/Transforms/gh-77420.mlir b/mlir/test/Transforms/gh-77420.mlir
new file mode 100644
index 00000000000000..0037cec0a96ab3
--- /dev/null
+++ b/mlir/test/Transforms/gh-77420.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt --canonicalize %s | FileCheck %s
+
+
+module {
+
+// CHECK:       func.func @f() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+  func.func @f() {
+    return
+  ^bb1:  // no predecessors
+    omp.parallel   {
+      %0 = llvm.intr.stacksave : !llvm.ptr
+      llvm.br ^bb1
+    ^bb1:  // pred: ^bb0
+      omp.terminator
+    }
+    return
+  }
+
+}
diff --git a/mlir/test/Transforms/test-strict-pattern-driver.mlir b/mlir/test/Transforms/test-strict-pattern-driver.mlir
index c87444cba8e1a2..cdade72b69e0aa 100644
--- a/mlir/test/Transforms/test-strict-pattern-driver.mlir
+++ b/mlir/test/Transforms/test-strict-pattern-driver.mlir
@@ -134,13 +134,13 @@ func.func @test_remove_cyclic_blocks() {
 
 // -----
 
-// CHECK-AN: notifyOperationErased: test.dummy_op
 // CHECK-AN: notifyOperationErased: test.bar
-// CHECK-AN: notifyOperationErased: test.qux
 // CHECK-AN: notifyOperationErased: test.qux_unreachable
+// CHECK-AN: notifyOperationErased: test.qux
 // CHECK-AN: notifyOperationErased: test.nested_dummy
 // CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.foo
+// CHECK-AN: notifyOperationErased: test.dummy_op
 // CHECK-AN: notifyOperationErased: test.erase_op
 // CHECK-AN-LABEL: func @test_remove_dead_blocks()
 //  CHECK-AN-NEXT:   return
@@ -172,13 +172,13 @@ func.func @test_remove_dead_blocks() {
 // CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.bar
 // CHECK-AN: notifyOperationErased: cf.br
-// CHECK-AN: notifyOperationErased: test.nested_b
-// CHECK-AN: notifyOperationErased: test.nested_a
-// CHECK-AN: notifyOperationErased: test.nested_d
 // CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.nested_e
+// CHECK-AN: notifyOperationErased: test.nested_d
 // CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.nested_c
+// CHECK-AN: notifyOperationErased: test.nested_b
+// CHECK-AN: notifyOperationErased: test.nested_a
 // CHECK-AN: notifyOperationErased: test.foo
 // CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.dummy_op
@@ -214,9 +214,9 @@ func.func @test_remove_nested_ops() {
 
 // CHECK-AN: notifyOperationErased: test.qux
 // CHECK-AN: notifyOperationErased: cf.br
-// CHECK-AN: notifyOperationErased: test.foo
-// CHECK-AN: notifyOperationErased: cf.br
 // CHECK-AN: notifyOperationErased: test.bar
+// CHECK-AN: notifyOperationErased: cf.br
+// CHECK-AN: notifyOperationErased: test.foo
 // CHECK-AN: notifyOperationErased: cf.cond_br
 // CHECK-AN-LABEL: func @test_remove_diamond(
 //  CHECK-AN-NEXT:   return