[Mlir-commits] [mlir] 9c16eef - [mlir][IR] Add ReverseDominanceIterator for IR walkers

[Matthias Springer]

Author: Matthias Springer
Date: 2023-03-22T09:01:58+01:00
New Revision: 9c16eef1ec46e10185713043663511d49ffff6b1

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

LOG: [mlir][IR] Add ReverseDominanceIterator for IR walkers

Blocks are enumerated depth-first, but post-order. I.e., a block is enumerated when its successors have been enumerated. This iteration style is suitable when deleting blocks in a regions: in the absence of cycles, uses are deleted before their definitions.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/Iterators.h
    mlir/test/IR/visitors.mlir
    mlir/test/lib/IR/TestVisitors.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Iterators.h b/mlir/include/mlir/IR/Iterators.h
index c16f7117f3dc..2c6137c72cf5 100644
--- a/mlir/include/mlir/IR/Iterators.h
+++ b/mlir/include/mlir/IR/Iterators.h
@@ -21,6 +21,7 @@
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/PostOrderIterator.h"
 
 namespace mlir {
 /// This iterator enumerates elements in "reverse" order. It is a wrapper around
@@ -37,7 +38,7 @@ struct ReverseIterator {
 /// This iterator enumerates elements according to their dominance relationship.
 /// Operations and regions are enumerated in "forward" order. Blocks are
 /// enumerated according to their successor relationship. Unreachable blocks are
-/// not enumerated.
+/// not enumerated. Blocks may not be erased during the traversal.
 ///
 /// Note: If `NoGraphRegions` is set to "true", this iterator asserts that each
 /// visited region has SSA dominance. In either case, the ops in such regions
@@ -70,6 +71,44 @@ struct ForwardDominanceIterator {
     return ForwardIterator::makeIterable(range);
   }
 };
+
+/// This iterator enumerates elements according to their reverse dominance
+/// relationship. Operations and regions are enumerated in "reverse" order.
+/// Blocks are enumerated according to their successor relationship, but
+/// post-order. I.e., a block is visited after its successors have been visited.
+/// Cycles in the block graph are broken in an unspecified way. Unreachable
+/// blocks are not enumerated. Blocks may not be erased during the traversal.
+///
+/// Note: If `NoGraphRegions` is set to "true", this iterator asserts that each
+/// visited region has SSA dominance.
+template <bool NoGraphRegions = false>
+struct ReverseDominanceIterator {
+  // llvm::reverse uses RangeT::rbegin and RangeT::rend.
+  static constexpr auto makeIterable(Block &range) {
+    return llvm::reverse(ForwardIterator::makeIterable(range));
+  }
+
+  static constexpr auto makeIterable(Operation &range) {
+    return llvm::reverse(ForwardIterator::makeIterable(range));
+  }
+
+  static auto makeIterable(Region &region) {
+    if (NoGraphRegions) {
+      // Only regions with SSA dominance are allowed.
+      assert(mayHaveSSADominance(region) && "graph regions are not allowed");
+    }
+
+    // Create post-order iterator. Blocks are enumerated according to their
+    // successor relationship.
+    Block *null = nullptr;
+    auto it = region.empty()
+                  ? llvm::make_range(llvm::po_end(null), llvm::po_end(null))
+                  : llvm::post_order(&region.front());
+
+    // Walk API expects Block references instead of pointers.
+    return llvm::make_pointee_range(it);
+  }
+};
 } // namespace mlir
 
 #endif // MLIR_IR_ITERATORS_H

diff  --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index ddbc334fa4ee..2d83d6922e0c 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -323,6 +323,36 @@ func.func @unordered_cfg_with_loop() {
 // CHECK:       Visiting region 0 from operation 'regionOp0'
 // CHECK:       Visiting region 0 from operation 'func.func'
 
+// CHECK-LABEL: Op reverse dominance post-order visits
+// CHECK:       Visiting op 'func.return'
+// CHECK-NOT:   Visiting op 'op6'
+// CHECK:       Visiting op 'op7'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op5'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op1'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op2'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op3'
+// CHECK:       Visiting op 'cf.cond_br'
+// CHECK:       Visiting op 'op0'
+// CHECK:       Visiting op 'regionOp0'
+// CHECK:       Visiting op 'func.func'
+
+// CHECK-LABEL: Block reverse dominance post-order visits
+// CHECK:       Visiting block ^bb7 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb5 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb1 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb2 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb3 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
+
+// CHECK-LABEL: Region reverse dominance post-order visits
+// CHECK:       Visiting region 0 from operation 'regionOp0'
+// CHECK:       Visiting region 0 from operation 'func.func'
+
 // CHECK-LABEL: Block pre-order erasures (skip)
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
 // CHECK:       Cannot erase block ^bb0 from region 0 from operation 'regionOp0', still has uses

diff  --git a/mlir/test/lib/IR/TestVisitors.cpp b/mlir/test/lib/IR/TestVisitors.cpp
index 6ed4abc71b7d..a14347ed2ec3 100644
--- a/mlir/test/lib/IR/TestVisitors.cpp
+++ b/mlir/test/lib/IR/TestVisitors.cpp
@@ -92,6 +92,19 @@ static void testPureCallbacks(Operation *op) {
                  << "\n";
     funcOp->walk<WalkOrder::PostOrder,
                  ForwardDominanceIterator</*NoGraphRegions=*/true>>(regionPure);
+
+    llvm::outs() << "Op reverse dominance post-order visits"
+                 << "\n";
+    funcOp->walk<WalkOrder::PostOrder,
+                 ReverseDominanceIterator</*NoGraphRegions=*/true>>(opPure);
+    llvm::outs() << "Block reverse dominance post-order visits"
+                 << "\n";
+    funcOp->walk<WalkOrder::PostOrder,
+                 ReverseDominanceIterator</*NoGraphRegions=*/true>>(blockPure);
+    llvm::outs() << "Region reverse dominance post-order visits"
+                 << "\n";
+    funcOp->walk<WalkOrder::PostOrder,
+                 ReverseDominanceIterator</*NoGraphRegions=*/true>>(regionPure);
   });
 }