[mlir][test] Generalize a bunch of FuncOp based passes to run on any operation/interfaces

[rriddle]

Summary

A lot of test passes are currently anchored on FuncOp, but this
dependency
is generally just historical. A majority of these test passes can run on
any operation, or can operate on a specific interface
(FunctionOpInterface/SymbolOpInterface).
This allows for greatly reducing the API dependency on FuncOp, which
is slated to be moved out of the Builtin dialect.

mlir/include/mlir/IR/FunctionInterfaces.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_FUNCTIONINTERFACES_H
 #define MLIR_IR_FUNCTIONINTERFACES_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
 
 namespace mlir {
 
 namespace function_interface_impl {
 
 /// Return the name of the attribute used for function types.

mlir/include/mlir/IR/FunctionInterfaces.td

let methods = [
     }]>,
     InterfaceMethod<[{
       Verify the type attribute of the function for derived op-specific
       invariants.
     }],
     "::mlir::LogicalResult", "verifyType">,
   ];
 
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // Name
+    //===------------------------------------------------------------------===//
+
+    /// Return the name of the function.
+    StringRef getName() { return SymbolTable::getSymbolName(*this); }
+  }];
   let extraSharedClassDeclaration = [{
     /// Block list iterator types.
     using BlockListType = Region::BlockListType;
     using iterator = BlockListType::iterator;
     using reverse_iterator = BlockListType::reverse_iterator;
 
     /// Block argument iterator types.
     using BlockArgListType = Region::BlockArgListType;

mlir/test/Analysis/test-dominance.mlir

-// RUN: mlir-opt %s -test-print-dominance -split-input-file 2>&1 | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-print-dominance)" -split-input-file 2>&1 | FileCheck %s
 
 // CHECK-LABEL: Testing : func_condBranch
 func @func_condBranch(%cond : i1) {
   cf.cond_br %cond, ^bb1, ^bb2
 ^bb1:
   cf.br ^exit
 ^bb2:
   cf.br ^exit

mlir/test/Analysis/test-liveness.mlir

-// RUN: mlir-opt %s -test-print-liveness -split-input-file 2>&1 | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-print-liveness)" -split-input-file 2>&1 | FileCheck %s
 
 // CHECK-LABEL: Testing : func_empty
 func @func_empty() {
   // CHECK: Block: 0
   // CHECK-NEXT: LiveIn:{{ *$}}
   // CHECK-NEXT: LiveOut:{{ *$}}
   // CHECK-NEXT: BeginLiveness
   // CHECK-NEXT: EndLiveness

mlir/test/Analysis/test-match-reduction.mlir

-// RUN: mlir-opt %s -test-match-reduction -verify-diagnostics -split-input-file
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-match-reduction)" -verify-diagnostics -split-input-file
 
 // Verify that the generic reduction detection utility works on different
 // dialects.
 
 // expected-remark@below {{Testing function}}
 func @linalg_red_add(%in0t : tensor<?xf32>, %out0t : tensor<1xf32>) {
   // expected-remark@below {{Reduction found in output #0!}}
   // expected-remark@below {{Reduced Value: <block argument> of type 'f32' at index: 0}}

mlir/test/Analysis/test-topoligical-sort.mlir

-// RUN: mlir-opt %s -test-print-topological-sort 2>&1 | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-print-topological-sort)" 2>&1 | FileCheck %s
 
 // CHECK-LABEL: Testing : region
 //       CHECK: arith.addi {{.*}} : index
 //  CHECK-NEXT: scf.for
 //       CHECK: } {__test_sort_original_idx__ = 2 : i64}
 //  CHECK-NEXT: arith.addi {{.*}} : i32
 //  CHECK-NEXT: arith.subi {{.*}} : i32
 func @region(

mlir/test/Dialect/Affine/loop-unswitch.mlir

-// RUN: mlir-opt %s -split-input-file -test-affine-loop-unswitch | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.func(test-affine-loop-unswitch)" | FileCheck %s
 
 // CHECK-DAG: #[[$SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>
 
 // CHECK-LABEL: func @if_else_imperfect
 func @if_else_imperfect(%A : memref<100xi32>, %B : memref<100xi32>, %v : i32) {
 // CHECK: %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref
   affine.for %i = 0 to 100 {
     affine.store %v, %A[%i] : memref<100xi32>

mlir/test/Dialect/Affine/memref-stride-calculation.mlir

-// RUN: mlir-opt %s -test-memref-stride-calculation -o /dev/null | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-memref-stride-calculation)" -o /dev/null | FileCheck %s
 
 func @f(%0: index) {
 // CHECK-LABEL: Testing: f
   %1 = memref.alloc() : memref<3x4x5xf32>
 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
   %2 = memref.alloc(%0) : memref<3x4x?xf32>
 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
   %3 = memref.alloc(%0) : memref<3x?x5xf32>

mlir/test/IR/diagnostic-handler-filter.mlir

-// RUN: mlir-opt %s -test-diagnostic-filter='filters=mysource1' -split-input-file -o - 2>&1 | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="builtin.func(test-diagnostic-filter{filters=mysource1})" -split-input-file -o - 2>&1 | FileCheck %s
 // This test verifies that diagnostic handler can emit the call stack successfully.
 
 // CHECK-LABEL: Test 'test1'
 // CHECK-NEXT: mysource2:1:0: error: test diagnostic
 // CHECK-NEXT: mysource3:2:0: note: called from
 func private @test1() attributes {
   test.loc = loc(callsite("foo"("mysource1":0:0) at callsite("mysource2":1:0 at "mysource3":2:0)))
 }
 
 // -----
 
 // CHECK-LABEL: Test 'test2'
 // CHECK-NEXT: mysource1:0:0: error: test diagnostic
 func private @test2() attributes {
   test.loc = loc("mysource1":0:0)
 }

mlir/test/IR/test-matchers.mlir

-// RUN: mlir-opt %s -mlir-disable-threading=true -test-matchers -o /dev/null 2>&1 | FileCheck %s
+// RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline="builtin.func(test-matchers)" -o /dev/null 2>&1 | FileCheck %s
 
 func @test1(%a: f32, %b: f32, %c: f32) {
   %0 = arith.addf %a, %b: f32
   %1 = arith.addf %a, %c: f32
   %2 = arith.addf %c, %b: f32
   %3 = arith.mulf %a, %2: f32
   %4 = arith.mulf %3, %1: f32
   %5 = arith.mulf %4, %4: f32

mlir/test/Transforms/parametric-mapping.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -test-mapping-to-processing-elements %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -pass-pipeline="builtin.func(test-mapping-to-processing-elements)" %s | FileCheck %s
 
 // CHECK: #[[mul_map:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
 // CHECK: #[[add_map:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
 
 //      CHECK: func @map1d
 // CHECK-SAME: (%[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
 func @map1d(%lb: index, %ub: index, %step: index) {
   // CHECK: %[[threads:.*]]:2 = "new_processor_id_and_range"() : () -> (index, index)

mlir/test/lib/Analysis/TestLiveness.cpp

-//===- TestLiveness.cpp - Test liveness construction and information
-//-------===//
+//===- TestLiveness.cpp - Test liveness construction and information ------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains test passes for constructing and resolving liveness
 // information.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Liveness.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 
 struct TestLivenessPass
-    : public PassWrapper<TestLivenessPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestLivenessPass, InterfacePass<SymbolOpInterface>> {
   StringRef getArgument() const final { return "test-print-liveness"; }
   StringRef getDescription() const final {
     return "Print the contents of a constructed liveness information.";
   }
   void runOnOperation() override {
     llvm::errs() << "Testing : " << getOperation().getName() << "\n";
     getAnalysis<Liveness>().print(llvm::errs());
   }

mlir/test/lib/Analysis/TestMatchReduction.cpp

 //===- TestMatchReduction.cpp - Test the match reduction utility ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains a test pass for the match reduction utility.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 
 void printReductionResult(Operation *redRegionOp, unsigned numOutput,
                           Value reducedValue,
                           ArrayRef<Operation *> combinerOps) {
   if (reducedValue) {
     redRegionOp->emitRemark("Reduction found in output #") << numOutput << "!";
     redRegionOp->emitRemark("Reduced Value: ") << reducedValue;
     for (Operation *combOp : combinerOps)
       redRegionOp->emitRemark("Combiner Op: ") << *combOp;
 
     return;
   }
 
   redRegionOp->emitRemark("Reduction NOT found in output #")
       << numOutput << "!";
 }
 
 struct TestMatchReductionPass
-    : public PassWrapper<TestMatchReductionPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestMatchReductionPass,
+                         InterfacePass<FunctionOpInterface>> {
   StringRef getArgument() const final { return "test-match-reduction"; }
   StringRef getDescription() const final {
     return "Test the match reduction utility.";
   }
 
   void runOnOperation() override {
-    FuncOp func = getOperation();
+    FunctionOpInterface func = getOperation();
     func->emitRemark("Testing function");
 
     func.walk<WalkOrder::PreOrder>([](Operation *op) {
-      if (isa<FuncOp>(op))
+      if (isa<FunctionOpInterface>(op))
         return;
 
       // Limit testing to ops with only one region.
       if (op->getNumRegions() != 1)
         return;
 
       Region &region = op->getRegion(0);
       if (!region.hasOneBlock())

mlir/test/lib/Analysis/TestMemRefBoundCheck.cpp

 #define DEBUG_TYPE "memref-bound-check"
 
 using namespace mlir;
 
 namespace {
 
 /// Checks for out of bound memref access subscripts..
 struct TestMemRefBoundCheck
-    : public PassWrapper<TestMemRefBoundCheck, OperationPass<FuncOp>> {
+    : public PassWrapper<TestMemRefBoundCheck, OperationPass<>> {
   StringRef getArgument() const final { return "test-memref-bound-check"; }
   StringRef getDescription() const final {
-    return "Check memref access bounds in a Function";
+    return "Check memref access bounds";
   }
   void runOnOperation() override;
 };
 
 } // namespace
 
 void TestMemRefBoundCheck::runOnOperation() {
-  getOperation().walk([](Operation *opInst) {
+  getOperation()->walk([](Operation *opInst) {
     TypeSwitch<Operation *>(opInst)
         .Case<AffineReadOpInterface, AffineWriteOpInterface>(
             [](auto op) { (void)boundCheckLoadOrStoreOp(op); });
 
     // TODO: do this for DMA ops as well.
   });
 }
 
 namespace mlir {
 namespace test {
 void registerMemRefBoundCheck() { PassRegistration<TestMemRefBoundCheck>(); }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Analysis/TestMemRefDependenceCheck.cpp

 
 using namespace mlir;
 
 namespace {
 
 // TODO: Add common surrounding loop depth-wise dependence checks.
 /// Checks dependences between all pairs of memref accesses in a Function.
 struct TestMemRefDependenceCheck
-    : public PassWrapper<TestMemRefDependenceCheck, OperationPass<FuncOp>> {
+    : public PassWrapper<TestMemRefDependenceCheck, OperationPass<>> {
   StringRef getArgument() const final { return "test-memref-dependence-check"; }
   StringRef getDescription() const final {
     return "Checks dependences between all pairs of memref accesses.";
   }
   SmallVector<Operation *, 4> loadsAndStores;
   void runOnOperation() override;
 };
 
for (unsigned j = 0; j < e; ++j) {
             << i << " to " << j << " at depth " << d << " = "
             << getDirectionVectorStr(ret, numCommonLoops, d,
                                      dependenceComponents);
       }
     }
   }
 }
 
-// Walks the Function 'f' adding load and store ops to 'loadsAndStores'.
-// Runs pair-wise dependence checks.
+/// Walks the operation adding load and store ops to 'loadsAndStores'. Runs
+/// pair-wise dependence checks.
 void TestMemRefDependenceCheck::runOnOperation() {
   // Collect the loads and stores within the function.
   loadsAndStores.clear();
-  getOperation().walk([&](Operation *op) {
+  getOperation()->walk([&](Operation *op) {
     if (isa<AffineLoadOp, AffineStoreOp>(op))
       loadsAndStores.push_back(op);
   });
 
   checkDependences(loadsAndStores);
 }
 
 namespace mlir {
 namespace test {
 void registerTestMemRefDependenceCheck() {
   PassRegistration<TestMemRefDependenceCheck>();
 }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Analysis/TestMemRefStrideCalculation.cpp

 //===- TestMemRefStrideCalculation.cpp - Pass to test strides computation--===//
 //
 // 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/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 struct TestMemRefStrideCalculation
-    : public PassWrapper<TestMemRefStrideCalculation, OperationPass<FuncOp>> {
+    : public PassWrapper<TestMemRefStrideCalculation,
+                         InterfacePass<SymbolOpInterface>> {
   StringRef getArgument() const final {
     return "test-memref-stride-calculation";
   }
   StringRef getDescription() const final {
     return "Test operation constant folding";
   }
   void runOnOperation() override;
 };

mlir/test/lib/Analysis/TestSlice.cpp

 
 using namespace mlir;
 
 static const StringLiteral kOrderMarker = "__test_sort_original_idx__";
 
 namespace {
 
 struct TestTopologicalSortPass
-    : public PassWrapper<TestTopologicalSortPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestTopologicalSortPass,
+                         InterfacePass<SymbolOpInterface>> {
   StringRef getArgument() const final { return "test-print-topological-sort"; }
   StringRef getDescription() const final {
     return "Print operations in topological order";
   }
   void runOnOperation() override {
     std::map<int, Operation *> ops;
     getOperation().walk([&ops](Operation *op) {
       if (auto originalOrderAttr = op->getAttrOfType<IntegerAttr>(kOrderMarker))

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

 //===- TestAffineDataCopy.cpp - Test affine data copy utility -------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to test affine data copy utility functions and
 // options.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 #define PASS_NAME "test-affine-data-copy"
 
 using namespace mlir;

mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp

 #define PASS_NAME "test-affine-loop-unswitch"
 
 using namespace mlir;
 
 namespace {
 
 /// This pass applies the permutation on the first maximal perfect nest.
 struct TestAffineLoopUnswitching
-    : public PassWrapper<TestAffineLoopUnswitching, OperationPass<FuncOp>> {
+    : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> {
   StringRef getArgument() const final { return PASS_NAME; }
   StringRef getDescription() const final {
     return "Tests affine loop unswitching / if/else hoisting";
   }
   TestAffineLoopUnswitching() = default;
   TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) = default;
 
   void runOnOperation() override;
 
   /// The maximum number of iterations to run this for.
   constexpr static unsigned kMaxIterations = 5;
 };
 
 } // namespace
 
 void TestAffineLoopUnswitching::runOnOperation() {
   // Each hoisting invalidates a lot of IR around. Just stop the walk after the
   // first if/else hoisting, and repeat until no more hoisting can be done, or
   // the maximum number of iterations have been run.
-  auto func = getOperation();
+  Operation *op = getOperation();
   unsigned i = 0;
   do {
     auto walkFn = [](AffineIfOp op) {
       return succeeded(hoistAffineIfOp(op)) ? WalkResult::interrupt()
                                             : WalkResult::advance();
     };
-    if (func.walk(walkFn).wasInterrupted())
+    if (op->walk(walkFn).wasInterrupted())
       break;
   } while (++i < kMaxIterations);
 }
 
 namespace mlir {
 void registerTestAffineLoopUnswitchingPass() {
   PassRegistration<TestAffineLoopUnswitching>();
 }
 } // namespace mlir

mlir/test/lib/Dialect/Affine/TestLoopMapping.cpp

 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 namespace {
 class TestLoopMappingPass
-    : public PassWrapper<TestLoopMappingPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestLoopMappingPass, OperationPass<>> {
 public:
   StringRef getArgument() const final {
     return "test-mapping-to-processing-elements";
   }
   StringRef getDescription() const final {
     return "test mapping a single loop on a virtual processor grid";
   }
   explicit TestLoopMappingPass() = default;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AffineDialect, scf::SCFDialect>();
   }
 
   void runOnOperation() override {
-    FuncOp func = getOperation();
-
     // SSA values for the transformation are created out of thin air by
     // unregistered "new_processor_id_and_range" operations. This is enough to
     // emulate mapping conditions.
     SmallVector<Value, 8> processorIds, numProcessors;
-    func.walk([&processorIds, &numProcessors](Operation *op) {
+    getOperation()->walk([&processorIds, &numProcessors](Operation *op) {
       if (op->getName().getStringRef() != "new_processor_id_and_range")
         return;
       processorIds.push_back(op->getResult(0));
       numProcessors.push_back(op->getResult(1));
     });
 
-    func.walk([&processorIds, &numProcessors](scf::ForOp op) {
+    getOperation()->walk([&processorIds, &numProcessors](scf::ForOp op) {
       // Ignore nested loops.
       if (op->getParentRegion()->getParentOfType<scf::ForOp>())
         return;
       mapLoopToProcessorIds(op, processorIds, numProcessors);
     });
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLoopMappingPass() { PassRegistration<TestLoopMappingPass>(); }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Pass/Pass.h"
 
 #define PASS_NAME "test-loop-permutation"
 
 using namespace mlir;
 
-static llvm::cl::OptionCategory clOptionsCategory(PASS_NAME " options");
-
 namespace {
 
 /// This pass applies the permutation on the first maximal perfect nest.
 struct TestLoopPermutation
-    : public PassWrapper<TestLoopPermutation, OperationPass<FuncOp>> {
+    : public PassWrapper<TestLoopPermutation, OperationPass<>> {
   StringRef getArgument() const final { return PASS_NAME; }
   StringRef getDescription() const final {
     return "Tests affine loop permutation utility";
   }
   TestLoopPermutation() = default;
   TestLoopPermutation(const TestLoopPermutation &pass) : PassWrapper(pass){};
 
   void runOnOperation() override;
 
 private:
   /// Permutation specifying loop i is mapped to permList[i] in
   /// transformed nest (with i going from outermost to innermost).
   ListOption<unsigned> permList{*this, "permutation-map",
                                 llvm::cl::desc("Specify the loop permutation"),
                                 llvm::cl::OneOrMore, llvm::cl::CommaSeparated};
 };
 
 } // namespace
 
 void TestLoopPermutation::runOnOperation() {
 
   SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
 
   SmallVector<AffineForOp, 2> forOps;
-  getOperation().walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
+  getOperation()->walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
 
   for (auto forOp : forOps) {
     SmallVector<AffineForOp, 6> nest;
     // Get the maximal perfect nest.
     getPerfectlyNestedLoops(nest, forOp);
     // Permute if the nest's size is consistent with the specified
     // permutation.
     if (nest.size() >= 2 && nest.size() == permMap.size()) {

mlir/test/lib/Dialect/GPU/TestGpuParallelLoopMapping.cpp

 
 using namespace mlir;
 
 namespace {
 /// Simple pass for testing the mapping of parallel loops to hardware ids using
 /// a greedy mapping strategy.
 class TestGpuGreedyParallelLoopMappingPass
     : public PassWrapper<TestGpuGreedyParallelLoopMappingPass,
-                         OperationPass<FuncOp>> {
+                         OperationPass<>> {
   StringRef getArgument() const final {
     return "test-gpu-greedy-parallel-loop-mapping";
   }
   StringRef getDescription() const final {
     return "Greedily maps all parallel loops to gpu hardware ids.";
   }
   void runOnOperation() override {
-    Operation *op = getOperation();
-    for (Region &region : op->getRegions())
+    for (Region &region : getOperation()->getRegions())
       greedilyMapParallelSCFToGPU(region);
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestGpuParallelLoopMappingPass() {
   PassRegistration<TestGpuGreedyParallelLoopMappingPass>();
 }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
 
 namespace {
 struct TestPadFusionPass
-    : public PassWrapper<TestPadFusionPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestPadFusionPass, OperationPass<>> {
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<AffineDialect, linalg::LinalgDialect, tensor::TensorDialect>();
   }
 
   StringRef getArgument() const final { return "test-linalg-pad-fusion"; }
   StringRef getDescription() const final { return "Test PadOp fusion"; }
 
   void runOnOperation() override {
     MLIRContext *context = &getContext();
-    FuncOp funcOp = getOperation();
     RewritePatternSet patterns(context);
     linalg::populateFuseTensorPadWithProducerLinalgOpPatterns(patterns);
-    if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
-                                            std::move(patterns))))
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       return signalPassFailure();
   }
 };
 } // namespace
 
 namespace test {
 void registerTestPadFusion() { PassRegistration<TestPadFusionPass>(); }
 } // namespace test
 
 } // namespace mlir

mlir/test/lib/Dialect/Math/TestAlgebraicSimplification.cpp

 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestMathAlgebraicSimplificationPass
-    : public PassWrapper<TestMathAlgebraicSimplificationPass,
-                         OperationPass<FuncOp>> {
+    : public PassWrapper<TestMathAlgebraicSimplificationPass, OperationPass<>> {
   void runOnOperation() override;
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect, math::MathDialect>();
   }
   StringRef getArgument() const final {
     return "test-math-algebraic-simplification";
   }
   StringRef getDescription() const final {

mlir/test/lib/Dialect/Math/TestExpandTanh.cpp

 #include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestExpandTanhPass
-    : public PassWrapper<TestExpandTanhPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestExpandTanhPass, OperationPass<>> {
   void runOnOperation() override;
   StringRef getArgument() const final { return "test-expand-tanh"; }
   StringRef getDescription() const final { return "Test expanding tanh"; }
 };
 } // namespace
 
 void TestExpandTanhPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());

mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp

 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestMathPolynomialApproximationPass
-    : public PassWrapper<TestMathPolynomialApproximationPass,
-                         OperationPass<FuncOp>> {
+    : public PassWrapper<TestMathPolynomialApproximationPass, OperationPass<>> {
   TestMathPolynomialApproximationPass() = default;
   TestMathPolynomialApproximationPass(
       const TestMathPolynomialApproximationPass &pass)
       : PassWrapper(pass) {}
 
   void runOnOperation() override;
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithmeticDialect, math::MathDialect,

mlir/test/lib/Dialect/MemRef/TestComposeSubView.cpp

 #include "mlir/Dialect/MemRef/Transforms/ComposeSubView.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestComposeSubViewPass
-    : public PassWrapper<TestComposeSubViewPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestComposeSubViewPass, OperationPass<>> {
   StringRef getArgument() const final { return "test-compose-subview"; }
   StringRef getDescription() const final {
     return "Test combining composed subviews";
   }
   void runOnOperation() override;
   void getDependentDialects(DialectRegistry &registry) const override;
 };
 

mlir/test/lib/Dialect/MemRef/TestMultiBuffer.cpp

 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 struct TestMultiBufferingPass
-    : public PassWrapper<TestMultiBufferingPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestMultiBufferingPass, OperationPass<>> {
   TestMultiBufferingPass() = default;
   TestMultiBufferingPass(const TestMultiBufferingPass &pass)
       : PassWrapper(pass) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AffineDialect>();
   }
   StringRef getArgument() const final { return "test-multi-buffering"; }
   StringRef getDescription() const final {
     return "Test multi buffering transformation";
   }
   void runOnOperation() override;
   Option<unsigned> multiplier{
       *this, "multiplier",
       llvm::cl::desc(
           "Decide how many versions of the buffer should be created,"),
       llvm::cl::init(2)};
 };
 
 void TestMultiBufferingPass::runOnOperation() {
   SmallVector<memref::AllocOp> allocs;
-  getOperation().walk(
+  getOperation()->walk(
       [&allocs](memref::AllocOp alloc) { allocs.push_back(alloc); });
   for (memref::AllocOp alloc : allocs)
     (void)multiBuffer(alloc, multiplier);
 }
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestMultiBuffering() {
   PassRegistration<TestMultiBufferingPass>();
 }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Dialect/SCF/TestLoopParametricTiling.cpp

 
 using namespace mlir;
 
 namespace {
 
 // Extracts fixed-range loops for top-level loop nests with ranges defined in
 // the pass constructor.  Assumes loops are permutable.
 class SimpleParametricLoopTilingPass
-    : public PassWrapper<SimpleParametricLoopTilingPass,
-                         OperationPass<FuncOp>> {
+    : public PassWrapper<SimpleParametricLoopTilingPass, OperationPass<>> {
 public:
   StringRef getArgument() const final {
     return "test-extract-fixed-outer-loops";
   }
   StringRef getDescription() const final {
     return "test application of parametric tiling to the outer loops so that "
            "the ranges of outer loops become static";
   }
   SimpleParametricLoopTilingPass() = default;
   SimpleParametricLoopTilingPass(const SimpleParametricLoopTilingPass &) {}
   explicit SimpleParametricLoopTilingPass(ArrayRef<int64_t> outerLoopSizes) {
     sizes = outerLoopSizes;
   }
 
   void runOnOperation() override {
-    FuncOp func = getOperation();
-    func.walk([this](scf::ForOp op) {
+    getOperation()->walk([this](scf::ForOp op) {
       // Ignore nested loops.
       if (op->getParentRegion()->getParentOfType<scf::ForOp>())
         return;
       extractFixedOuterLoops(op, sizes);
     });
   }
 
   ListOption<int64_t> sizes{

mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp

static unsigned getNestingDepth(Operation *op) {
   while ((currOp = currOp->getParentOp())) {
     if (isa<scf::ForOp>(currOp))
       depth++;
   }
   return depth;
 }
 
 class TestLoopUnrollingPass
-    : public PassWrapper<TestLoopUnrollingPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestLoopUnrollingPass, OperationPass<>> {
 public:
   StringRef getArgument() const final { return "test-loop-unrolling"; }
   StringRef getDescription() const final {
     return "Tests loop unrolling transformation";
   }
   TestLoopUnrollingPass() = default;
   TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
   explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
                                  unsigned loopDepthParam,
                                  bool annotateLoopParam) {
     unrollFactor = unrollFactorParam;
     loopDepth = loopDepthParam;
     annotateLoop = annotateLoopParam;
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithmeticDialect>();
   }
 
   void runOnOperation() override {
-    FuncOp func = getOperation();
     SmallVector<scf::ForOp, 4> loops;
-    func.walk([&](scf::ForOp forOp) {
+    getOperation()->walk([&](scf::ForOp forOp) {
       if (getNestingDepth(forOp) == loopDepth)
         loops.push_back(forOp);
     });
     auto annotateFn = [this](unsigned i, Operation *op, OpBuilder b) {
       if (annotateLoop) {
         op->setAttr("unrolled_iteration", b.getUI32IntegerAttr(i));
       }
     };

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestTensorTransforms
-    : public PassWrapper<TestTensorTransforms, OperationPass<FuncOp>> {
+    : public PassWrapper<TestTensorTransforms, OperationPass<>> {
   TestTensorTransforms() = default;
   TestTensorTransforms(const TestTensorTransforms &pass) : PassWrapper(pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithmeticDialect, scf::SCFDialect>();
   }
 
   StringRef getArgument() const final {
struct TestTensorTransforms
 
   Option<bool> testFoldConstantExtractSlice{
       *this, "test-fold-constant-extract-slice",
       llvm::cl::desc("Test folding arith.constant and tensor.extract_slice"),
       llvm::cl::init(false)};
 };
 } // namespace
 
-static void applySplitPaddingPatterns(FuncOp funcOp) {
-  RewritePatternSet patterns(funcOp.getContext());
+static void applySplitPaddingPatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
   tensor::populateSplitPaddingPatterns(patterns);
-  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
-static void applyFoldConstantExtractSlicePatterns(FuncOp funcOp) {
-  RewritePatternSet patterns(funcOp.getContext());
+static void applyFoldConstantExtractSlicePatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
   tensor::ControlConstantExtractSliceFusionFn controlFn =
       [](tensor::ExtractSliceOp op) {
         if (!op.source().hasOneUse())
           return false;
 
         auto resultType = op.result().getType().cast<ShapedType>();
         constexpr int64_t kConstantFoldingMaxNumElements = 1024;
         return resultType.getNumElements() <= kConstantFoldingMaxNumElements;
       };
 
   tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
-  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 void TestTensorTransforms::runOnOperation() {
-  FuncOp func = getOperation();
+  Operation *rootOp = getOperation();
   if (testSplitPaddingPatterns)
-    applySplitPaddingPatterns(func);
+    applySplitPaddingPatterns(rootOp);
   if (testFoldConstantExtractSlice)
-    applyFoldConstantExtractSlicePatterns(func);
+    applyFoldConstantExtractSlicePatterns(rootOp);
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir

mlir/test/lib/IR/TestDiagnostics.cpp

 
 #include "mlir/Pass/Pass.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 namespace {
 struct TestDiagnosticFilterPass
-    : public PassWrapper<TestDiagnosticFilterPass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestDiagnosticFilterPass,
+                         InterfacePass<SymbolOpInterface>> {
   StringRef getArgument() const final { return "test-diagnostic-filter"; }
   StringRef getDescription() const final {
     return "Test diagnostic filtering support.";
   }
   TestDiagnosticFilterPass() = default;
   TestDiagnosticFilterPass(const TestDiagnosticFilterPass &) {}
 
   void runOnOperation() override {

mlir/test/lib/IR/TestDominance.cpp

public:
   }
 
 private:
   Operation *operation;
   DenseMap<Block *, size_t> blockIds;
 };
 
 struct TestDominancePass
-    : public PassWrapper<TestDominancePass, OperationPass<FuncOp>> {
+    : public PassWrapper<TestDominancePass, InterfacePass<SymbolOpInterface>> {
   StringRef getArgument() const final { return "test-print-dominance"; }
   StringRef getDescription() const final {
     return "Print the dominance information for multiple regions.";
   }
 
   void runOnOperation() override {
     llvm::errs() << "Testing : " << getOperation().getName() << "\n";
     DominanceTest dominanceTest(getOperation());

mlir/test/lib/IR/TestFunc.cpp

 //===- TestFunc.cpp - Pass to test helpers on function 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/BuiltinOps.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
-/// This is a test pass for verifying FuncOp's insertArgument method.
+/// This is a test pass for verifying FunctionOpInterface's insertArgument
+/// method.
 struct TestFuncInsertArg
     : public PassWrapper<TestFuncInsertArg, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-insert-arg"; }
   StringRef getDescription() const final { return "Test inserting func args."; }
   void runOnOperation() override {
     auto module = getOperation();
 
     UnknownLoc unknownLoc = UnknownLoc::get(module.getContext());
-    for (FuncOp func : module.getOps<FuncOp>()) {
+    for (auto func : module.getOps<FunctionOpInterface>()) {
       auto inserts = func->getAttrOfType<ArrayAttr>("test.insert_args");
       if (!inserts || inserts.empty())
         continue;
       SmallVector<unsigned, 4> indicesToInsert;
       SmallVector<Type, 4> typesToInsert;
       SmallVector<DictionaryAttr, 4> attrsToInsert;
       SmallVector<Location, 4> locsToInsert;
       for (auto insert : inserts.getAsRange<ArrayAttr>()) {
for (auto func : module.getOps<FunctionOpInterface>()) {
       }
       func->removeAttr("test.insert_args");
       func.insertArguments(indicesToInsert, typesToInsert, attrsToInsert,
                            locsToInsert);
     }
   }
 };
 
-/// This is a test pass for verifying FuncOp's insertResult method.
+/// This is a test pass for verifying FunctionOpInterface's insertResult method.
 struct TestFuncInsertResult
     : public PassWrapper<TestFuncInsertResult, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-insert-result"; }
   StringRef getDescription() const final {
     return "Test inserting func results.";
   }
   void runOnOperation() override {
     auto module = getOperation();
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
+    for (auto func : module.getOps<FunctionOpInterface>()) {
       auto inserts = func->getAttrOfType<ArrayAttr>("test.insert_results");
       if (!inserts || inserts.empty())
         continue;
       SmallVector<unsigned, 4> indicesToInsert;
       SmallVector<Type, 4> typesToInsert;
       SmallVector<DictionaryAttr, 4> attrsToInsert;
       for (auto insert : inserts.getAsRange<ArrayAttr>()) {
         indicesToInsert.push_back(
             insert[0].cast<IntegerAttr>().getValue().getZExtValue());
         typesToInsert.push_back(insert[1].cast<TypeAttr>().getValue());
         attrsToInsert.push_back(insert.size() > 2
                                     ? insert[2].cast<DictionaryAttr>()
                                     : DictionaryAttr::get(&getContext()));
       }
       func->removeAttr("test.insert_results");
       func.insertResults(indicesToInsert, typesToInsert, attrsToInsert);
     }
   }
 };
 
-/// This is a test pass for verifying FuncOp's eraseArgument method.
+/// This is a test pass for verifying FunctionOpInterface's eraseArgument
+/// method.
 struct TestFuncEraseArg
     : public PassWrapper<TestFuncEraseArg, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-erase-arg"; }
   StringRef getDescription() const final { return "Test erasing func args."; }
   void runOnOperation() override {
     auto module = getOperation();
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
+    for (auto func : module.getOps<FunctionOpInterface>()) {
       BitVector indicesToErase(func.getNumArguments());
       for (auto argIndex : llvm::seq<int>(0, func.getNumArguments()))
         if (func.getArgAttr(argIndex, "test.erase_this_arg"))
           indicesToErase.set(argIndex);
       func.eraseArguments(indicesToErase);
     }
   }
 };
 
-/// This is a test pass for verifying FuncOp's eraseResult method.
+/// This is a test pass for verifying FunctionOpInterface's eraseResult method.
 struct TestFuncEraseResult
     : public PassWrapper<TestFuncEraseResult, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-erase-result"; }
   StringRef getDescription() const final {
     return "Test erasing func results.";
   }
   void runOnOperation() override {
     auto module = getOperation();
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
+    for (auto func : module.getOps<FunctionOpInterface>()) {
       BitVector indicesToErase(func.getNumResults());
       for (auto resultIndex : llvm::seq<int>(0, func.getNumResults()))
         if (func.getResultAttr(resultIndex, "test.erase_this_result"))
           indicesToErase.set(resultIndex);
       func.eraseResults(indicesToErase);
     }
   }
 };
 
-/// This is a test pass for verifying FuncOp's setType method.
+/// This is a test pass for verifying FunctionOpInterface's setType method.
 struct TestFuncSetType
     : public PassWrapper<TestFuncSetType, OperationPass<ModuleOp>> {
   StringRef getArgument() const final { return "test-func-set-type"; }
-  StringRef getDescription() const final { return "Test FuncOp::setType."; }
+  StringRef getDescription() const final {
+    return "Test FunctionOpInterface::setType.";
+  }
   void runOnOperation() override {
     auto module = getOperation();
     SymbolTable symbolTable(module);
 
-    for (FuncOp func : module.getOps<FuncOp>()) {
+    for (auto func : module.getOps<FunctionOpInterface>()) {
       auto sym = func->getAttrOfType<FlatSymbolRefAttr>("test.set_type_from");
       if (!sym)
         continue;
-      func.setType(symbolTable.lookup<FuncOp>(sym.getValue()).getType());
+      func.setType(
+          symbolTable.lookup<FunctionOpInterface>(sym.getValue()).getType());
     }
   }
 };
 } // namespace
 
 namespace mlir {
 void registerTestFunc() {
   PassRegistration<TestFuncInsertArg>();

mlir/test/lib/IR/TestMatchers.cpp

 //===- TestMatchers.cpp - Pass to test matchers ---------------------------===//
 //
 // 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/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 /// This is a test pass for verifying matchers.
-struct TestMatchers : public PassWrapper<TestMatchers, OperationPass<FuncOp>> {
+struct TestMatchers
+    : public PassWrapper<TestMatchers, InterfacePass<FunctionOpInterface>> {
   void runOnOperation() override;
   StringRef getArgument() const final { return "test-matchers"; }
   StringRef getDescription() const final {
     return "Test C++ pattern matchers.";
   }
 };
 } // namespace
 
 // This could be done better but is not worth the variadic template trouble.
 template <typename Matcher>
-static unsigned countMatches(FuncOp f, Matcher &matcher) {
+static unsigned countMatches(FunctionOpInterface f, Matcher &matcher) {
   unsigned count = 0;
   f.walk([&count, &matcher](Operation *op) {
     if (matcher.match(op))
       ++count;
   });
   return count;
 }
 
 using mlir::matchers::m_Any;
 using mlir::matchers::m_Val;
-static void test1(FuncOp f) {
+static void test1(FunctionOpInterface f) {
   assert(f.getNumArguments() == 3 && "matcher test funcs must have 3 args");
 
   auto a = m_Val(f.getArgument(0));
   auto b = m_Val(f.getArgument(1));
   auto c = m_Val(f.getArgument(2));
 
   auto p0 = m_Op<arith::AddFOp>(); // using 0-arity matcher
   llvm::outs() << "Pattern add(*) matched " << countMatches(f, p0)
static void test1(FunctionOpInterface f) {
   llvm::outs() << "Pattern mul(mul(a, *), add(a, c)) matched "
                << countMatches(f, p16) << " times\n";
 
   auto p17 = m_Op<arith::MulFOp>(mulOfAany, m_Op<arith::AddFOp>(c, b));
   llvm::outs() << "Pattern mul(mul(a, *), add(c, b)) matched "
                << countMatches(f, p17) << " times\n";
 }
 
-void test2(FuncOp f) {
+void test2(FunctionOpInterface f) {
   auto a = m_Val(f.getArgument(0));
   FloatAttr floatAttr;
   auto p =
       m_Op<arith::MulFOp>(a, m_Op<arith::AddFOp>(a, m_Constant(&floatAttr)));
   auto p1 = m_Op<arith::MulFOp>(a, m_Op<arith::AddFOp>(a, m_Constant()));
   // Last operation that is not the terminator.
   Operation *lastOp = f.getBody().front().back().getPrevNode();
   if (p.match(lastOp))

mlir/test/lib/IR/TestOpaqueLoc.cpp

getOperation().getBody()->walk([&](Operation *op) {
 
       Location loc = op->getLoc();
 
       /// Set opaque location without fallback location to test the
       /// corresponding get method.
       op->setLoc(
           OpaqueLoc::get<MyLocation *>(myLocs.back().get(), &getContext()));
 
-      if (isa<FuncOp>(op) || op->hasTrait<OpTrait::IsTerminator>())
+      if (isa<ModuleOp>(op->getParentOp()) ||
+          op->hasTrait<OpTrait::IsTerminator>())
         return;
 
       OpBuilder builder(op);
 
       /// Add the same operation but with fallback location to test the
       /// corresponding get method and serialization.
       Operation *opCloned1 = builder.clone(*op);
       opCloned1->setLoc(OpaqueLoc::get<MyLocation *>(myLocs.back().get(), loc));

mlir/test/lib/IR/TestVisitors.cpp

static void testPureCallbacks(Operation *op) {
   llvm::outs() << "Region post-order visits"
                << "\n";
   op->walk<WalkOrder::PostOrder>(regionPure);
 }
 
 /// Tests erasure callbacks that skip the walk.
 static void testSkipErasureCallbacks(Operation *op) {
   auto skipOpErasure = [](Operation *op) {
-    // Do not erase module and function op. Otherwise there wouldn't be too
-    // much to test in pre-order.
-    if (isa<ModuleOp>(op) || isa<FuncOp>(op))
+    // Do not erase module and module children operations. Otherwise, there
+    // wouldn't be too much to test in pre-order.
+    if (isa<ModuleOp>(op) || isa<ModuleOp>(op->getParentOp()))
       return WalkResult::advance();
 
     llvm::outs() << "Erasing ";
     printOperation(op);
     llvm::outs() << "\n";
     op->dropAllUses();
     op->erase();
     return WalkResult::skip();
   };
   auto skipBlockErasure = [](Block *block) {
-    // Do not erase module and function blocks. Otherwise there wouldn't be
-    // too much to test in pre-order.
+    // Do not erase module and module children blocks. Otherwise there wouldn't
+    // be too much to test in pre-order.
     Operation *parentOp = block->getParentOp();
-    if (isa<ModuleOp>(parentOp) || isa<FuncOp>(parentOp))
+    if (isa<ModuleOp>(parentOp) || isa<ModuleOp>(parentOp->getParentOp()))
       return WalkResult::advance();
 
     llvm::outs() << "Erasing ";
     printBlock(block);
     llvm::outs() << "\n";
     block->erase();
     return WalkResult::skip();
   };

mlir/test/lib/Transforms/TestConstantFold.cpp

 //===- TestConstantFold.cpp - Pass to test constant folding ---------------===//
 //
 // 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/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple constant folding pass.
 struct TestConstantFold
-    : public PassWrapper<TestConstantFold, OperationPass<FuncOp>> {
+    : public PassWrapper<TestConstantFold, OperationPass<>> {
   StringRef getArgument() const final { return "test-constant-fold"; }
   StringRef getDescription() const final {
     return "Test operation constant folding";
   }
-  // All constants in the function post folding.
-  SmallVector<Operation *, 8> existingConstants;
+  // All constants in the operation post folding.
+  SmallVector<Operation *> existingConstants;
 
   void foldOperation(Operation *op, OperationFolder &helper);
   void runOnOperation() override;
 };
 } // namespace
 
 void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) {
   auto processGeneratedConstants = [this](Operation *op) {
     existingConstants.push_back(op);
   };
 
   // Attempt to fold the specified operation, including handling unused or
   // duplicated constants.
   (void)helper.tryToFold(op, processGeneratedConstants);
 }
 
-// For now, we do a simple top-down pass over a function folding constants.  We
-// don't handle conditional control flow, block arguments, folding conditional
-// branches, or anything else fancy.
 void TestConstantFold::runOnOperation() {
   existingConstants.clear();
 
-  // Collect and fold the operations within the function.
+  // Collect and fold the operations within the operation.
   SmallVector<Operation *, 8> ops;
-  getOperation().walk([&](Operation *op) { ops.push_back(op); });
+  getOperation()->walk([&](Operation *op) { ops.push_back(op); });
 
   // Fold the constants in reverse so that the last generated constants from
   // folding are at the beginning. This creates somewhat of a linear ordering to
   // the newly generated constants that matches the operation order and improves
   // the readability of test cases.
   OperationFolder helper(&getContext());
   for (Operation *op : llvm::reverse(ops))
     foldOperation(op, helper);
 
   // By the time we are done, we may have simplified a bunch of code, leaving
-  // around dead constants.  Check for them now and remove them.
+  // around dead constants. Check for them now and remove them.
   for (auto *cst : existingConstants) {
     if (cst->use_empty())
       cst->erase();
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestConstantFold() { PassRegistration<TestConstantFold>(); }
 } // namespace test
 } // namespace mlir