[mlir] Extend Operation visitor with pre-order traversal

[dcaballe]

Summary

This patch extends the Region, Block and Operation visitors to also support pre-order walks.
We introduce a new template argument that dictates the walk order (only pre-order and
post-order are supported for now). The default order for Regions, Blocks and Operations is
post-order. Mixed orders (e.g., Region/Block pre-order + Operation post-order) could easily
be implemented, as shown in NumberOfExecutions.cpp.

[dcaballe]

mlir/lib/IR/Visitors.cpp

43

Walk implementation for Region and Block callbacks above seems to be already happening in pre-order. Is this expected? I found this a bit confusing since the traversal for Operation callbacks is in post-order by default.

1. We could change the default traversal for Region and Block callbacks to post-order, if that made any sense, and let them take the IsPreOrder argument.

2. Otherwise, we could keep the current behavior and separate the public Region/Block walk implementation in Visitor.h from the Operation implementation so that only the latter takes the IsPreOrder template argument.

WDYT?

[jurahul]

Please also see: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/utils/visitor_util.h

[rriddle]

Looks mostly good. My main comment is that I would not use a boolean for the template argument, it is very ambiguous what true/false means. Can you use something like an enum instead?

[rriddle]

mlir/lib/IR/Visitors.cpp

43

I would ensure that everything has the same default behavior. So I think option 1 (If I'm reading that right).

[dcaballe]

Please also see: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/utils/visitor_util.h

Thanks for the pointer! It would be great to have something as generic as that visitor in the mlir repo!

Can you use something like an enum instead?

Sure!

mlir/lib/IR/Visitors.cpp

43

Yes, that's #1. I'll give it a try. Hopefully, changing the current default behavior doesn't require too many changes.

[frgossen]

Very useful. Thanks!

mlir/lib/IR/Visitors.cpp

43

I found this strange, too. I believe the liveness analysis, or at least its tests, depend on this behaviour.

#1 sounds like the way to go.

[dcaballe]

Addressing feedback:

• Use enum to encode walk order.
• Set region, block and operation walks to post-order by default.
• Adjust walks in Liveness and NumberOfExecutions to align with the new default traversal.

[rriddle]

LGTM after resolving last comment.

mlir/lib/IR/Visitors.cpp

16

Can you just pass the order as an extra parameter? These methods represent the internal implementation of the walk methods, so I'm fine with adding more parameters here. That would also remove all of the templates within this file.

[bondhugula]

mlir/lib/IR/Visitors.cpp

85–86

There is a typo here: op -> block.

92

Side note: why aren't we using early increment here? This can be done in separate revision but it could be dangerous/annoying to have subtly different iterator invalidation semantics for these different versions. Also, I think these semantics aren't documented on the walk methods.

We have the same issue at at least one more instance above as well.

[bondhugula]

This is nice! Thanks. Just a typo I caught. I think we should fix the iteration invalidation semantics in a subsequent revision (I can do that if you want.)

[dcaballe]

• Passing order as an extra parameter.
• Fixing typo.

I'll commit this tomorrow if no further comments.

mlir/lib/IR/Visitors.cpp

85–86

removed

92

Good point! IIUC, the callback here is on the block. We invoke the callback before or after iterating over all the block's operations so the operation iteration wouldn't be invalidated even if the callback erased block's operations. It could invalidate the block iterator, though, if the block is erased... Yeah, it's a bit inconsistent. Maybe we should use early increments for all the region/block/operation iterators? I could quickly take care of it in a separate patch.

[dcaballe]

Let me revisit a few things in this patch and come back.

[dcaballe]

This new diff introduces important changes to the previous version.
Hopefully, @rriddle/@frgossen/@bondhugula can take another look.

• Added new TestIRVisitors pass with basic testing for the IR visitors (I didn't feel very comfortable without it).
• Introduced 'Prune' as a new WalkResult. It should be used, for example, to skip the walk of an operation that has been erased in pre-order. Otherwise, the walk would recurse into the erased operation (thanks for reporting the issue, @sgrechanik!).
• Addressed @bondhugula's comment on early increment ranges. They are only needed for block and operation ranges since regions can't be directly erased (right?).
• Updated documentation to include erasure information and walk order information. The walk order information was misleading since it only refers to nested operations/block/regions. Operations in a block are always visited in lexicographical order regardless of the nesting order. Same for all the blocks in a region or all the regions in an op.

Please, let me know if the new doc makes more sense. We can iterate on it.
Thanks!

[rriddle]

Can you separate out the prune functionality into a different commit?

[dcaballe]

Can you separate out the prune functionality into a different commit?

I could commit the previous version of the patch as is (already approved) and then create a new review with the new changes. Would that work better?
If I removed only the prune part, the new testing pass would fail since we need the prune functionality for pre-order walks that erase ops.

[rriddle]

In D97217#2597502, @dcaballe wrote:

Can you separate out the prune functionality into a different commit?

I could commit the previous version of the patch as is (already approved) and then create a new review with the new changes. Would that work better?
If I removed only the prune part, the new testing pass would fail since we need the prune functionality for pre-order walks that erase ops.

Yes, I would prefer to review the pruning stuff separately.

[dcaballe]

Restoring previous approved version. I'll commit it if no more comments.
I'll create a separate review for the prune changes.

[rriddle]

LGTM

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 in postorder, calling the callback for
-  /// each operation.
+  /// 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).
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return walk(begin(), end(), std::forward<FnT>(callback));
+    return walk<Order>(begin(), end(), std::forward<FnT>(callback));
   }
 
-  /// Walk the operations in the specified [begin, end) range of this block in
-  /// postorder, calling the callback for each operation. This method is invoked
-  /// for void return callbacks.
+  /// 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.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  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(&op, callback);
+      detail::walk<Order>(&op, callback);
   }
 
-  /// Walk the operations in the specified [begin, end) range of this block in
-  /// postorder, calling the callback for each operation. This method is invoked
-  /// for interruptible callbacks.
+  /// 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.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  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(&op, callback).wasInterrupted())
+      if (detail::walk<Order>(&op, callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 

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 in postorder, calling the callback for each nested
-  /// operation(including this one).
+  /// 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).
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return state->walk(std::forward<FnT>(callback));
+    return state->walk<Order>(std::forward<FnT>(callback));
   }
 
   // These are default implementations of customization hooks.
 public:
   /// This hook returns any canonicalization pattern rewrites that the operation
   /// supports, for use by the canonicalization pass.
   static void getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {}

mlir/include/mlir/IR/Operation.h

template <template <typename T> class Trait> bool mightHaveTrait() {
     const AbstractOperation *abstractOp = getAbstractOperation();
     return abstractOp ? abstractOp->hasTrait<Trait>() : true;
   }
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
-  /// Walk the operation in postorder, calling the callback for each nested
-  /// operation(including this one). The callback method can take any of the
-  /// following forms:
+  /// 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:
   ///   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.
   ///     * op->walk([](... op) {
   ///         // Interrupt, i.e cancel, the walk based on some invariant.
   ///         if (some_invariant)
   ///           return WalkResult::interrupt();
   ///         return WalkResult::advance();
   ///       });
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return detail::walk(this, std::forward<FnT>(callback));
+    return detail::walk<Order>(this, std::forward<FnT>(callback));
   }
 
   //===--------------------------------------------------------------------===//
   // Uses
   //===--------------------------------------------------------------------===//
 
   /// Drop all uses of results of this operation.
   void dropAllUses() {

mlir/include/mlir/IR/Region.h

public:
   /// they are to be deleted.
   void dropAllReferences();
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operations in this region in postorder, calling the callback for
-  /// each operation. This method is invoked for void-returning callbacks.
+  /// 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.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  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(callback);
+      block.walk<Order>(callback);
   }
 
   /// Walk the operations in this region in postorder, calling the callback for
-  /// each operation. This method is invoked for interruptible callbacks.
+  /// 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.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  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(callback).wasInterrupted())
+      if (block.walk<Order>(callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // CFG view utilities
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Visitors.h

public:
 
   static WalkResult interrupt() { return {Interrupt}; }
   static WalkResult advance() { return {Advance}; }
 
   /// Returns true if the walk was interrupted.
   bool wasInterrupted() const { return result == Interrupt; }
 };
 
+/// 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.
-void walk(Operation *op, function_ref<void(Region *)> callback);
-void walk(Operation *op, function_ref<void(Block *)> callback);
-void walk(Operation *op, function_ref<void(Operation *)> callback);
-
+/// the given operation. The walk order is specified by 'order'.
+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. These functions walk until an interrupt result is
-/// returned by the callback.
-WalkResult walk(Operation *op, function_ref<WalkResult(Region *)> callback);
-WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback);
-WalkResult walk(Operation *op, function_ref<WalkResult(Operation *)> callback);
+/// the given operation. The walk order is specified by 'order'. These functions
+/// walk until an interrupt result is returned by the callback.
+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. This method is selected for callbacks that operate on
+/// 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*.
 ///
 /// Example:
 ///   op->walk([](Region *r) { ... });
 ///   op->walk([](Block *b) { ... });
 ///   op->walk([](Operation *op) { ... });
 template <
-    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    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 walk(op, function_ref<RetT(ArgT)>(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. This method is selected for void returning callbacks that
-/// operate on a specific derived operation type.
+/// 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.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) { ... });
 template <
-    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    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));
+  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. This method is selected for WalkReturn returning
-/// interruptible callbacks that operate on a specific derived operation type.
+/// 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.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) {
 ///     if (some_invariant)
 ///       return WalkResult::interrupt();
 ///     return WalkResult::advance();
 ///   });
 template <
-    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    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, WalkResult>::value,
     RetT>::type
 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));
+  return detail::walk(op, function_ref<RetT(Operation *)>(wrapperFn), Order);
 }
 
 /// Utility to provide the return type of a templated walk method.
 template <typename FnT>
 using walkResultType = decltype(walk(nullptr, std::declval<FnT>()));
 } // end namespace detail
 
 } // namespace mlir
 
 #endif

mlir/lib/Analysis/Liveness.cpp

 };
 } // namespace
 
 /// Builds the internal liveness block mapping.
 static void buildBlockMapping(Operation *operation,
                               DenseMap<Block *, BlockInfoBuilder> &builders) {
   llvm::SetVector<Block *> toProcess;
 
-  operation->walk([&](Block *block) {
+  operation->walk<WalkOrder::PreOrder>([&](Block *block) {
     BlockInfoBuilder &builder =
         builders.try_emplace(block, block).first->second;
 
     if (builder.updateLiveIn())
       toProcess.insert(block->pred_begin(), block->pred_end());
   });
 
   // Propagate the in and out-value sets (fixpoint iteration)

 /// Dumps the liveness information to the given stream.
 void Liveness::print(raw_ostream &os) const {
   os << "// ---- Liveness -----\n";
 
   // Builds unique block/value mappings for testing purposes.
   DenseMap<Block *, size_t> blockIds;
   DenseMap<Operation *, size_t> operationIds;
   DenseMap<Value, size_t> valueIds;
-  operation->walk([&](Block *block) {
+  operation->walk<WalkOrder::PreOrder>([&](Block *block) {
     blockIds.insert({block, blockIds.size()});
     for (BlockArgument argument : block->getArguments())
       valueIds.insert({argument, valueIds.size()});
     for (Operation &operation : *block) {
       operationIds.insert({&operation, operationIds.size()});
       for (Value result : operation.getResults())
         valueIds.insert({result, valueIds.size()});
     }
std::sort(orderedValues.begin(), orderedValues.end(),
               [&](Value left, Value right) {
                 return valueIds[left] < valueIds[right];
               });
     for (Value value : orderedValues)
       printValueRef(value);
   };
 
   // Dump information about in and out values.
-  operation->walk([&](Block *block) {
+  operation->walk<WalkOrder::PreOrder>([&](Block *block) {
     os << "// - Block: " << blockIds[block] << "\n";
     const auto *liveness = getLiveness(block);
     os << "// --- LiveIn: ";
     printValueRefs(liveness->inValues);
     os << "\n// --- LiveOut: ";
     printValueRefs(liveness->outValues);
     os << "\n";
 

mlir/lib/Analysis/NumberOfExecutions.cpp

if (blockInfo.count(&block) == 0)
       blockInfo.insert({&block, BlockNumberOfExecutionsInfo(
                                     &block, numRegionInvocations,
                                     /*numberOfBlockExecutions=*/None)});
 }
 
 /// Creates a new NumberOfExecutions analysis that computes how many times a
 /// block within a region is executed for all associated regions.
 NumberOfExecutions::NumberOfExecutions(Operation *op) : operation(op) {
-  operation->walk([&](Region *region) {
+  operation->walk<WalkOrder::PreOrder>([&](Region *region) {
     computeRegionBlockNumberOfExecutions(*region, blockNumbersOfExecution);
   });
 }
 
 Optional<int64_t>
 NumberOfExecutions::getNumberOfExecutions(Operation *op,
                                           Region *perEntryOfThisRegion) const {
   // Assuming that all operations complete in a finite amount of time (do not
NumberOfExecutions::getNumberOfExecutions(Block *block,
 
   return numberOfExecutions;
 }
 
 void NumberOfExecutions::printBlockExecutions(
     raw_ostream &os, Region *perEntryOfThisRegion) const {
   unsigned blockId = 0;
 
-  operation->walk([&](Block *block) {
+  operation->walk<WalkOrder::PreOrder>([&](Block *block) {
     llvm::errs() << "Block: " << blockId++ << "\n";
     llvm::errs() << "Number of executions: ";
     if (auto n = getNumberOfExecutions(block, perEntryOfThisRegion))
       llvm::errs() << *n << "\n";
     else
       llvm::errs() << "<unknown>\n";
   });
 }
 
 void NumberOfExecutions::printOperationExecutions(
     raw_ostream &os, Region *perEntryOfThisRegion) const {
-  operation->walk([&](Block *block) {
+  operation->walk<WalkOrder::PreOrder>([&](Block *block) {
     block->walk([&](Operation *operation) {
       // Skip the operation that was used to build the analysis.
       if (operation == this->operation)
         return;
 
       llvm::errs() << "Operation: " << operation->getName() << "\n";
       llvm::errs() << "Number of executions: ";
       if (auto n = getNumberOfExecutions(operation, perEntryOfThisRegion))

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.
-void detail::walk(Operation *op, function_ref<void(Region *)> callback) {
+/// given operation. The walk order is specified by 'Order'.
+
+void detail::walk(Operation *op, function_ref<void(Region *)> callback,
+                  WalkOrder order) {
   for (auto &region : op->getRegions()) {
+    if (order == WalkOrder::PreOrder)
       callback(&region);
     for (auto &block : region) {
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk(&nestedOp, callback, order);
     }
+    if (order == WalkOrder::PostOrder)
+      callback(&region);
   }
 }
 
-void detail::walk(Operation *op, function_ref<void(Block *)> callback) {
+void detail::walk(Operation *op, function_ref<void(Block *)> callback,
+                  WalkOrder order) {
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
+      if (order == WalkOrder::PreOrder)
         callback(&block);
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk(&nestedOp, callback, order);
+      if (order == WalkOrder::PostOrder)
+        callback(&block);
     }
   }
 }
 
-void detail::walk(Operation *op, function_ref<void(Operation *op)> callback) {
+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) {
       // Early increment here in the case where the operation is erased.
       for (auto &nestedOp : llvm::make_early_inc_range(block))
-        walk(&nestedOp, callback);
+        walk(&nestedOp, callback, order);
     }
   }
+
+  if (order == WalkOrder::PostOrder)
     callback(op);
 }
 
 /// Walk all of the regions/blocks/operations nested under and including the
-/// given operation. These functions walk operations until an interrupt result
-/// is returned by the callback.
+/// given operation. The walk order is specified by 'order'. These functions
+/// walk operations until an interrupt result is returned by the callback.
 WalkResult detail::walk(Operation *op,
-                        function_ref<WalkResult(Region *op)> callback) {
+                        function_ref<WalkResult(Region *)> callback,
+                        WalkOrder order) {
   for (auto &region : op->getRegions()) {
+    if (order == WalkOrder::PreOrder)
       if (callback(&region).wasInterrupted())
         return WalkResult::interrupt();
     for (auto &block : region) {
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk(&nestedOp, callback, order);
     }
+    if (order == WalkOrder::PostOrder)
+      if (callback(&region).wasInterrupted())
+        return WalkResult::interrupt();
   }
   return WalkResult::advance();
 }
 
 WalkResult detail::walk(Operation *op,
-                        function_ref<WalkResult(Block *op)> callback) {
+                        function_ref<WalkResult(Block *)> callback,
+                        WalkOrder order) {
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
+      if (order == WalkOrder::PreOrder)
         if (callback(&block).wasInterrupted())
           return WalkResult::interrupt();
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk(&nestedOp, callback, order);
+      if (order == WalkOrder::PostOrder)
+        if (callback(&block).wasInterrupted())
+          return WalkResult::interrupt();
     }
   }
   return WalkResult::advance();
 }
 
 WalkResult detail::walk(Operation *op,
-                        function_ref<WalkResult(Operation *op)> callback) {
+                        function_ref<WalkResult(Operation *)> callback,
+                        WalkOrder order) {
+  if (order == WalkOrder::PreOrder)
+    if (callback(op).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).wasInterrupted())
+        if (walk(&nestedOp, callback, order).wasInterrupted())
           return WalkResult::interrupt();
       }
     }
   }
+
+  if (order == WalkOrder::PostOrder)
     return callback(op);
+  return WalkResult::advance();
 }