Refactor the 'walk' methods for operations.

This change refactors and cleans up the implementation of the operation walk methods. After this refactoring is that the explicit template parameter for the operation type is no longer needed for the explicit op walks. For example:

    op->walk<AffineForOp>([](AffineForOp op) { ... });

is now accomplished via:

    op->walk([](AffineForOp op) { ... });

PiperOrigin-RevId: 266209552

mlir/examples/Linalg/Linalg2/Example.cpp

@@ -104,7 +104,7 @@ TEST_FUNC(linalg_ops_folded_slices) {
   //  CHECK-NEXT: linalg.dot({{.*}}, {{.*}}, {{.*}}) : !linalg.view<f32>
   // clang-format on
 
-  f.walk<SliceOp>([](SliceOp slice) {
+  f.walk([](SliceOp slice) {
     auto *sliceResult = slice.getResult();
     auto viewOp = emitAndReturnFullyComposedView(sliceResult);
     sliceResult->replaceAllUsesWith(viewOp.getResult());

mlir/examples/Linalg/Linalg3/lib/Transforms.cpp

@@ -36,7 +36,7 @@ using namespace linalg;
 using namespace linalg::intrinsics;
 
 void linalg::composeSliceOps(mlir::FuncOp f) {
-  f.walk<SliceOp>([](SliceOp sliceOp) {
+  f.walk([](SliceOp sliceOp) {
     auto *sliceResult = sliceOp.getResult();
     auto viewOp = emitAndReturnFullyComposedView(sliceResult);
     sliceResult->replaceAllUsesWith(viewOp.getResult());

mlir/include/mlir/IR/Block.h

@@ -23,6 +23,7 @@
 #define MLIR_IR_BLOCK_H
 
 #include "mlir/IR/Value.h"
+#include "mlir/IR/Visitors.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/ilist_node.h"
@@ -289,20 +290,20 @@ public:
 
   /// Walk the operations in this block in postorder, calling the callback for
   /// each operation.
-  void walk(llvm::function_ref<void(Operation *)> callback);
-
-  /// Specialization of walk to only visit operations of 'OpTy'.
-  template <typename OpTy> void walk(llvm::function_ref<void(OpTy)> callback) {
-    walk([&](Operation *opInst) {
-      if (auto op = dyn_cast<OpTy>(opInst))
-        callback(op);
-    });
+  /// See Operation::walk for more details.
+  template <typename FnT> void walk(FnT &&callback) {
+    return walk(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.
-  void walk(Block::iterator begin, Block::iterator end,
-            llvm::function_ref<void(Operation *)> callback);
+  /// postorder, calling the callback for each operation. This method is invoked
+  /// for void return callbacks.
+  /// See Operation::walk for more details.
+  template <typename FnT>
+  void walk(Block::iterator begin, Block::iterator end, FnT &&callback) {
+    for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
+      detail::walkOperations(&op, std::forward<FnT>(callback));
+  }
 
   //===--------------------------------------------------------------------===//
   // Other

mlir/include/mlir/IR/OpDefinition.h

@@ -189,16 +189,9 @@ public:
 
   /// Walk the operation in postorder, calling the callback for each nested
   /// operation(including this one).
-  void walk(llvm::function_ref<void(Operation *)> callback) {
-    state->walk(callback);
-  }
-
-  /// Specialization of walk to only visit operations of 'OpTy'.
-  template <typename OpTy> void walk(llvm::function_ref<void(OpTy)> callback) {
-    walk([&](Operation *opInst) {
-      if (auto op = dyn_cast<OpTy>(opInst))
-        callback(op);
-    });
+  /// See Operation::walk for more details.
+  template <typename FnT> void walk(FnT &&callback) {
+    state->walk(std::forward<FnT>(callback));
   }
 
   // These are default implementations of customization hooks.

mlir/include/mlir/IR/Operation.h

@@ -491,16 +491,15 @@ public:
   // Operation Walkers
   //===--------------------------------------------------------------------===//
 
-  /// Walk this operation in postorder, calling the callback for each operation
-  /// including this one.
-  void walk(llvm::function_ref<void(Operation *)> callback);
-
-  /// Specialization of walk to only visit operations of 'T'.
-  template <typename T> void walk(llvm::function_ref<void(T)> callback) {
-    walk([&](Operation *op) {
-      if (auto derivedOp = dyn_cast<T>(op))
-        callback(derivedOp);
-    });
+  /// 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:
+  ///   (void)(Operation*) : Walk all operations opaquely.
+  ///     * op->walk([](Operation *nestedOp) { ...});
+  ///   (void)(OpT) : Walk all operations of the given derived type.
+  ///     * op->walk([](ReturnOp returnOp) { ...});
+  template <typename FnT> void walk(FnT &&callback) {
+    detail::walkOperations(this, std::forward<FnT>(callback));
   }
 
   //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/Region.h

@@ -124,9 +124,13 @@ public:
   /// they are to be deleted.
   void dropAllReferences();
 
-  /// Walk the operations in this block in postorder, calling the callback for
+  /// Walk the operations in this region in postorder, calling the callback for
   /// each operation.
-  void walk(llvm::function_ref<void(Operation *)> callback);
+  /// See Operation::walk for more details.
+  template <typename FnT> void walk(FnT &&callback) {
+    for (auto &block : *this)
+      block.walk(callback);
+  }
 
   /// Displays the CFG in a window. This is for use from the debugger and
   /// depends on Graphviz to generate the graph.

mlir/include/mlir/IR/Visitors.h

@@ -0,0 +1,89 @@
+//===- Visitors.h - Utilities for visiting operations -----------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines utilities for walking and visiting operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_VISITORS_H
+#define MLIR_IR_VISITORS_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace mlir {
+class Operation;
+
+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 operations nested under and including the given operation.
+void walkOperations(Operation *op, function_ref<void(Operation *op)> 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 dynamically dispatching based
+// upon the type of the callback function.
+
+/// Walk all of the operations nested under and including the given operation.
+/// This method is selected for callbacks that operation on Operation*.
+///
+/// Example:
+///   op->walk([](Operation *op) { ... });
+template <
+    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
+typename std::enable_if<std::is_same<ArgT, Operation *>::value, RetT>::type
+walkOperations(Operation *op, FuncTy &&callback) {
+  return detail::walkOperations(op, llvm::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
+/// operation on a specific derived operation type.
+///
+/// Example:
+///   op->walk([](ReturnOp op) { ... });
+template <
+    typename FuncTy, typename ArgT = detail::first_argument<FuncTy>,
+    typename RetT = decltype(std::declval<FuncTy>()(std::declval<ArgT>()))>
+typename std::enable_if<!std::is_same<ArgT, Operation *>::value &&
+                            std::is_same<RetT, void>::value,
+                        RetT>::type
+walkOperations(Operation *op, FuncTy &&callback) {
+  auto wrapperFn = [&](Operation *op) {
+    if (auto derivedOp = dyn_cast<ArgT>(op))
+      callback(derivedOp);
+  };
+  return detail::walkOperations(op, function_ref<RetT(Operation *)>(wrapperFn));
+}
+} // end namespace detail
+
+} // namespace mlir
+
+#endif

mlir/include/mlir/Support/LLVM.h

@@ -55,6 +55,7 @@ template <typename T> struct DenseMapInfo;
 template <typename ValueT, typename ValueInfoT> class DenseSet;
 template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
 class DenseMap;
+template <typename Fn> class function_ref;
 
 // Other common classes.
 class raw_ostream;
@@ -80,6 +81,7 @@ template <typename KeyT, typename ValueT,
 using DenseMap = llvm::DenseMap<KeyT, ValueT, KeyInfoT, BucketT>;
 template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
 using DenseSet = llvm::DenseSet<ValueT, ValueInfoT>;
+template <typename Fn> using function_ref = llvm::function_ref<Fn>;
 using llvm::MutableArrayRef;
 using llvm::None;
 using llvm::Optional;

mlir/lib/Analysis/TestParallelismDetection.cpp

@@ -45,7 +45,7 @@ FunctionPassBase *mlir::createParallelismDetectionTestPass() {
 void TestParallelismDetection::runOnFunction() {
   FuncOp f = getFunction();
   OpBuilder b(f.getBody());
-  f.walk<AffineForOp>([&](AffineForOp forOp) {
+  f.walk([&](AffineForOp forOp) {
     if (isLoopParallel(forOp))
       forOp.emitRemark("parallel loop");
     else

mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp

@@ -131,7 +131,7 @@ public:
     initializeCachedTypes();
 
     for (auto func : getModule().getOps<FuncOp>()) {
-      func.walk<mlir::gpu::LaunchFuncOp>(
+      func.walk(
           [this](mlir::gpu::LaunchFuncOp op) { translateGpuLaunchCalls(op); });
     }
   }

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

@@ -71,7 +71,7 @@ static FuncOp outlineKernelFunc(gpu::LaunchOp launchOp) {
   outlinedFunc.setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
                        builder.getUnitAttr());
   injectGpuIndexOperations(loc, outlinedFunc);
-  outlinedFunc.walk<mlir::gpu::Return>([](mlir::gpu::Return op) {
+  outlinedFunc.walk([](mlir::gpu::Return op) {
     OpBuilder replacer(op);
     replacer.create<ReturnOp>(op.getLoc());
     op.erase();
@@ -98,7 +98,7 @@ public:
   void runOnModule() override {
     ModuleManager moduleManager(getModule());
     for (auto func : getModule().getOps<FuncOp>()) {
-      func.walk<mlir::gpu::LaunchOp>([&](mlir::gpu::LaunchOp op) {
+      func.walk([&](mlir::gpu::LaunchOp op) {
         FuncOp outlinedFunc = outlineKernelFunc(op);
         moduleManager.insert(outlinedFunc);
         convertToLaunchFuncOp(op, outlinedFunc);

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

@@ -239,8 +239,7 @@ static void fuseLinalgOps(FuncOp f, ArrayRef<int64_t> tileSizes) {
 
   // 1. Record the linalg ops so we can traverse them in reverse order.
   SmallVector<Operation *, 8> linalgOps;
-  f.walk<LinalgOp>(
-      [&](LinalgOp op) { linalgOps.push_back(op.getOperation()); });
+  f.walk([&](LinalgOp op) { linalgOps.push_back(op.getOperation()); });
 
   // 2. Setup the dependences graph, aliases are populated lazily.
   Aliases aliases;

mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp

@@ -877,7 +877,7 @@ struct LowerLinalgToLLVMPass : public ModulePass<LowerLinalgToLLVMPass> {
 // affine will look different than lowering to LLVM and it is still unclear how
 // everything will be eventually structured.
 static void lowerLinalgSubViewOps(FuncOp &f) {
-  f.walk<SubViewOp>([&](SubViewOp op) {
+  f.walk([&](SubViewOp op) {
     OpBuilder b(op);
     ScopedContext scope(b, op.getLoc());
     auto *view = op.getView();

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

@@ -490,7 +490,7 @@ mlir::linalg::tileLinalgOp(LinalgOp op, ArrayRef<int64_t> tileSizes,
 static void tileLinalgOps(FuncOp f, ArrayRef<int64_t> tileSizes,
                           bool promoteViews) {
   OperationFolder folder;
-  f.walk<LinalgOp>([promoteViews, tileSizes, &folder](LinalgOp op) {
+  f.walk([promoteViews, tileSizes, &folder](LinalgOp op) {
     // TODO(ntv) some heuristic here to decide what to promote. Atm it is all or
     // nothing.
     SmallVector<bool, 8> viewsToPromote(op.getNumInputsAndOutputs(),
@@ -500,7 +500,7 @@ static void tileLinalgOps(FuncOp f, ArrayRef<int64_t> tileSizes,
     if (opLoopsPair)
       op.erase();
   });
-  f.walk<LinalgOp>([](LinalgOp op) {
+  f.walk([](LinalgOp op) {
     if (!op.getOperation()->hasNoSideEffect())
       return;
     if (op.getOperation()->use_empty())

mlir/lib/IR/Block.cpp

@@ -224,22 +224,6 @@ Block *Block::getSinglePredecessor() {
   return it == pred_end() ? firstPred : nullptr;
 }
 
-//===----------------------------------------------------------------------===//
-// Operation Walkers
-//===----------------------------------------------------------------------===//
-
-void Block::walk(llvm::function_ref<void(Operation *)> callback) {
-  walk(begin(), end(), callback);
-}
-
-/// Walk the operations in the specified [begin, end) range of this block,
-/// calling the callback for each operation.
-void Block::walk(Block::iterator begin, Block::iterator end,
-                 llvm::function_ref<void(Operation *)> callback) {
-  for (auto &op : llvm::make_early_inc_range(llvm::make_range(begin, end)))
-    op.walk(callback);
-}
-
 //===----------------------------------------------------------------------===//
 // Other
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Operation.cpp

@@ -289,19 +289,6 @@ void Operation::replaceUsesOfWith(Value *from, Value *to) {
       operand.set(to);
 }
 
-//===----------------------------------------------------------------------===//
-// Operation Walkers
-//===----------------------------------------------------------------------===//
-
-void Operation::walk(llvm::function_ref<void(Operation *)> callback) {
-  // Visit any internal operations.
-  for (auto &region : getRegions())
-    region.walk(callback);
-
-  // Visit the current operation.
-  callback(this);
-}
-
 //===----------------------------------------------------------------------===//
 // Other
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Region.cpp

@@ -168,13 +168,6 @@ bool Region::isIsolatedFromAbove(llvm::Optional<Location> noteLoc) {
   return isIsolatedAbove(*this, *this, noteLoc);
 }
 
-/// Walk the operations in this block in postorder, calling the callback for
-/// each operation.
-void Region::walk(llvm::function_ref<void(Operation *)> callback) {
-  for (auto &block : *this)
-    block.walk(callback);
-}
-
 Region *llvm::ilist_traits<::mlir::Block>::getParentRegion() {
   size_t Offset(
       size_t(&((Region *)nullptr->*Region::getSublistAccess(nullptr))));

mlir/lib/IR/Visitors.cpp

@@ -0,0 +1,34 @@
+//===- Visitors.cpp - MLIR Visitor Utilties -------------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "mlir/IR/Visitors.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+
+/// Walk all of the operations nested under and including the given operations.
+void detail::walkOperations(Operation *op,
+                            function_ref<void(Operation *op)> callback) {
+  // TODO(b/140235992) 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))
+        walkOperations(&nestedOp, callback);
+
+  callback(op);
+}

mlir/lib/Transforms/LoopCoalescing.cpp

@@ -34,7 +34,7 @@ public:
   void runOnFunction() override {
     FuncOp func = getFunction();
 
-    func.walk<loop::ForOp>([](loop::ForOp op) {
+    func.walk([](loop::ForOp op) {
       // Ignore nested loops.
       if (op.getParentOfType<loop::ForOp>())
         return;

mlir/lib/Transforms/LoopInvariantCodeMotion.cpp

@@ -240,7 +240,7 @@ void LoopInvariantCodeMotion::runOnFunction() {
   // Walk through all loops in a function in innermost-loop-first order.  This
   // way, we first LICM from the inner loop, and place the ops in
   // the outer loop, which in turn can be further LICM'ed.
-  getFunction().walk<AffineForOp>([&](AffineForOp op) {
+  getFunction().walk([&](AffineForOp op) {
     LLVM_DEBUG(op.getOperation()->print(llvm::dbgs() << "\nOriginal loop\n"));
     runOnAffineForOp(op);
   });

mlir/lib/Transforms/LoopUnroll.cpp

@@ -128,7 +128,7 @@ void LoopUnroll::runOnFunction() {
     // Gathers all loops with trip count <= minTripCount. Do a post order walk
     // so that loops are gathered from innermost to outermost (or else unrolling
     // an outer one may delete gathered inner ones).
-    getFunction().walk<AffineForOp>([&](AffineForOp forOp) {
+    getFunction().walk([&](AffineForOp forOp) {
       Optional<uint64_t> tripCount = getConstantTripCount(forOp);
       if (tripCount.hasValue() && tripCount.getValue() <= clUnrollFullThreshold)
         loops.push_back(forOp);

mlir/lib/Transforms/MemRefDataFlowOpt.cpp

@@ -226,8 +226,7 @@ void MemRefDataFlowOpt::runOnFunction() {
   memrefsToErase.clear();
 
   // Walk all load's and perform load/store forwarding.
-  f.walk<AffineLoadOp>(
-      [&](AffineLoadOp loadOp) { forwardStoreToLoad(loadOp); });
+  f.walk([&](AffineLoadOp loadOp) { forwardStoreToLoad(loadOp); });
 
   // Erase all load op's whose results were replaced with store fwd'ed ones.
   for (auto *loadOp : loadOpsToErase) {

mlir/lib/Transforms/PipelineDataTransfer.cpp

@@ -144,8 +144,7 @@ void PipelineDataTransfer::runOnFunction() {
   // gets deleted and replaced by a prologue, a new steady-state loop and an
   // epilogue).
   forOps.clear();
-  getFunction().walk<AffineForOp>(
-      [&](AffineForOp forOp) { forOps.push_back(forOp); });
+  getFunction().walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
   for (auto forOp : forOps)
     runOnAffineForOp(forOp);
 }

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

@@ -258,7 +258,7 @@ FusionResult mlir::canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
 /// returns false otherwise.
 bool mlir::getLoopNestStats(AffineForOp forOpRoot, LoopNestStats *stats) {
   bool ret = true;
-  forOpRoot.getOperation()->walk<AffineForOp>([&](AffineForOp forOp) {
+  forOpRoot.walk([&](AffineForOp forOp) {
     auto *childForOp = forOp.getOperation();
     auto *parentForOp = forOp.getOperation()->getParentOp();
     if (!llvm::isa<FuncOp>(parentForOp)) {

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -165,8 +165,7 @@ LogicalResult mlir::promoteIfSingleIteration(AffineForOp forOp) {
 /// their body into the containing Block.
 void mlir::promoteSingleIterationLoops(FuncOp f) {
   // Gathers all innermost loops through a post order pruned walk.
-  f.walk<AffineForOp>(
-      [](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
+  f.walk([](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
 }
 
 /// Generates a 'affine.for' op with the specified lower and upper bounds

mlir/lib/Transforms/Vectorize.cpp

@@ -1240,7 +1240,7 @@ void Vectorize::runOnFunction() {
   NestedPatternContext mlContext;
 
   llvm::DenseSet<Operation *> parallelLoops;
-  f.walk<AffineForOp>([&parallelLoops](AffineForOp loop) {
+  f.walk([&parallelLoops](AffineForOp loop) {
     if (isLoopParallel(loop))
       parallelLoops.insert(loop);
   });

mlir/test/lib/Transforms/TestLoopMapping.cpp

@@ -49,7 +49,7 @@ public:
       numProcessors.push_back(op->getResult(1));
     });
 
-    func.walk<loop::ForOp>([&processorIds, &numProcessors](loop::ForOp op) {
+    func.walk([&processorIds, &numProcessors](loop::ForOp op) {
       // Ignore nested loops.
       if (op.getParentRegion()->getParentOfType<loop::ForOp>())
         return;

mlir/test/lib/Transforms/TestLoopParametricTiling.cpp

@@ -43,7 +43,7 @@ public:
 
   void runOnFunction() override {
     FuncOp func = getFunction();
-    func.walk<loop::ForOp>([this](loop::ForOp op) {
+    func.walk([this](loop::ForOp op) {
       // Ignore nested loops.
       if (op.getParentRegion()->getParentOfType<loop::ForOp>())
         return;