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

28 files changed, 175 insertions, 93 deletions

diff --git a/mlir/examples/Linalg/Linalg2/Example.cpp b/mlir/examples/Linalg/Linalg2/Example.cpp
index 1490e2bced1..8b9eee49ddb 100644
--- a/mlir/examples/Linalg/Linalg2/Example.cpp
+++ b/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());
diff --git a/mlir/examples/Linalg/Linalg3/lib/Transforms.cpp b/mlir/examples/Linalg/Linalg3/lib/Transforms.cpp
index 8731138bba5..ce2656520fa 100644
--- a/mlir/examples/Linalg/Linalg3/lib/Transforms.cpp
+++ b/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());
diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 84144b89c36..31b51204a4c 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/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
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 570990a2c58..18c4615728d 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/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.
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index d61dc067936..c3f4c0eb4cf 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/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));
   }
 
   //===--------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index c6f97c31fb8..8b7db9a6147 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/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.
diff --git a/mlir/include/mlir/IR/Visitors.h b/mlir/include/mlir/IR/Visitors.h
new file mode 100644
index 00000000000..f1876dfaebc
--- /dev/null
+++ b/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
diff --git a/mlir/include/mlir/Support/LLVM.h b/mlir/include/mlir/Support/LLVM.h
index f0dd1216652..6ae8d5c075c 100644
--- a/mlir/include/mlir/Support/LLVM.h
+++ b/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;
diff --git a/mlir/lib/Analysis/TestParallelismDetection.cpp b/mlir/lib/Analysis/TestParallelismDetection.cpp
index 351a6a7a191..75982a8e0c5 100644
--- a/mlir/lib/Analysis/TestParallelismDetection.cpp
+++ b/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
diff --git a/mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp b/mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp
index d4293ba5a1c..ba0bc475168 100644
--- a/mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp
+++ b/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); });
     }
   }
diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 481ed247e81..dae8ae8ec55 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/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);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index 9f44b1cffa6..d4860642bb0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/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;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp b/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp
index 0bc355ae91d..3c4123eb309 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp
+++ b/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();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index d2495bdfeac..11b3334f250 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/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())
diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 28614ca8bdc..a225c4d03bc 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -225,22 +225,6 @@ Block *Block::getSinglePredecessor() {
 }
 
 //===----------------------------------------------------------------------===//
-// 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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index a623e39aa0c..56c74312007 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -290,19 +290,6 @@ void Operation::replaceUsesOfWith(Value *from, Value *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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index 0947ddd04f3..df0cf077d28 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/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))));
diff --git a/mlir/lib/IR/Visitors.cpp b/mlir/lib/IR/Visitors.cpp
new file mode 100644
index 00000000000..4622098f19f
--- /dev/null
+++ b/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);
+}
diff --git a/mlir/lib/Transforms/LoopCoalescing.cpp b/mlir/lib/Transforms/LoopCoalescing.cpp
index c4024fe303f..8e220607f06 100644
--- a/mlir/lib/Transforms/LoopCoalescing.cpp
+++ b/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;
diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
index 293e565cda7..be39297a4b2 100644
--- a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
+++ b/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);
   });
diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index 2acc5a90f5f..5e132794149 100644
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/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);
diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
index 9b71ada100c..f922d508c69 100644
--- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
+++ b/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) {
diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp
index a814af92a5f..74b06aa368d 100644
--- a/mlir/lib/Transforms/PipelineDataTransfer.cpp
+++ b/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);
 }
diff --git a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp b/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
index 8b314780c9f..99f315e3fd0 100644
--- a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
+++ b/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)) {
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index d6a31f92aed..91f72f3aa1d 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/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
diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp
index cbf616eae10..89e3da7477d 100644
--- a/mlir/lib/Transforms/Vectorize.cpp
+++ b/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);
   });
diff --git a/mlir/test/lib/Transforms/TestLoopMapping.cpp b/mlir/test/lib/Transforms/TestLoopMapping.cpp
index f4aa6469a99..c25fea9aa13 100644
--- a/mlir/test/lib/Transforms/TestLoopMapping.cpp
+++ b/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;
diff --git a/mlir/test/lib/Transforms/TestLoopParametricTiling.cpp b/mlir/test/lib/Transforms/TestLoopParametricTiling.cpp
index cf68ec1b9a7..6dc0bfde371 100644
--- a/mlir/test/lib/Transforms/TestLoopParametricTiling.cpp
+++ b/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;