Make the CodeGenRewrite pass able to operate at any level of the IR (NFC)

flang/include/flang/Optimizer/CodeGen/CGPasses.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_OPTIMIZER_CODEGEN_FIR_PASSES
 #define FORTRAN_OPTIMIZER_CODEGEN_FIR_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def CodeGenRewrite : Pass<"cg-rewrite", "mlir::ModuleOp"> {
+def CodeGenRewrite : Pass<"cg-rewrite"> {
   let summary = "Rewrite some FIR ops into their code-gen forms.";
   let description = [{
     Fuse specific subgraphs into single Ops for code generation.
   }];
   let constructor = "fir::createFirCodeGenRewritePass()";
   let dependentDialects = [
     "fir::FIROpsDialect", "fir::FIRCodeGenDialect", "mlir::BuiltinDialect",
     "mlir::LLVM::LLVMDialect", "mlir::omp::OpenMPDialect"
   ];
   let statistics = [
     Statistic<"numDCE", "num-dce'd", "Number of operations eliminated">
   ];
 }
 
 #endif // FORTRAN_OPTIMIZER_CODEGEN_FIR_PASSES

flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp

LLVM_DEBUG(llvm::dbgs()
                << "rewriting " << arrCoor << " to " << xArrCoor << '\n');
     rewriter.replaceOp(arrCoor, xArrCoor.getOperation()->getResults());
     return mlir::success();
   }
 };
 
 class CodeGenRewrite : public CodeGenRewriteBase<CodeGenRewrite> {
 public:
-  void runOn(mlir::Operation *op, mlir::Region &region) {
+  void runOnOperation() override final {
     auto &context = getContext();
     mlir::OpBuilder rewriter(&context);
     mlir::ConversionTarget target(context);
     target.addLegalDialect<FIROpsDialect, FIRCodeGenDialect,
                            mlir::StandardOpsDialect>();
     target.addIllegalOp<ArrayCoorOp>();
     target.addIllegalOp<ReboxOp>();
     target.addDynamicallyLegalOp<EmboxOp>([](EmboxOp embox) {
       return !(embox.getShape() ||
                embox.getType().cast<BoxType>().getEleTy().isa<SequenceType>());
     });
-    mlir::OwningRewritePatternList patterns;
+    mlir::OwningRewritePatternList patterns(&context);
     patterns.insert<EmboxConversion, ArrayCoorConversion, ReboxConversion>(
         &context);
-    if (mlir::failed(
-            mlir::applyPartialConversion(op, target, std::move(patterns)))) {
+    if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
+                                                  std::move(patterns)))) {
       mlir::emitError(mlir::UnknownLoc::get(&context),
                       "error in running the pre-codegen conversions");
       signalPassFailure();
     }
     // Erase any residual.
+    for (Region &region : getOperation()->getRegions())
       simplifyRegion(region);
   }
 
-  void runOnOperation() override final {
-    // Call runOn on all top level regions that may contain emboxOp/arrayCoorOp.
-    auto mod = getOperation();
-    for (auto func : mod.getOps<mlir::FuncOp>())
-      runOn(func, func.getBody());
-    for (auto global : mod.getOps<fir::GlobalOp>())
-      runOn(global, global.getRegion());
-  }
-
   // Clean up the region.
   void simplifyRegion(mlir::Region &region) {
     for (auto &block : region.getBlocks())
       for (auto &op : block.getOperations()) {
         for (auto &reg : op.getRegions())
           simplifyRegion(reg);
         maybeEraseOp(&op);
       }

flang/test/Fir/cg-ops.fir

-// RUN: fir-opt --cg-rewrite %s | FileCheck %s
+// RUN: fir-opt --pass-pipeline="func(cg-rewrite),fir.global(cg-rewrite)" %s | FileCheck %s
 
 // CHECK-LABEL: func @codegen(
 // CHECK-SAME: %[[arg:.*]]: !fir
 func @codegen(%addr : !fir.ref<!fir.array<?xi32>>) {
   // CHECK: %[[zero:.*]] = constant 0 : index
   %0 = constant 0 : index
   %1 = fir.shape_shift %0, %0 : (index, index) -> !fir.shapeshift<1>
   %2 = fir.slice %0, %0, %0 : (index, index, index) -> !fir.slice<1>