[mlir][IR] Add ReverseDominanceIterator for IR walkers

[springerm]

Summary

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.

mlir/include/mlir/IR/Iterators.h

 #define MLIR_IR_ITERATORS_H
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/IR/RegionKindInterface.h"
 #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
 /// llvm::reverse.
 struct ReverseIterator {
   // llvm::reverse uses RangeT::rbegin and RangeT::rend.
   template <typename RangeT>
   static constexpr auto makeIterable(RangeT &&range) {
     return llvm::reverse(
         ForwardIterator::makeIterable(std::forward<RangeT>(range)));
   }
 };
 
 /// 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
 /// are visited in forward order, but for regions without SSA dominance this
 /// does not guarantee that defining ops are visited before their users.
 template <bool NoGraphRegions = false>
 struct ForwardDominanceIterator {
   static Block &makeIterable(Block &range) {
static auto makeIterable(Region &region) {
     // Walk API expects Block references instead of pointers.
     return llvm::make_pointee_range(it);
   }
 
   static MutableArrayRef<Region> makeIterable(Operation &range) {
     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

mlir/test/IR/visitors.mlir

 // CHECK:       Visiting block ^bb7 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 'func.func'
 
 // CHECK-LABEL: Region forward dominance post-order visits
 // 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
 // CHECK:       Cannot erase block ^bb1 from region 0 from operation 'regionOp0', still has uses
 // CHECK:       Erasing block ^bb2 from region 0 from operation 'regionOp0'
 // CHECK:       Erasing block ^bb2 from region 0 from operation 'regionOp0'
 // CHECK:       Cannot erase block ^bb2 from region 0 from operation 'regionOp0', still has uses
 // CHECK:       Cannot erase block ^bb3 from region 0 from operation 'regionOp0', still has uses
 // CHECK:       Cannot erase block ^bb4 from region 0 from operation 'regionOp0', still has uses

mlir/test/lib/IR/TestVisitors.cpp

op->walk([&](FunctionOpInterface funcOp) {
     llvm::outs() << "Block forward dominance post-order visits"
                  << "\n";
     funcOp->walk<WalkOrder::PostOrder,
                  ForwardDominanceIterator</*NoGraphRegions=*/true>>(blockPure);
     llvm::outs() << "Region forward dominance post-order visits"
                  << "\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);
   });
 }
 
 /// Tests erasure callbacks that skip the walk.
 static void testSkipErasureCallbacks(Operation *op) {
   auto skipOpErasure = [](Operation *op) {
     // Do not erase module and module children operations. Otherwise, there
     // wouldn't be too much to test in pre-order.