[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

39

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

39

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

39

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

39

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

17

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

86

There is a typo here: op -> block.

91

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

86

removed

91

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:
 
   //===--------------------------------------------------------------------===//
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
   /// Walk the operations in this block in postorder, calling the callback for
   /// each operation.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <bool IsPreOrder = false, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return walk(begin(), end(), std::forward<FnT>(callback));
+    return walk<IsPreOrder>(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.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <bool IsPreOrder = false, 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<IsPreOrder>(&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.
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <bool IsPreOrder = false, 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<IsPreOrder>(&op, callback).wasInterrupted())
         return WalkResult::interrupt();
     return WalkResult::advance();
   }
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/OpDefinition.h

public:
 
   /// 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).
   /// See Operation::walk for more details.
-  template <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  template <bool IsPreOrder = false, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return state->walk(std::forward<FnT>(callback));
+    return state->walk<IsPreOrder>(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

public:
   ///   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 <bool IsPreOrder = false, typename FnT,
+            typename RetT = detail::walkResultType<FnT>>
   RetT walk(FnT &&callback) {
-    return detail::walk(this, std::forward<FnT>(callback));
+    return detail::walk<IsPreOrder>(this, std::forward<FnT>(callback));
   }
 
   //===--------------------------------------------------------------------===//
   // Uses
   //===--------------------------------------------------------------------===//
 
   /// Drop all uses of results of this operation.
   void dropAllUses() {

mlir/include/mlir/IR/Region.h

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

mlir/include/mlir/IR/Visitors.h

 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.
+template <bool IsPreOrder>
 void walk(Operation *op, function_ref<void(Region *)> callback);
+template <bool IsPreOrder>
 void walk(Operation *op, function_ref<void(Block *)> callback);
+template <bool IsPreOrder>
 void walk(Operation *op, function_ref<void(Operation *)> callback);
 
 /// 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.
+template <bool IsPreOrder>
 WalkResult walk(Operation *op, function_ref<WalkResult(Region *)> callback);
+template <bool IsPreOrder>
 WalkResult walk(Operation *op, function_ref<WalkResult(Block *)> callback);
+template <bool IsPreOrder>
 WalkResult walk(Operation *op, function_ref<WalkResult(Operation *)> callback);
 
 // 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
 /// 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>,
+    bool IsPreOrder = false, 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 walk<IsPreOrder>(op, function_ref<RetT(ArgT)>(callback));
 }
 
 /// 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.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) { ... });
 template <
-    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    bool IsPreOrder = false, 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<IsPreOrder>(op,
+                                  function_ref<RetT(Operation *)>(wrapperFn));
 }
 
 /// 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.
 ///
 /// Example:
 ///   op->walk([](ReturnOp op) {
 ///     if (some_invariant)
 ///       return WalkResult::interrupt();
 ///     return WalkResult::advance();
 ///   });
 template <
-    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    bool IsPreOrder = false, 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<IsPreOrder>(op,
+                                  function_ref<RetT(Operation *)>(wrapperFn));
 }
 
 /// 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/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.
+
+template <bool IsPreOrder>
 void detail::walk(Operation *op, function_ref<void(Region *)> callback) {
   for (auto &region : op->getRegions()) {
     callback(&region);
     for (auto &block : region) {
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk<IsPreOrder>(&nestedOp, callback);
     }
   }
 }
+template void detail::walk<true>(Operation *, function_ref<void(Region *)>);
+template void detail::walk<false>(Operation *, function_ref<void(Region *)>);
 
+template <bool IsPreOrder>
 void detail::walk(Operation *op, function_ref<void(Block *)> callback) {
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
       callback(&block);
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk<IsPreOrder>(&nestedOp, callback);
     }
   }
 }
+template void detail::walk<true>(Operation *, function_ref<void(Block *)>);
+template void detail::walk<false>(Operation *, function_ref<void(Block *)>);
 
+template <bool IsPreOrder>
 void detail::walk(Operation *op, function_ref<void(Operation *op)> callback) {
+  if (IsPreOrder)
+    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<IsPreOrder>(&nestedOp, callback);
     }
   }
+
+  if (!IsPreOrder)
     callback(op);
 }
+template void detail::walk<true>(Operation *, function_ref<void(Operation *)>);
+template void detail::walk<false>(Operation *, function_ref<void(Operation *)>);
 
 /// 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.
+template <bool IsPreOrder>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Region *op)> callback) {
   for (auto &region : op->getRegions()) {
     if (callback(&region).wasInterrupted())
       return WalkResult::interrupt();
     for (auto &block : region) {
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk<IsPreOrder>(&nestedOp, callback);
     }
   }
   return WalkResult::advance();
 }
+template WalkResult detail::walk<true>(Operation *,
+                                       function_ref<WalkResult(Region *)>);
+template WalkResult detail::walk<false>(Operation *,
+                                        function_ref<WalkResult(Region *)>);
 
+template <bool IsPreOrder>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Block *op)> callback) {
   for (auto &region : op->getRegions()) {
     for (auto &block : region) {
       if (callback(&block).wasInterrupted())
         return WalkResult::interrupt();
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback);
+        walk<IsPreOrder>(&nestedOp, callback);
     }
   }
   return WalkResult::advance();
 }
+template WalkResult detail::walk<true>(Operation *,
+                                       function_ref<WalkResult(Block *)>);
+template WalkResult detail::walk<false>(Operation *,
+                                        function_ref<WalkResult(Block *)>);
 
+template <bool IsPreOrder>
 WalkResult detail::walk(Operation *op,
                         function_ref<WalkResult(Operation *op)> callback) {
+  if (IsPreOrder)
+    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<IsPreOrder>(&nestedOp, callback).wasInterrupted())
           return WalkResult::interrupt();
       }
     }
   }
+
+  if (!IsPreOrder)
     return callback(op);
+  return WalkResult::advance();
 }
+template WalkResult detail::walk<true>(Operation *,
+                                       function_ref<WalkResult(Operation *)>);
+template WalkResult detail::walk<false>(Operation *,
+                                        function_ref<WalkResult(Operation *)>);