[Mlir-commits] [mlir] 46f65e4 - [mlir]use correct iterator when eraseOp (#83444)

[llvmlistbot at llvm.org]

Author: Congcong Cai
Date: 2024-03-05T15:33:49+08:00
New Revision: 46f65e45e0f5ce4cc0edabceebee681231d24687

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

LOG: [mlir]use correct iterator when eraseOp (#83444)

#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.

Added: 
    mlir/test/Transforms/gh-77420.mlir

Modified: 
    mlir/lib/IR/PatternMatch.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5ba5328f14b89e..8796289d725707 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -229,7 +229,10 @@ 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())) {
+        // Some blocks may have invalid successor, use a set including nullptr
+        // to avoid null pointer.
+        llvm::SmallPtrSet<Block *, 4> visited{nullptr};
+        for (Block *b : llvm::post_order_ext(&r.front(), visited)) {
           // Visit ops in reverse order.
           for (Operation &op :
                llvm::make_early_inc_range(ReverseIterator::makeIterable(*b)))

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
+  }
+
+}