[Mlir-commits] [mlir] adb6494 - [MLIR] Add generic walk support to OpState

[Rahul Joshi]

Author: Rahul Joshi
Date: 2022-01-24T13:35:29-08:00
New Revision: adb6494660eb234d009fe333e65bf94e8becf955

URL: https://github.com/llvm/llvm-project/commit/adb6494660eb234d009fe333e65bf94e8becf955
DIFF: https://github.com/llvm/llvm-project/commit/adb6494660eb234d009fe333e65bf94e8becf955.diff

LOG: [MLIR] Add generic walk support to OpState

- This allows calling the generic walkers on specific operation instances.

Differential Revision: https://reviews.llvm.org/D117949

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpDefinition.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index c68c6b12bf0a..d1400c855783 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -157,10 +157,40 @@ class OpState {
   /// See Operation::walk for more details.
   template <WalkOrder Order = WalkOrder::PostOrder, typename FnT,
             typename RetT = detail::walkResultType<FnT>>
-  RetT walk(FnT &&callback) {
+  typename std::enable_if<
+      llvm::function_traits<std::decay_t<FnT>>::num_args == 1, RetT>::type
+  walk(FnT &&callback) {
     return state->walk<Order>(std::forward<FnT>(callback));
   }
 
+  /// Generic walker with a stage aware callback. Walk the operation by calling
+  /// the callback for each nested operation (including this one) N+1 times,
+  /// where N is the number of regions attached to that operation.
+  ///
+  /// The callback method can take any of the following forms:
+  ///   void(Operation *, const WalkStage &) : Walk all operation opaquely
+  ///     * op.walk([](Operation *nestedOp, const WalkStage &stage) { ...});
+  ///   void(OpT, const WalkStage &) : Walk all operations of the given derived
+  ///                                  type.
+  ///     * op.walk([](ReturnOp returnOp, const WalkStage &stage) { ...});
+  ///   WalkResult(Operation*|OpT, const WalkStage &stage) : Walk operations,
+  ///          but allow for interruption/skipping.
+  ///     * op.walk([](... op, const WalkStage &stage) {
+  ///         // 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 <typename FnT, typename RetT = detail::walkResultType<FnT>>
+  typename std::enable_if<
+      llvm::function_traits<std::decay_t<FnT>>::num_args == 2, RetT>::type
+  walk(FnT &&callback) {
+    return state->walk(std::forward<FnT>(callback));
+  }
+
   // These are default implementations of customization hooks.
 public:
   /// This hook returns any canonicalization pattern rewrites that the operation