[mlir] Add 'Skip' result to Operation visitor

[dcaballe]

Summary

This patch is a follow-up on D97217. It adds a new 'Skip' result to the Operation visitor
so that a callback can stop the ongoing visit of an operation/block/region and
continue visiting the next one without fully interrupting the walk. Skipping is
needed to be able to erase an operation/block in pre-order and do not continue
visiting the internals of that operation/block.

Related to the skipping mechanism, the patch also introduces the following changes:

• Added new TestIRVisitors pass with basic testing for the IR visitors.
• Fixed missing early increment ranges in visitor implementation.
• Updated documentation of walk methods to include erasure information and walk order information.

[sgrechanik]

mlir/include/mlir/IR/Block.h

259

May be reformat this sentence for better clarity? Something like this:

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 pruned after the erasure.

mlir/lib/IR/Visitors.cpp

139

It's not obvious to me whether it matters that we return prune() here, and not advance(). I think it shouldn't matter but may be I miss something, so maybe a clarifying comment would help.

[rriddle]

Just to clarify, "prune" is ignored for post order walks?

mlir/include/mlir/IR/Visitors.h

31

Why not just Skip? Seems much easier to understand.

mlir/lib/IR/Visitors.cpp

15–20

Remove the newline here.

80

Same here and elsewhere.

mlir/test/lib/IR/TestVisitors.cpp

14

Anonymous namespaces should only contain classes, the functions should be top-level.

[dcaballe]

Addressed feedback.

[dcaballe]

Just to clarify, "prune" is ignored for post order walks?

Yes, it's ignored in the sense that it just happens that because of the post-order traversal, when we come to invoke the callback for, let's say, an operation, all its nested regions/blocks/operations has been already visited so there wouldn't be anything to prune. Does it make sense?

mlir/include/mlir/IR/Block.h

259

Much better, thanks!

mlir/include/mlir/IR/Visitors.h

31

It makes sense, thanks!

mlir/lib/IR/Visitors.cpp

139

Good point! Yeah, I think it makes more sense to return 'advance()' since the pruning for this op is already happening here and the caller shouldn't have to do anything about it. I'll change that and update the related comments. Thanks!

mlir/test/lib/IR/TestVisitors.cpp

14

Good catch!

[rriddle]

LGTM

mlir/include/mlir/IR/Operation.h

502

Missed change?

mlir/lib/IR/Visitors.cpp

139

Pulled this comment above the if.

mlir/test/lib/IR/TestVisitors.cpp

123

Weird phrasing here, can you revise?

125

The naming on these functions and several variables should be camelCase

166

Can you move this to the test namespace?

[dcaballe]

Addressed feedback.

[dcaballe]

Thanks! I will commit it tomorrow if no more comments.

mlir/test/lib/IR/TestVisitors.cpp

125

Sorry about that.

[nicolasvasilache]

mlir/include/mlir/IR/Block.h

255

Did you mean "topological" order ?
If it's really lexicographic, can you describe the alphabet?

mlir/include/mlir/IR/Operation.h

502

Locally, skip the walk of this op's region?

505

Globally interrupt the walk ?

mlir/lib/IR/Visitors.cpp

41

@rriddle in general, is there value of not always making region and block range early increment ?

83

@rriddle Is this a distinction we should want to generally worried about?
Wouldn't it be easier to just always early increment and be done with it?

100

is there an invariant we can / should assert ?

[mehdi_amini]

mlir/lib/IR/Visitors.cpp

41

If you iterate the operation in a block and you're modifying the current operation, what is visited at the next iteration?

We could make it an implicit early_inc but that would make it different from the convention of every other C++ iterators, which runs the risk of unintended behavior for the user.
(now one may argue that the lack of early increment is non intuitive in most cases and is thus the original issue, but that ship has sailed in C++...)

[dcaballe]

mlir/include/mlir/IR/Block.h

255

I'm trying to say "in program order or in order of appearance", which is the order achieved by traversing the lists they are contained in. I've seen "lexicographical order" used for this matter in the past but maybe I'm wrong. I found "topological order" misleading since one might think we are traversing instructions following def-use chains, which would provide one of the many potential topological orders for a set of instructions in a block. Any other suggestion?

mlir/include/mlir/IR/Operation.h

502

Ok, I can add these comments. Thanks!

mlir/lib/IR/Visitors.cpp

83

I didn't do it because I couldn't find a good way to test it. Maybe I'm missing something but I couldn't find public APIs to erase/insert regions from/to ops (probably for good reasons?). Therefore, I leaned towards preserving the existing behavior. No strong opinion, though. I buy the consistency argument. Just let me know and I'll change it.

100

Not that I can think of... What this comment is saying is that if the region was skipped, there's nothing we can skip at this point since the region was already visited (post-order). We could do:

if (result.wasSkipped())
  continue;

But this would be redundant since we are at the end of the loop.

mlir/include/mlir/IR/Block.h

public:
   succ_iterator succ_begin() { return getSuccessors().begin(); }
   succ_iterator succ_end() { return getSuccessors().end(); }
   SuccessorRange getSuccessors() { return SuccessorRange(this); }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
-  /// Walk the operations in this block, calling the callback for each
-  /// operation. The walk order for regions, blocks and operations is specified
-  /// by 'Order' (post-order by default).
+  /// 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'
+  /// (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,
             typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
     return walk<Order>(begin(), end(), std::forward<FnT>(callback));
   }
 
-  /// Walk the operations in the specified [begin, end) range of this block,
-  /// calling the callback for each operation. The walk order for regions,
-  /// blocks and operations is specified by 'Order' (post-order by default).
-  /// This method is invoked for void return callbacks.
+  /// 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
+  /// 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,
             typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, void>::value, RetT>::type
   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);
   }
 
-  /// Walk the operations in the specified [begin, end) range of this block,
-  /// calling the callback for each operation. The walk order for regions,
-  /// blocks and operations is specified by 'Order' (post-order by default).
-  /// This method is invoked for interruptible callbacks.
+  /// 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
+  /// 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,
             typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, WalkResult>::value, RetT>::type
   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())
         return WalkResult::interrupt();

mlir/include/mlir/IR/OpDefinition.h

public:
   /// Emit a warning about this operation, reporting up to any diagnostic
   /// handlers that may be listening.
   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). The walk order for regions, blocks and
-  /// operations is specified by 'Order' (post-order by default).
+  /// 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'
+  /// (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,
             typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
     return state->walk<Order>(std::forward<FnT>(callback));
   }
 
   // These are default implementations of customization hooks.

mlir/include/mlir/IR/Operation.h

template <template <typename T> class Trait> bool mightHaveTrait() {
     return abstractOp ? abstractOp->hasTrait<Trait>() : true;
   }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operation by calling the callback for each nested operation
-  /// (including this one). The walk order for regions, blocks and operations is
-  /// specified by 'Order' (post-order by default). The callback method can take
-  /// any of the following forms:
+  /// (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'
+  /// (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/cancellation.
+  ///                                interruption/skipping.
   ///     * op->walk([](... op) {
-  ///         // Interrupt, i.e cancel, the walk based on some invariant.
+  ///         // 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,
             typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
     return detail::walk<Order>(this, std::forward<FnT>(callback));
   }

mlir/include/mlir/IR/Region.h

public:
   /// an essential step in breaking cyclic dependences between references when
   /// they are to be deleted.
   void dropAllReferences();
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
-  /// Walk the operations in this region in postorder, calling the callback for
-  /// each operation. The walk order for regions, blocks and operations is
-  /// specified by 'Order' (post-order by default). This method is invoked for
-  /// void-returning callbacks.
-  /// See Operation::walk for more details.
+  /// 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'
+  /// (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,
             typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, void>::value, RetT>::type
   walk(FnT &&callback) {
     for (auto &block : *this)
       block.walk<Order>(callback);
   }
 
-  /// Walk the operations in this region in postorder, calling the callback for
-  /// each operation. The walk order for regions, blocks and operations is
-  /// specified by 'Order' (post-order by default). This method is invoked for
-  /// interruptible callbacks.
+  /// 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'
+  /// (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,
             typename RetT = detail::walkResultType<FnT>>
   typename std::enable_if<std::is_same<RetT, WalkResult>::value, RetT>::type
   walk(FnT &&callback) {
     for (auto &block : *this)
       if (block.walk<Order>(callback).wasInterrupted())
         return WalkResult::interrupt();

mlir/include/mlir/IR/Visitors.h

 
 namespace mlir {
 class Diagnostic;
 class InFlightDiagnostic;
 class Operation;
 class Block;
 class Region;
 
-/// A utility result that is used to signal if a walk method should be
-/// interrupted or advance.
+/// A utility result that is used to signal how to proceed with an ongoing walk:
+///   * Interrupt: the walk will be interrupted and no more operations, regions
+///   or blocks will be visited.
+///   * Advance: the walk will continue.
+///   * Skip: the walk of the current operation, region or block and their
+///   nested elements that haven't been visited already will be skipped and will
+///   continue with the next operation, region or block.
 class WalkResult {
-  enum ResultEnum { Interrupt, Advance } result;
+  enum ResultEnum { Interrupt, Advance, Skip } result;
 
 public:
   WalkResult(ResultEnum result) : result(result) {}
 
   /// Allow LogicalResult to interrupt the walk on failure.
   WalkResult(LogicalResult result)
       : result(failed(result) ? Interrupt : Advance) {}
 
   /// Allow diagnostics to interrupt the walk.
   WalkResult(Diagnostic &&) : result(Interrupt) {}
   WalkResult(InFlightDiagnostic &&) : result(Interrupt) {}
 
   bool operator==(const WalkResult &rhs) const { return result == rhs.result; }
 
   static WalkResult interrupt() { return {Interrupt}; }
   static WalkResult advance() { return {Advance}; }
+  static WalkResult skip() { return {Skip}; }
 
   /// Returns true if the walk was interrupted.
   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 };
 
 namespace detail {
 /// Helper templates to deduce the first argument of a callback parameter.
 template <typename Ret, typename Arg> Arg first_argument_type(Ret (*)(Arg));
 template <typename Ret, typename F, typename Arg>
 Arg first_argument_type(Ret (F::*)(Arg));
 template <typename Ret, typename F, typename Arg>
 Arg first_argument_type(Ret (F::*)(Arg) 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>()));
 
 /// Walk all of the regions, blocks, or operations nested under (and including)
-/// the given operation. The walk order is specified by 'order'.
+/// 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.
 void walk(Operation *op, function_ref<void(Region *)> callback,
           WalkOrder order);
 void walk(Operation *op, function_ref<void(Block *)> callback, WalkOrder order);
 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. The walk order is specified by 'order'. These functions
-/// walk until an interrupt result is returned by the callback.
+/// 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 walk is in post-order, or
+///   * the walk is in pre-order and the walk is skipped after the erasure.
 WalkResult walk(Operation *op, function_ref<WalkResult(Region *)> callback,
                 WalkOrder order);
 WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback,
                 WalkOrder order);
 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. The walk order is specified by 'Order' (post-order
-/// by default). This method is selected for callbacks that operate on
-/// Region*, Block*, and Operation*.
+/// 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
+/// 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,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 typename std::enable_if<
     llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value, RetT>::type
 walk(Operation *op, FuncTy &&callback) {
   return detail::walk(op, function_ref<RetT(ArgT)>(callback), Order);
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
-/// given operation. The walk order for regions, blocks and operations is
-/// specified by 'Order' (post-order by default). This method is selected for
-/// void returning callbacks that operate on a specific derived operation type.
+/// 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
+/// 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,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 typename std::enable_if<
     !llvm::is_one_of<ArgT, Operation *, Region *, Block *>::value &&
         std::is_same<RetT, void>::value,
     RetT>::type
 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);
 }
 
 /// Walk all of the operations of type 'ArgT' nested under and including the
-/// given operation. The walk order for regions, blocks and operations is
-/// specified by 'Order' (post-order by default). This method is selected for
-/// WalkReturn returning interruptible callbacks that operate on a specific
-/// derived operation type.
+/// 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:
+///   * 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,
     typename ArgT = detail::first_argument<FuncTy>,
     typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
 typename std::enable_if<

mlir/lib/IR/Visitors.cpp

 //===- Visitors.cpp - MLIR Visitor Utilities ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Visitors.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
 /// Walk all of the regions/blocks/operations nested under and including the
-/// given operation. The walk order is specified by 'Order'.
-
+/// 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.
 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()) {
     if (order == WalkOrder::PreOrder)
       callback(&region);
     for (auto &block : region) {
       for (auto &nestedOp : block)
         walk(&nestedOp, callback, order);
     }
     if (order == WalkOrder::PostOrder)
       callback(&region);
   }
 }
 
 void detail::walk(Operation *op, function_ref<void(Block *)> callback,
                   WalkOrder order) {
   for (auto &region : op->getRegions()) {
-    for (auto &block : region) {
+    // Early increment here in the case where the block is erased.
+    for (auto &block : llvm::make_early_inc_range(region)) {
       if (order == WalkOrder::PreOrder)
         callback(&block);
       for (auto &nestedOp : block)
         walk(&nestedOp, callback, order);
       if (order == WalkOrder::PostOrder)
         callback(&block);
     }
   }
for (auto &region : op->getRegions()) {
     }
   }
 
   if (order == WalkOrder::PostOrder)
     callback(op);
 }
 
 /// Walk all of the regions/blocks/operations nested under and including the
-/// given operation. The walk order is specified by 'order'. These functions
-/// walk operations until an interrupt result is returned by the callback.
+/// given operation. These functions walk operations until an interrupt result
+/// 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.
 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()) {
-    if (order == WalkOrder::PreOrder)
-      if (callback(&region).wasInterrupted())
+    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)
         walk(&nestedOp, callback, order);
     }
-    if (order == WalkOrder::PostOrder)
+    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();
 }
 
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Block *)> callback,
                         WalkOrder order) {
   for (auto &region : op->getRegions()) {
-    for (auto &block : region) {
-      if (order == WalkOrder::PreOrder)
-        if (callback(&block).wasInterrupted())
+    // Early increment here in the case where the block is erased.
+    for (auto &block : llvm::make_early_inc_range(region)) {
+      if (order == WalkOrder::PreOrder) {
+        WalkResult result = callback(&block);
+        if (result.wasSkipped())
+          continue;
+        if (result.wasInterrupted())
           return WalkResult::interrupt();
+      }
       for (auto &nestedOp : block)
         walk(&nestedOp, callback, order);
-      if (order == WalkOrder::PostOrder)
+      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();
 }
 
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Operation *)> callback,
                         WalkOrder order) {
-  if (order == WalkOrder::PreOrder)
-    if (callback(op).wasInterrupted())
+  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) {
       // 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())
           return WalkResult::interrupt();
       }
     }
   }
 
   if (order == WalkOrder::PostOrder)
     return callback(op);
   return WalkResult::advance();
 }

mlir/test/IR/visitors.mlir

+// RUN: mlir-opt -test-ir-visitors -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+
+// Verify the different configurations of IR visitors.
+// Constant, yield and other terminator ops are not matched for simplicity.
+// Module and function op and their immediately nested blocks are not erased in
+// callbacks with return so that the output includes more cases in pre-order.
+
+func @structured_cfg() {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c10 = constant 10 : index
+  scf.for %i = %c1 to %c10 step %c1 {
+    %cond = "use0"(%i) : (index) -> (i1)
+    scf.if %cond {
+      "use1"(%i) : (index) -> ()
+    } else {
+      "use2"(%i) : (index) -> ()
+    }
+    "use3"(%i) : (index) -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: Op pre-order visit
+// CHECK:       Visiting op 'module'
+// CHECK:       Visiting op 'func'
+// CHECK:       Visiting op 'scf.for'
+// CHECK:       Visiting op 'use0'
+// CHECK:       Visiting op 'scf.if'
+// CHECK:       Visiting op 'use1'
+// CHECK:       Visiting op 'use2'
+// CHECK:       Visiting op 'use3'
+// CHECK:       Visiting op 'std.return'
+
+// CHECK-LABEL: Block pre-order visits
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'module'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'func'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
+// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'
+
+// CHECK-LABEL: Region pre-order visits
+// CHECK:       Visiting region 0 from operation 'module'
+// CHECK:       Visiting region 0 from operation 'func'
+// CHECK:       Visiting region 0 from operation 'scf.for'
+// CHECK:       Visiting region 0 from operation 'scf.if'
+// CHECK:       Visiting region 1 from operation 'scf.if'
+
+// CHECK-LABEL: Op post-order visits
+// CHECK:       Visiting op 'use0'
+// CHECK:       Visiting op 'use1'
+// CHECK:       Visiting op 'use2'
+// CHECK:       Visiting op 'scf.if'
+// CHECK:       Visiting op 'use3'
+// CHECK:       Visiting op 'scf.for'
+// CHECK:       Visiting op 'std.return'
+// CHECK:       Visiting op 'func'
+// CHECK:       Visiting op 'module'
+
+// CHECK-LABEL: Block post-order visits
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
+// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'func'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'module'
+
+// CHECK-LABEL: Region post-order visits
+// CHECK:       Visiting region 0 from operation 'scf.if'
+// CHECK:       Visiting region 1 from operation 'scf.if'
+// CHECK:       Visiting region 0 from operation 'scf.for'
+// CHECK:       Visiting region 0 from operation 'func'
+// CHECK:       Visiting region 0 from operation 'module'
+
+// CHECK-LABEL: Op pre-order erasures
+// CHECK:       Erasing op 'scf.for'
+// CHECK:       Erasing op 'std.return'
+
+// CHECK-LABEL: Block pre-order erasures
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'
+
+// CHECK-LABEL: Op post-order erasures (skip)
+// CHECK:       Erasing op 'use0'
+// CHECK:       Erasing op 'use1'
+// CHECK:       Erasing op 'use2'
+// CHECK:       Erasing op 'scf.if'
+// CHECK:       Erasing op 'use3'
+// CHECK:       Erasing op 'scf.for'
+// CHECK:       Erasing op 'std.return'
+
+// CHECK-LABEL: Block post-order erasures (skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.if'
+// CHECK:       Erasing block ^bb0 from region 1 from operation 'scf.if'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'
+
+// CHECK-LABEL: Op post-order erasures (no skip)
+// CHECK:       Erasing op 'use0'
+// CHECK:       Erasing op 'use1'
+// CHECK:       Erasing op 'use2'
+// CHECK:       Erasing op 'scf.if'
+// CHECK:       Erasing op 'use3'
+// CHECK:       Erasing op 'scf.for'
+// CHECK:       Erasing op 'std.return'
+// CHECK:       Erasing op 'func'
+// CHECK:       Erasing op 'module'
+
+// CHECK-LABEL: Block post-order erasures (no skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.if'
+// CHECK:       Erasing block ^bb0 from region 1 from operation 'scf.if'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'func'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'module'
+
+// -----
+
+func @unstructured_cfg() {
+  "regionOp0"() ({
+    ^bb0:
+      "op0"() : () -> ()
+      br ^bb2
+    ^bb1:
+      "op1"() : () -> ()
+      br ^bb2
+    ^bb2:
+      "op2"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK-LABEL: Op pre-order visits
+// CHECK:       Visiting op 'module'
+// CHECK:       Visiting op 'func'
+// CHECK:       Visiting op 'regionOp0'
+// CHECK:       Visiting op 'op0'
+// CHECK:       Visiting op 'std.br'
+// CHECK:       Visiting op 'op1'
+// CHECK:       Visiting op 'std.br'
+// CHECK:       Visiting op 'op2'
+// CHECK:       Visiting op 'std.return'
+
+// CHECK-LABEL: Block pre-order visits
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'module'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'func'
+// CHECK:       Visiting block ^bb0 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-LABEL: Region pre-order visits
+// CHECK:       Visiting region 0 from operation 'module'
+// CHECK:       Visiting region 0 from operation 'func'
+// CHECK:       Visiting region 0 from operation 'regionOp0'
+
+// CHECK-LABEL: Op post-order visits
+// CHECK:       Visiting op 'op0'
+// CHECK:       Visiting op 'std.br'
+// CHECK:       Visiting op 'op1'
+// CHECK:       Visiting op 'std.br'
+// CHECK:       Visiting op 'op2'
+// CHECK:       Visiting op 'regionOp0'
+// CHECK:       Visiting op 'std.return'
+// CHECK:       Visiting op 'func'
+// CHECK:       Visiting op 'module'
+
+// CHECK-LABEL: Block post-order visits
+// CHECK:       Visiting block ^bb0 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 ^bb0 from region 0 from operation 'func'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'module'
+
+// CHECK-LABEL: Region post-order visits
+// CHECK:       Visiting region 0 from operation 'regionOp0'
+// CHECK:       Visiting region 0 from operation 'func'
+// CHECK:       Visiting region 0 from operation 'module'
+
+// CHECK-LABEL: Op pre-order erasures (skip)
+// CHECK:       Erasing op 'regionOp0'
+// CHECK:       Erasing op 'std.return'
+
+// CHECK-LABEL: Block pre-order erasures (skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+
+// CHECK-LABEL: Op post-order erasures (skip)
+// CHECK:       Erasing op 'op0'
+// CHECK:       Erasing op 'std.br'
+// CHECK:       Erasing op 'op1'
+// CHECK:       Erasing op 'std.br'
+// CHECK:       Erasing op 'op2'
+// CHECK:       Erasing op 'regionOp0'
+// CHECK:       Erasing op 'std.return'
+
+// CHECK-LABEL: Block post-order erasures (skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+
+// CHECK-LABEL: Op post-order erasures (no skip)
+// CHECK:       Erasing op 'op0'
+// CHECK:       Erasing op 'std.br'
+// CHECK:       Erasing op 'op1'
+// CHECK:       Erasing op 'std.br'
+// CHECK:       Erasing op 'op2'
+// CHECK:       Erasing op 'regionOp0'
+// CHECK:       Erasing op 'std.return'
+
+// CHECK-LABEL: Block post-order erasures (no skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'func'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'module'

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestPrintDefUse.cpp
   TestPrintNesting.cpp
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp
   TestTypes.cpp
+  TestVisitors.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/IR/TestVisitors.cpp

+//===- TestIRVisitors.cpp - Pass to test the IR visitors ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+static void printRegion(Region *region) {
+  llvm::outs() << "region " << region->getRegionNumber() << " from operation '"
+               << region->getParentOp()->getName() << "'";
+}
+
+static void printBlock(Block *block) {
+  llvm::outs() << "block ";
+  block->printAsOperand(llvm::outs(), /*printType=*/false);
+  llvm::outs() << " from ";
+  printRegion(block->getParent());
+}
+
+static void printOperation(Operation *op) {
+  llvm::outs() << "op '" << op->getName() << "'";
+}
+
+/// Tests pure callbacks.
+static void testPureCallbacks(Operation *op) {
+  auto opPure = [](Operation *op) {
+    llvm::outs() << "Visiting ";
+    printOperation(op);
+    llvm::outs() << "\n";
+  };
+  auto blockPure = [](Block *block) {
+    llvm::outs() << "Visiting ";
+    printBlock(block);
+    llvm::outs() << "\n";
+  };
+  auto regionPure = [](Region *region) {
+    llvm::outs() << "Visiting ";
+    printRegion(region);
+    llvm::outs() << "\n";
+  };
+
+  llvm::outs() << "Op pre-order visits"
+               << "\n";
+  op->walk<WalkOrder::PreOrder>(opPure);
+  llvm::outs() << "Block pre-order visits"
+               << "\n";
+  op->walk<WalkOrder::PreOrder>(blockPure);
+  llvm::outs() << "Region pre-order visits"
+               << "\n";
+  op->walk<WalkOrder::PreOrder>(regionPure);
+
+  llvm::outs() << "Op post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder>(opPure);
+  llvm::outs() << "Block post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder>(blockPure);
+  llvm::outs() << "Region post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder>(regionPure);
+}
+
+/// Tests erasure callbacks that skip the walk.
+static void testSkipErasureCallbacks(Operation *op) {
+  auto skipOpErasure = [](Operation *op) {
+    // Do not erase module and function op. Otherwise there wouldn't be too
+    // much to test in pre-order.
+    if (isa<ModuleOp>(op) || isa<FuncOp>(op))
+      return WalkResult::advance();
+
+    llvm::outs() << "Erasing ";
+    printOperation(op);
+    llvm::outs() << "\n";
+    op->dropAllUses();
+    op->erase();
+    return WalkResult::skip();
+  };
+  auto skipBlockErasure = [](Block *block) {
+    // Do not erase module and function blocks. Otherwise there wouldn't be
+    // too much to test in pre-order.
+    Operation *parentOp = block->getParentOp();
+    if (isa<ModuleOp>(parentOp) || isa<FuncOp>(parentOp))
+      return WalkResult::advance();
+
+    llvm::outs() << "Erasing ";
+    printBlock(block);
+    llvm::outs() << "\n";
+    block->erase();
+    return WalkResult::skip();
+  };
+
+  llvm::outs() << "Op pre-order erasures (skip)"
+               << "\n";
+  Operation *cloned = op->clone();
+  cloned->walk<WalkOrder::PreOrder>(skipOpErasure);
+  cloned->erase();
+
+  llvm::outs() << "Block pre-order erasures (skip)"
+               << "\n";
+  cloned = op->clone();
+  cloned->walk<WalkOrder::PreOrder>(skipBlockErasure);
+  cloned->erase();
+
+  llvm::outs() << "Op post-order erasures (skip)"
+               << "\n";
+  cloned = op->clone();
+  cloned->walk<WalkOrder::PostOrder>(skipOpErasure);
+  cloned->erase();
+
+  llvm::outs() << "Block post-order erasures (skip)"
+               << "\n";
+  cloned = op->clone();
+  cloned->walk<WalkOrder::PostOrder>(skipBlockErasure);
+  cloned->erase();
+}
+
+/// Tests callbacks that erase the op or block but don't return 'Skip'. This
+/// callbacks are only valid in post-order.
+static void testNoSkipErasureCallbacks(Operation *op) {
+  auto noSkipOpErasure = [](Operation *op) {
+    llvm::outs() << "Erasing ";
+    printOperation(op);
+    llvm::outs() << "\n";
+    op->dropAllUses();
+    op->erase();
+  };
+  auto noSkipBlockErasure = [](Block *block) {
+    llvm::outs() << "Erasing ";
+    printBlock(block);
+    llvm::outs() << "\n";
+    block->erase();
+  };
+
+  llvm::outs() << "Op post-order erasures (no skip)"
+               << "\n";
+  Operation *cloned = op->clone();
+  cloned->walk<WalkOrder::PostOrder>(noSkipOpErasure);
+
+  llvm::outs() << "Block post-order erasures (no skip)"
+               << "\n";
+  cloned = op->clone();
+  cloned->walk<WalkOrder::PostOrder>(noSkipBlockErasure);
+  cloned->erase();
+}
+
+namespace {
+/// This pass exercises the different configurations of the IR visitors.
+struct TestIRVisitorsPass
+    : public PassWrapper<TestIRVisitorsPass, OperationPass<>> {
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    testPureCallbacks(op);
+    testSkipErasureCallbacks(op);
+    testNoSkipErasureCallbacks(op);
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestIRVisitorsPass() {
+  PassRegistration<TestIRVisitorsPass>("test-ir-visitors",
+                                       "Test various visitors.");
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDialect(DialectRegistry &);
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestGpuParallelLoopMappingPass();
+void registerTestIRVisitorsPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
#if MLIR_ROCM_CONVERSIONS_ENABLED
   test::registerTestConvertGPUKernelToHsacoPass();
 #endif
   test::registerTestConvVectorization();
   test::registerTestDecomposeCallGraphTypes();
   test::registerTestDominancePass();
   test::registerTestDynamicPipelinePass();
   test::registerTestExpandTanhPass();
   test::registerTestGpuParallelLoopMappingPass();
+  test::registerTestIRVisitorsPass();
   test::registerTestInterfaces();
   test::registerTestLinalgCodegenStrategy();
   test::registerTestLinalgFusionTransforms();
   test::registerTestLinalgTensorFusionTransforms();
   test::registerTestLinalgGreedyFusion();
   test::registerTestLinalgHoisting();
   test::registerTestLinalgTileAndFuseSequencePass();
   test::registerTestLinalgTransforms();