[mlir][IR] Add Iterator template option to IR walkers

[springerm]

Summary

This allows users to specify a top-down or bottom-up traversal of the IR, in addition to the already existing WalkOrder.

Certain transformations work better with a forward traversal. E.g., when cloning a piece of IR, operations should be cloned top-down so that all uses are defined when creating an op.

Certain transformations work better with a reverse traversal. E.g., when erasing a piece of IR, operations should be erased bottom-up to avoid erasing operations that still have users.

[springerm]

rebase

[springerm]

update

[springerm]

mlir/lib/IR/Visitors.cpp

42

Note: Explicit instantiations would not be needed if detail::walk would be defined in the header file.

[springerm]

simplify

[springerm]

update

[Mogball]

This change overall makes sense to me

mlir/lib/IR/Visitors.cpp

31

Can you implement this with if constexpr instead? Walking is very performance sensitive and it's better not to regress existing code.

[springerm]

address comments

[springerm]

mlir/lib/IR/Visitors.cpp

31

I was able to simplify TopDownIterator::makeRange and make it a constexpr function that simplify passes through the argument as a reference. Together with Iterator being a template argument, there should be no difference in performance in case of a top-down traversal.

[rriddle]

mlir/include/mlir/IR/Visitors.h

66–83

Why not just name them "ForwardIterator" and "ReverseIterator" instead of TopDown/BottomUp (which aren't correct, given the hierarchal structure of the IR).

[springerm]

address comments

mlir/include/mlir/IR/Block.h

public:
   SuccessorRange getSuccessors() { return SuccessorRange(this); }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operations in this block. The callback method is called for each
   /// nested region, block or operation, depending on the callback provided.
-  /// Regions, blocks and operations at the same nesting level are visited in
-  /// lexicographical order. The walk order for enclosing regions, blocks and
-  /// operations with respect to their nested ones is specified by 'Order'
+  /// The order in which regions, blocks and operations at the same nesting
+  /// level are visited (e.g., lexicographical or reverse lexicographical order)
+  /// is determined by 'ItOrder'. The walk order for enclosing regions, blocks
+  /// and operations with respect to their nested ones is specified by 'Order'
   /// (post-order by default). A callback on a block or operation is allowed to
   /// erase that block or operation if either:
   ///   * the walk is in post-order, or
   ///   * the walk is in pre-order and the walk is skipped after the erasure.
   /// See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return walk<Order>(begin(), end(), std::forward<FnT>(callback));
+    return walk<Order, ItOrder>(begin(), end(), std::forward<FnT>(callback));
   }
 
   /// Walk the operations in the specified [begin, end) range of this block. The
   /// callback method is called for each nested region, block or operation,
-  /// depending on the callback provided. Regions, blocks and operations at the
-  /// same nesting level are visited in lexicographical order. The walk order
+  /// depending on the callback provided. The order in which regions, blocks and
+  /// operations at the same nesting level are visited (e.g., lexicographical or
+  /// reverse lexicographical order) is determined by 'ItOrder'. The walk order
   /// for enclosing regions, blocks and operations with respect to their nested
   /// ones is specified by 'Order' (post-order by default). This method is
   /// invoked for void-returning callbacks. A callback on a block or operation
   /// is allowed to erase that block or operation only if the walk is in
   /// post-order. See non-void method for pre-order erasure.
   /// See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<std::is_same<RetT, void>::value, RetT>
   walk(Block::iterator begin, Block::iterator end, FnT &&callback) {
     for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
-      detail::walk<Order>(&op, callback);
+      detail::walk<Order, ItOrder>(&op, callback);
   }
 
   /// Walk the operations in the specified [begin, end) range of this block. The
   /// callback method is called for each nested region, block or operation,
-  /// depending on the callback provided. Regions, blocks and operations at the
-  /// same nesting level are visited in lexicographical order. The walk order
+  /// depending on the callback provided. The order in which regions, blocks and
+  /// operations at the same nesting level are visited (e.g., lexicographical or
+  /// reverse lexicographical order) is determined by 'ItOrder'. The walk order
   /// for enclosing regions, blocks and operations with respect to their nested
   /// ones is specified by 'Order' (post-order by default). This method is
   /// invoked for skippable or interruptible callbacks. A callback on a block or
   /// operation is allowed to erase that block or operation if either:
   ///   * the walk is in post-order, or
   ///   * the walk is in pre-order and the walk is skipped after the erasure.
   /// See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<std::is_same<RetT, WalkResult>::value, RetT>
   walk(Block::iterator begin, Block::iterator end, FnT &&callback) {
     for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
-      if (detail::walk<Order>(&op, callback).wasInterrupted())
+      if (detail::walk<Order, ItOrder>(&op, callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/OpDefinition.h

public:
   InFlightDiagnostic emitWarning(const Twine &message = {});
 
   /// Emit a remark about this operation, reporting up to any diagnostic
   /// handlers that may be listening.
   InFlightDiagnostic emitRemark(const Twine &message = {});
 
   /// Walk the operation by calling the callback for each nested operation
   /// (including this one), block or region, depending on the callback provided.
-  /// Regions, blocks and operations at the same nesting level are visited in
-  /// lexicographical order. The walk order for enclosing regions, blocks and
-  /// operations with respect to their nested ones is specified by 'Order'
+  /// The order in which regions, blocks and operations the same nesting level
+  /// are visited (e.g., lexicographical or reverse lexicographical order) is
+  /// determined by 'ItOrder'. The walk order for enclosing regions, blocks
+  /// and operations with respect to their nested ones is specified by 'Order'
   /// (post-order by default). A callback on a block or operation is allowed to
   /// erase that block or operation if either:
   ///   * the walk is in post-order, or
   ///   * the walk is in pre-order and the walk is skipped after the erasure.
   /// See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<llvm::function_traits<std::decay_t<FnT>>::num_args == 1,
                    RetT>
   walk(FnT &&callback) {
-    return state->walk<Order>(std::forward<FnT>(callback));
+    return state->walk<Order, ItOrder>(std::forward<FnT>(callback));
   }
 
   /// Generic walker with a stage aware callback. Walk the operation by calling
   /// the callback for each nested operation (including this one) N+1 times,
   /// where N is the number of regions attached to that operation.
   ///
   /// The callback method can take any of the following forms:
   ///   void(Operation *, const WalkStage &) : Walk all operation opaquely

mlir/include/mlir/IR/Operation.h

public:
   }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operation by calling the callback for each nested operation
   /// (including this one), block or region, depending on the callback provided.
-  /// Regions, blocks and operations at the same nesting level are visited in
-  /// lexicographical order. The walk order for enclosing regions, blocks and
-  /// operations with respect to their nested ones is specified by 'Order'
+  /// The order in which regions, blocks and operations at the same nesting
+  /// level are visited (e.g., lexicographical or reverse lexicographical order)
+  /// is determined by 'ItOrder'. The walk order for enclosing regions, blocks
+  /// and operations with respect to their nested ones is specified by 'Order'
   /// (post-order by default). A callback on a block or operation is allowed to
   /// erase that block or operation if either:
   ///   * the walk is in post-order, or
   ///   * the walk is in pre-order and the walk is skipped after the erasure.
   ///
   /// The callback method can take any of the following forms:
   ///   void(Operation*) : Walk all operations opaquely.
   ///     * op->walk([](Operation *nestedOp) { ...});
   ///   void(OpT) : Walk all operations of the given derived type.
   ///     * op->walk([](ReturnOp returnOp) { ...});
   ///   WalkResult(Operation*|OpT) : Walk operations, but allow for
   ///                                interruption/skipping.
   ///     * op->walk([](... op) {
   ///         // Skip the walk of this op based on some invariant.
   ///         if (some_invariant)
   ///           return WalkResult::skip();
   ///         // Interrupt, i.e cancel, the walk based on some invariant.
   ///         if (another_invariant)
   ///           return WalkResult::interrupt();
   ///         return WalkResult::advance();
   ///       });
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<llvm::function_traits<std::decay_t<FnT>>::num_args == 1,
                    RetT>
   walk(FnT &&callback) {
-    return detail::walk<Order>(this, std::forward<FnT>(callback));
+    return detail::walk<Order, ItOrder>(this, std::forward<FnT>(callback));
   }
 
   /// Generic walker with a stage aware callback. Walk the operation by calling
   /// the callback for each nested operation (including this one) N+1 times,
   /// where N is the number of regions attached to that operation.
   ///
   /// The callback method can take any of the following forms:
   ///   void(Operation *, const WalkStage &) : Walk all operation opaquely

mlir/include/mlir/IR/Region.h

public:
   void dropAllReferences();
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operations in this region. The callback method is called for each
   /// nested region, block or operation, depending on the callback provided.
-  /// Regions, blocks and operations at the same nesting level are visited in
-  /// lexicographical order. The walk order for enclosing regions, blocks and
-  /// operations with respect to their nested ones is specified by 'Order'
+  /// The order in which regions, blocks and operations at the same nesting
+  /// level are visited (e.g., lexicographical or reverse lexicographical order)
+  /// is determined by 'ItOrder'. The walk order for enclosing regions, blocks
+  /// and operations with respect to their nested ones is specified by 'Order'
   /// (post-order by default). This method is invoked for void-returning
   /// callbacks. A callback on a block or operation is allowed to erase that
   /// block or operation only if the walk is in post-order. See non-void method
   /// for pre-order erasure. See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<std::is_same<RetT, void>::value, RetT> walk(FnT &&callback) {
     for (auto &block : *this)
-      block.walk<Order>(callback);
+      block.walk<Order, ItOrder>(callback);
   }
 
   /// Walk the operations in this region. The callback method is called for each
   /// nested region, block or operation, depending on the callback provided.
-  /// Regions, blocks and operations at the same nesting level are visited in
-  /// lexicographical order. The walk order for enclosing regions, blocks and
-  /// operations with respect to their nested ones is specified by 'Order'
+  /// The order in which regions, blocks and operations at the same nesting
+  /// level are visited (e.g., lexicographical or reverse lexicographical order)
+  /// is determined by 'ItOrder'. The walk order for enclosing regions, blocks
+  /// and operations with respect to their nested ones is specified by 'Order'
   /// (post-order by default). This method is invoked for skippable or
   /// interruptible callbacks. A callback on a block or operation is allowed to
   /// erase that block or operation if either:
   ///   * the walk is in post-order,
   ///   * or the walk is in pre-order and the walk is skipped after the erasure.
   /// See Operation::walk for more details.
-  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+  template <WalkOrder Order = WalkOrder::PostOrder,
+            IteratorOrder ItOrder = IteratorOrder::Forward, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
   std::enable_if_t<std::is_same<RetT, WalkResult>::value, RetT>
   walk(FnT &&callback) {
     for (auto &block : *this)
-      if (block.walk<Order>(callback).wasInterrupted())
+      if (block.walk<Order, ItOrder>(callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // CFG view utilities
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Visitors.h

public:
   bool wasInterrupted() const { return result == Interrupt; }
 
   /// Returns true if the walk was skipped.
   bool wasSkipped() const { return result == Skip; }
 };
 
 /// Traversal order for region, block and operation walk utilities.
 enum class WalkOrder { PreOrder, PostOrder };
+enum class IteratorOrder { Forward, Reverse };
 
 /// A utility class to encode the current walk stage for "generic" walkers.
 /// When walking an operation, we can either choose a Pre/Post order walker
 /// which invokes the callback on an operation before/after all its attached
 /// regions have been visited, or choose a "generic" walker where the callback
 /// is invoked on the operation N+1 times where N is the number of regions
 /// attached to that operation. The `WalkStage` class below encodes the current
 /// stage of the walk, i.e., which regions have already been visited, and the
 /// callback accepts an additional argument for the current stage. Such
 /// generic walkers that accept stage-aware callbacks are only applicable when
 /// the callback operates on an operation (i.e., not applicable for callbacks
 /// on Blocks or Regions).
 class WalkStage {
 public:
   explicit WalkStage(Operation *op);
 
   /// Return true if parent operation is being visited before all regions.
   bool isBeforeAllRegions() const { return nextRegion == 0; }
   /// Returns true if parent operation is being visited just before visiting
   /// region number `region`.
   bool isBeforeRegion(int region) const { return nextRegion == region; }
   /// Returns true if parent operation is being visited just after visiting
   /// region number `region`.
   bool isAfterRegion(int region) const { return nextRegion == region + 1; }
   /// Return true if parent operation is being visited after all regions.
   bool isAfterAllRegions() const { return nextRegion == numRegions; }
   /// Advance the walk stage.

 Arg first_argument_type(Ret (F::*)(Arg, Rest...) const);
 template <typename F>
 decltype(first_argument_type(&F::operator())) first_argument_type(F);
 
 /// Type definition of the first argument to the given callable 'T'.
 template <typename T>
 using first_argument = decltype(first_argument_type(std::declval<T>()));
 
+/// This iterator enumerates elements in "forward" order.
+struct ForwardIterator {
+  template <typename RangeT> static auto makeRange(RangeT &&range) {
+    return llvm::make_range(std::begin(std::forward<RangeT>(range)),
+                            std::end(std::forward<RangeT>(range)));
+  }
+};
+
+/// This iterator enumerates elements in "reverse" order. It is a wrapper around
+/// llvm::reverse.
+struct ReverseIterator {
+  template <typename RangeT> static auto makeRange(RangeT &&range) {
+    return llvm::reverse(
+        llvm::make_range(std::begin(std::forward<RangeT>(range)),
+                         std::end(std::forward<RangeT>(range))));
+  }
+};
+
 /// Walk all of the regions, blocks, or operations nested under (and including)
-/// the given operation. Regions, blocks and operations at the same nesting
-/// level are visited in lexicographical order. The walk order for enclosing
-/// regions, blocks and operations with respect to their nested ones is
-/// specified by 'order'. These methods are invoked for void-returning
+/// the given operation. The order in which regions, blocks and operations at
+/// the same nesting level are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'Iterator'. The walk order for
+/// enclosing regions, blocks and operations with respect to their nested ones
+/// is specified by 'order'. These methods are invoked for void-returning
 /// callbacks. A callback on a block or operation is allowed to erase that block
 /// or operation only if the walk is in post-order. See non-void method for
 /// pre-order erasure.
+template <typename Iterator>
 void walk(Operation *op, function_ref<void(Region *)> callback,
           WalkOrder order);
+template <typename Iterator>
 void walk(Operation *op, function_ref<void(Block *)> callback, WalkOrder order);
+template <typename Iterator>
 void walk(Operation *op, function_ref<void(Operation *)> callback,
           WalkOrder order);
 /// Walk all of the regions, blocks, or operations nested under (and including)
-/// the given operation. Regions, blocks and operations at the same nesting
-/// level are visited in lexicographical order. The walk order for enclosing
-/// regions, blocks and operations with respect to their nested ones is
-/// specified by 'order'. This method is invoked for skippable or interruptible
-/// callbacks. A callback on a block or operation is allowed to erase that block
-/// or operation if either:
+/// the given operation. The order in which regions, blocks and operations at
+/// the same nesting level are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'Iterator'. The walk order for
+/// enclosing regions, blocks and operations with respect to their nested ones
+/// is specified by 'order'. This method is invoked for skippable or
+/// interruptible callbacks. A callback on a block or operation is allowed to
+/// erase that block or operation if either:
 ///   * the walk is in post-order, or
 ///   * the walk is in pre-order and the walk is skipped after the erasure.
+template <typename Iterator>
 WalkResult walk(Operation *op, function_ref<WalkResult(Region *)> callback,
                 WalkOrder order);
+template <typename Iterator>
 WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback,
                 WalkOrder order);
+template <typename Iterator>
 WalkResult walk(Operation *op, function_ref<WalkResult(Operation *)> callback,
                 WalkOrder order);
 
 // Below are a set of functions to walk nested operations. Users should favor
 // the direct `walk` methods on the IR classes(Operation/Block/etc) over these
 // methods. They are also templated to allow for statically dispatching based
 // upon the type of the callback function.
 
 /// Walk all of the regions, blocks, or operations nested under (and including)
-/// the given operation. Regions, blocks and operations at the same nesting
-/// level are visited in lexicographical order. The walk order for enclosing
-/// regions, blocks and operations with respect to their nested ones is
-/// specified by 'Order' (post-order by default). A callback on a block or
+/// the given operation. The order in which regions, blocks and operations at
+/// the same nesting level are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'IteratorOrder'. The walk order for
+/// enclosing regions, blocks and operations with respect to their nested ones
+/// is specified by 'Order' (post-order by default). A callback on a block or
 /// operation is allowed to erase that block or operation if either:
 ///   * the walk is in post-order, or
 ///   * the walk is in pre-order and the walk is skipped after the erasure.
 /// This method is selected for callbacks that operate on Region*, Block*, and
 /// Operation*.
 ///
 /// Example:
 ///   op->walk([](Region *r) { ... });
 ///   op->walk([](Block *b) { ... });
 ///   op->walk([](Operation *op) { ... });
 template <
-    WalkOrder Order = WalkOrder::PostOrder, typename FuncTy,
+    WalkOrder Order = WalkOrder::PostOrder,
+    IteratorOrder ItOrder = IteratorOrder::Forward, typename FuncTy,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 std::enable_if_t<llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value,
                  RetT>
 walk(Operation *op, FuncTy &&callback) {
-  return detail::walk(op, function_ref<RetT(ArgT)>(callback), Order);
+  if (ItOrder == IteratorOrder::Forward) {
+    return detail::walk<ForwardIterator>(op, function_ref<RetT(ArgT)>(callback),
+                                         Order);
+  }
+  if (ItOrder == IteratorOrder::Reverse) {
+    return detail::walk<ReverseIterator>(op, function_ref<RetT(ArgT)>(callback),
+                                         Order);
+  }
+  llvm_unreachable("unsupported iterator order");
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
-/// given operation. Regions, blocks and operations at the same nesting
-/// level are visited in lexicographical order. The walk order for enclosing
-/// regions, blocks and operations with respect to their nested ones is
-/// specified by 'order' (post-order by default). This method is selected for
+/// given operation. The order in which regions, blocks and operations at
+/// the same nesting are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'IteratorOrder'. The walk order for
+/// enclosing regions, blocks and operations with respect to their nested ones
+/// is specified by 'order' (post-order by default). This method is selected for
 /// void-returning callbacks that operate on a specific derived operation type.
 /// A callback on an operation is allowed to erase that operation only if the
 /// walk is in post-order. See non-void method for pre-order erasure.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) { ... });
 template <
-    WalkOrder Order = WalkOrder::PostOrder, typename FuncTy,
+    WalkOrder Order = WalkOrder::PostOrder,
+    IteratorOrder ItOrder = IteratorOrder::Forward, typename FuncTy,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 std::enable_if_t<
     !llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value &&
         std::is_same<RetT, void>::value,
     RetT>
 walk(Operation *op, FuncTy &&callback) {
   auto wrapperFn = [&](Operation *op) {
     if (auto derivedOp = dyn_cast<ArgT>(op))
       callback(derivedOp);
   };
-  return detail::walk(op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  if (ItOrder == IteratorOrder::Forward) {
+    return detail::walk<ForwardIterator>(
+        op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  }
+  if (ItOrder == IteratorOrder::Reverse) {
+    return detail::walk<ReverseIterator>(
+        op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  }
+  llvm_unreachable("unsupported iterator order");
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
-/// given operation. Regions, blocks and operations at the same nesting level
-/// are visited in lexicographical order. The walk order for enclosing regions,
-/// blocks and operations with respect to their nested ones is specified by
-/// 'Order' (post-order by default). This method is selected for WalkReturn
-/// returning skippable or interruptible callbacks that operate on a specific
-/// derived operation type. A callback on an operation is allowed to erase that
-/// operation if either:
+/// given operation. The order in which regions, blocks and operations at
+/// the same nesting are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'IteratorOrder'. The walk order for
+/// enclosing regions, blocks and operations with respect to their nested ones
+/// is specified by 'Order' (post-order by default). This method is selected for
+/// WalkReturn returning skippable or interruptible callbacks that operate on a
+/// specific derived operation type. A callback on an operation is allowed to
+/// erase that operation if either:
 ///   * the walk is in post-order, or
 ///   * the walk is in pre-order and the walk is skipped after the erasure.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) {
 ///     if (some_invariant)
 ///       return WalkResult::skip();
 ///     if (another_invariant)
 ///       return WalkResult::interrupt();
 ///     return WalkResult::advance();
 ///   });
 template <
-    WalkOrder Order = WalkOrder::PostOrder, typename FuncTy,
+    WalkOrder Order = WalkOrder::PostOrder,
+    IteratorOrder ItOrder = IteratorOrder::Forward, typename FuncTy,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 std::enable_if_t<
     !llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value &&
         std::is_same<RetT, WalkResult>::value,
     RetT>
 walk(Operation *op, FuncTy &&callback) {
   auto wrapperFn = [&](Operation *op) {
     if (auto derivedOp = dyn_cast<ArgT>(op))
       return callback(derivedOp);
     return WalkResult::advance();
   };
-  return detail::walk(op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  if (ItOrder == IteratorOrder::Forward) {
+    return detail::walk<ForwardIterator>(
+        op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  }
+  if (ItOrder == IteratorOrder::Reverse) {
+    return detail::walk<ReverseIterator>(
+        op, function_ref<RetT(Operation *)>(wrapperFn), Order);
+  }
+  llvm_unreachable("unsupported iterator order");
 }
 
 /// Generic walkers with stage aware callbacks.
 
 /// Walk all the operations nested under (and including) the given operation,
 /// with the callback being invoked on each operation N+1 times, where N is the
 /// number of regions attached to the operation. The `stage` input to the
 /// callback indicates the current walk stage. This method is invoked for void

mlir/lib/IR/Visitors.cpp

 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
 WalkStage::WalkStage(Operation *op)
     : numRegions(op->getNumRegions()), nextRegion(0) {}
 
 /// Walk all of the regions/blocks/operations nested under and including the
-/// given operation. Regions, blocks and operations at the same nesting level
-/// are visited in lexicographical order. The walk order for enclosing regions,
-/// blocks and operations with respect to their nested ones is specified by
-/// 'order'. These methods are invoked for void-returning callbacks. A callback
-/// on a block or operation is allowed to erase that block or operation only if
-/// the walk is in post-order. See non-void method for pre-order erasure.
+/// given operation. The order in which regions, blocks and operations at the
+/// same nesting level are visited (e.g., lexicographical or reverse
+/// lexicographical order) is determined by 'Iterator'. The walk order for
+/// enclosing regions,  blocks and operations with respect to their nested ones
+/// is specified by 'order'. These methods are invoked for void-returning
+/// callbacks. A callback on a block or operation is allowed to erase that block
+/// or operation only if the walk is in post-order. See non-void method for
+/// pre-order erasure.
+template <typename Iterator>
 void detail::walk(Operation *op, function_ref<void(Region *)> callback,
                   WalkOrder order) {
   // We don't use early increment for regions because they can't be erased from
   // a callback.
-  for (auto &region : op->getRegions()) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
     if (order == WalkOrder::PreOrder)
       callback(&region);
-    for (auto &block : region) {
-      for (auto &nestedOp : block)
-        walk(&nestedOp, callback, order);
+    for (auto &block : Iterator::makeRange(region)) {
+      for (auto &nestedOp : Iterator::makeRange(block))
+        walk<Iterator>(&nestedOp, callback, order);
     }
     if (order == WalkOrder::PostOrder)
       callback(&region);
   }
 }
+// Explicit template instantiations for all supported iterators.
+template void
+detail::walk<detail::ForwardIterator>(Operation *, function_ref<void(Region *)>,
+                                      WalkOrder);
+template void
+detail::walk<detail::ReverseIterator>(Operation *, function_ref<void(Region *)>,
+                                      WalkOrder);
 
+template <typename Iterator>
 void detail::walk(Operation *op, function_ref<void(Block *)> callback,
                   WalkOrder order) {
-  for (auto &region : op->getRegions()) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
     // Early increment here in the case where the block is erased.
-    for (auto &block : llvm::make_early_inc_range(region)) {
+    for (auto &block :
+         llvm::make_early_inc_range(Iterator::makeRange(region))) {
       if (order == WalkOrder::PreOrder)
         callback(&block);
-      for (auto &nestedOp : block)
-        walk(&nestedOp, callback, order);
+      for (auto &nestedOp : Iterator::makeRange(block))
+        walk<Iterator>(&nestedOp, callback, order);
       if (order == WalkOrder::PostOrder)
         callback(&block);
     }
   }
 }
+// Explicit template instantiations for all supported iterators.
+template void detail::walk<detail::ForwardIterator>(Operation *,
+                                                    function_ref<void(Block *)>,
+                                                    WalkOrder);
+template void detail::walk<detail::ReverseIterator>(Operation *,
+                                                    function_ref<void(Block *)>,
+                                                    WalkOrder);
 
+template <typename Iterator>
 void detail::walk(Operation *op, function_ref<void(Operation *)> callback,
                   WalkOrder order) {
   if (order == WalkOrder::PreOrder)
     callback(op);
 
   // TODO: This walk should be iterative over the operations.
-  for (auto &region : op->getRegions()) {
-    for (auto &block : region) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
+    for (auto &block : Iterator::makeRange(region)) {
       // Early increment here in the case where the operation is erased.
-      for (auto &nestedOp : llvm::make_early_inc_range(block))
-        walk(&nestedOp, callback, order);
+      for (auto &nestedOp :
+           llvm::make_early_inc_range(Iterator::makeRange(block)))
+        walk<Iterator>(&nestedOp, callback, order);
     }
   }
 
   if (order == WalkOrder::PostOrder)
     callback(op);
 }
+// Explicit template instantiations for all supported iterators.
+template void detail::walk<detail::ForwardIterator>(
+    Operation *, function_ref<void(Operation *)>, WalkOrder);
+template void detail::walk<detail::ReverseIterator>(
+    Operation *, function_ref<void(Operation *)>, WalkOrder);
 
 void detail::walk(Operation *op,
                   function_ref<void(Operation *, const WalkStage &)> callback) {
   WalkStage stage(op);
 
   for (Region &region : op->getRegions()) {
     // Invoke callback on the parent op before visiting each child region.
     callback(op, stage);

 /// is returned by the callback. Walks on regions, blocks and operations may
 /// also be skipped if the callback returns a skip result. Regions, blocks and
 /// operations at the same nesting level are visited in lexicographical order.
 /// The walk order for enclosing regions, blocks and operations with respect to
 /// their nested ones is specified by 'order'. A callback on a block or
 /// operation is allowed to erase that block or operation if either:
 ///   * the walk is in post-order, or
 ///   * the walk is in pre-order and the walk is skipped after the erasure.
+template <typename Iterator>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Region *)> callback,
                         WalkOrder order) {
   // We don't use early increment for regions because they can't be erased from
   // a callback.
-  for (auto &region : op->getRegions()) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
     if (order == WalkOrder::PreOrder) {
       WalkResult result = callback(&region);
       if (result.wasSkipped())
         continue;
       if (result.wasInterrupted())
         return WalkResult::interrupt();
     }
-    for (auto &block : region) {
-      for (auto &nestedOp : block)
-        if (walk(&nestedOp, callback, order).wasInterrupted())
+    for (auto &block : Iterator::makeRange(region)) {
+      for (auto &nestedOp : Iterator::makeRange(block))
+        if (walk<Iterator>(&nestedOp, callback, order).wasInterrupted())
           return WalkResult::interrupt();
     }
     if (order == WalkOrder::PostOrder) {
       if (callback(&region).wasInterrupted())
         return WalkResult::interrupt();
       // We don't check if this region was skipped because its walk already
       // finished and the walk will continue with the next region.
     }
   }
   return WalkResult::advance();
 }
+// Explicit template instantiations for all supported iterators.
+template WalkResult detail::walk<detail::ForwardIterator>(
+    Operation *, function_ref<WalkResult(Region *)>, WalkOrder);
+template WalkResult detail::walk<detail::ReverseIterator>(
+    Operation *, function_ref<WalkResult(Region *)>, WalkOrder);
 
+template <typename Iterator>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Block *)> callback,
                         WalkOrder order) {
-  for (auto &region : op->getRegions()) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
     // Early increment here in the case where the block is erased.
-    for (auto &block : llvm::make_early_inc_range(region)) {
+    for (auto &block :
+         llvm::make_early_inc_range(Iterator::makeRange(region))) {
       if (order == WalkOrder::PreOrder) {
         WalkResult result = callback(&block);
         if (result.wasSkipped())
           continue;
         if (result.wasInterrupted())
           return WalkResult::interrupt();
       }
-      for (auto &nestedOp : block)
-        if (walk(&nestedOp, callback, order).wasInterrupted())
+      for (auto &nestedOp : Iterator::makeRange(block))
+        if (walk<Iterator>(&nestedOp, callback, order).wasInterrupted())
           return WalkResult::interrupt();
       if (order == WalkOrder::PostOrder) {
         if (callback(&block).wasInterrupted())
           return WalkResult::interrupt();
         // We don't check if this block was skipped because its walk already
         // finished and the walk will continue with the next block.
       }
     }
   }
   return WalkResult::advance();
 }
+// Explicit template instantiations for all supported iterators.
+template WalkResult detail::walk<detail::ForwardIterator>(
+    Operation *, function_ref<WalkResult(Block *)>, WalkOrder);
+template WalkResult detail::walk<detail::ReverseIterator>(
+    Operation *, function_ref<WalkResult(Block *)>, WalkOrder);
 
+template <typename Iterator>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Operation *)> callback,
                         WalkOrder order) {
   if (order == WalkOrder::PreOrder) {
     WalkResult result = callback(op);
     // If skipped, caller will continue the walk on the next operation.
     if (result.wasSkipped())
       return WalkResult::advance();
     if (result.wasInterrupted())
       return WalkResult::interrupt();
   }
 
   // TODO: This walk should be iterative over the operations.
-  for (auto &region : op->getRegions()) {
-    for (auto &block : region) {
+  for (auto &region : Iterator::makeRange(op->getRegions())) {
+    for (auto &block : Iterator::makeRange(region)) {
       // Early increment here in the case where the operation is erased.
-      for (auto &nestedOp : llvm::make_early_inc_range(block)) {
-        if (walk(&nestedOp, callback, order).wasInterrupted())
+      for (auto &nestedOp :
+           llvm::make_early_inc_range(Iterator::makeRange(block))) {
+        if (walk<Iterator>(&nestedOp, callback, order).wasInterrupted())
           return WalkResult::interrupt();
       }
     }
   }
 
   if (order == WalkOrder::PostOrder)
     return callback(op);
   return WalkResult::advance();
 }
+// Explicit template instantiations for all supported iterators.
+template WalkResult detail::walk<detail::ForwardIterator>(
+    Operation *, function_ref<WalkResult(Operation *)>, WalkOrder);
+template WalkResult detail::walk<detail::ReverseIterator>(
+    Operation *, function_ref<WalkResult(Operation *)>, WalkOrder);
 
 WalkResult detail::walk(
     Operation *op,
     function_ref<WalkResult(Operation *, const WalkStage &)> callback) {
   WalkStage stage(op);
 
   for (Region &region : op->getRegions()) {
     // Invoke callback on the parent op before visiting each child region.