Using explicit operation types in passes. (#21971)

@kuhar noted that using `auto op = getOperation();` can mask errors if
the Passes.td changes and I've personally had that happen before -
updating all our usage to use the explicit types so that we'll get
compile errors if types don't match (or aren't downcast-compatible).

Author: benvanik

compiler/plugins/input/StableHLO/Conversion/Preprocessing/FlattenTuplesInCFG.cpp

@@ -316,15 +316,15 @@ LogicalResult convertFunction(func::FuncOp oldFunction,
 struct FlattenTuplesInCFG final
     : impl::FlattenTuplesInCFGBase<FlattenTuplesInCFG> {
   void runOnOperation() override {
-    ModuleOp module = getOperation();
-    MLIRContext *ctx = module.getContext();
+    mlir::ModuleOp moduleOp = getOperation();
+    MLIRContext *ctx = moduleOp.getContext();
     Builder builder(ctx);
 
     // Build a list of (oldFunction, newFunction) for all functions we need to
     // replace. This will ensure that when we go to convert function bodies we
     // have only new functions defined.
     SmallVector<std::pair<func::FuncOp, func::FuncOp>> convertedFunctions;
-    for (auto oldFunction : module.getOps<func::FuncOp>()) {
+    for (auto oldFunction : moduleOp.getOps<func::FuncOp>()) {
       FunctionType oldFunctionType = oldFunction.getFunctionType();
 
       llvm::SmallVector<Type> newInputTypes;
@@ -349,7 +349,7 @@ struct FlattenTuplesInCFG final
     // Replace functions in the module.
     for (auto [oldFunction, newFunction] : convertedFunctions) {
       oldFunction.erase();
-      module.push_back(newFunction);
+      moduleOp.push_back(newFunction);
     }
 
     // Run canonicalization patterns to cancel out remaining tuple ops. We need

compiler/plugins/input/StableHLO/Conversion/Preprocessing/FlattenTuplesInSCF.cpp

@@ -243,8 +243,7 @@ struct FlattenTuplesInSCF final
   }
 
   void runOnOperation() override {
-    ModuleOp module = getOperation();
-    MLIRContext *ctx = module.getContext();
+    MLIRContext *ctx = &getContext();
     Builder b(ctx);
 
     // Run canonicalization patterns to cancel out remaining tuple ops. We need

compiler/plugins/input/TOSA/InputConversion/Converti48Toi64.cpp

@@ -166,14 +166,14 @@ void Converti48Toi64Pass::runOnOperation() {
   });
 
   auto *ctx = &getContext();
-  auto func = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
 
   RewritePatternSet patterns(&getContext());
   patterns.add<GenericTypeConvert>(ctx, converter);
   populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                  converter);
 
-  if (failed(applyFullConversion(func, target, std::move(patterns)))) {
+  if (failed(applyFullConversion(funcOp, target, std::move(patterns)))) {
     signalPassFailure();
   }
 }

compiler/plugins/input/TOSA/InputConversion/TosaToLinalgExt.cpp

@@ -157,9 +157,9 @@ class TosaToLinalgExtPass final
     target.addIllegalOp<tosa::ScatterOp>();
     target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
 
-    FunctionOpInterface func = getOperation();
+    mlir::FunctionOpInterface funcOp = getOperation();
     mlir::iree_compiler::populateTosaToLinalgExtPatterns(&patterns);
-    if (failed(applyFullConversion(func, target, std::move(patterns))))
+    if (failed(applyFullConversion(funcOp, target, std::move(patterns))))
       signalPassFailure();
   }
 };

compiler/src/iree/compiler/Codegen/Common/BufferizeCopyOnlyDispatchesPass.cpp

@@ -51,7 +51,7 @@ struct BufferizeCopyOnlyDispatchesPass final
 } // namespace
 
 void BufferizeCopyOnlyDispatchesPass::runOnOperation() {
-  auto funcOp = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
 
   /// Check if the dispatch has all sources for
   /// `iree_tensor_ext.dispatch.tensor.store` operations coming from

compiler/src/iree/compiler/Codegen/Common/CPU/CPUPrepareUkernels.cpp

@@ -411,7 +411,7 @@ struct CPUPrepareUkernelsPass
 void CPUPrepareUkernelsPass::runOnOperation() {
   MLIRContext *ctx = &getContext();
   RewritePatternSet patterns(ctx);
-  auto funcOp = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
   IRRewriter rewriter(ctx);
   auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(funcOp);
 

compiler/src/iree/compiler/Codegen/Common/ConcretizePadResultShape.cpp

@@ -134,7 +134,7 @@ class ConcretizePadResultShapePass final
 public:
   void runOnOperation() override {
     MLIRContext *context = &getContext();
-    auto funcOp = getOperation();
+    mlir::FunctionOpInterface funcOp = getOperation();
 
     ConfigTrackingListener listener;
     GreedyRewriteConfig config;

compiler/src/iree/compiler/Codegen/Common/ConvertToDestinationPassingStylePass.cpp

@@ -599,7 +599,7 @@ struct SwitchStoreOfIfResultValue
 } // namespace
 
 void ConvertToDestinationPassingStylePass::runOnOperation() {
-  auto funcOp = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
   MLIRContext *context = &getContext();
 
   // Dont do anything for functions that have multiple blocks for now.

compiler/src/iree/compiler/Codegen/Common/DropVectorUnitDims.cpp

@@ -31,7 +31,7 @@ class DropVectorUnitDimsPass
 
 void DropVectorUnitDimsPass::runOnOperation() {
   MLIRContext *ctx = &getContext();
-  auto funcOp = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
 
   // Apply transfer ops write to read forwarding and dead transfer write
   // optimizations.

compiler/src/iree/compiler/Codegen/Common/EraseDeadAllocAndStores.cpp

@@ -31,7 +31,7 @@ class EraseDeadAllocAndStoresPass final
 };
 
 void EraseDeadAllocAndStoresPass::runOnOperation() {
-  auto funcOp = getOperation();
+  mlir::FunctionOpInterface funcOp = getOperation();
   IRRewriter rewriter(&getContext());
   memref::eraseDeadAllocAndStores(rewriter, funcOp);
 }