[flang] Generalize `AbstractResultOpt` pass

[unterumarmung]

Summary

This change decouples common functionality for convering abstract
results, so it can be reused later.

Depends on D129485

[clementval]

It would be good to have the patch for GlobalOp as well so we can see how it all fits together.

flang/include/flang/Tools/CLOptions.inc

74

Unrelated change

[unterumarmung]

Removed unrelated change & rebased

[unterumarmung]

Rebased & resolved merge conflicts

[jeanPerier]

Looks good to me, thanks

[clementval]

LGTM. Just fix couple of typos and punctuations before merging.

flang/lib/Optimizer/Transforms/AbstractResult.cpp

195

196

197–200

258

263

[unterumarmung]

Rebased and fixed style issues

flang/include/flang/Optimizer/Transforms/Passes.h

 } // namespace mlir
 
 namespace fir {
 
 //===----------------------------------------------------------------------===//
 // Passes defined in Passes.td
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<mlir::Pass> createAbstractResultOptPass();
+std::unique_ptr<mlir::Pass> createAbstractResultOnFuncOptPass();
 std::unique_ptr<mlir::Pass> createAffineDemotionPass();
 std::unique_ptr<mlir::Pass> createArrayValueCopyPass();
 std::unique_ptr<mlir::Pass> createFirToCfgPass();
 std::unique_ptr<mlir::Pass> createCharacterConversionPass();
 std::unique_ptr<mlir::Pass> createExternalNameConversionPass();
 std::unique_ptr<mlir::Pass> createMemDataFlowOptPass();
 std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 std::unique_ptr<mlir::Pass> createMemoryAllocationPass();

flang/include/flang/Optimizer/Transforms/Passes.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FLANG_OPTIMIZER_TRANSFORMS_PASSES
 #define FLANG_OPTIMIZER_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def AbstractResultOpt : Pass<"abstract-result-opt", "mlir::func::FuncOp"> {
+class AbstractResultOptBase<string optExt, string operation> 
+  : Pass<"abstract-result-on-" # optExt # "-opt", operation> {
   let summary = "Convert fir.array, fir.box and fir.rec function result to "
                 "function argument";
   let description = [{
     This pass is required before code gen to the LLVM IR dialect,
     including the pre-cg rewrite pass.
   }];
-  let constructor = "::fir::createAbstractResultOptPass()";
   let dependentDialects = [
     "fir::FIROpsDialect", "mlir::func::FuncDialect"
   ];
   let options = [
     Option<"passResultAsBox", "abstract-result-as-box",
            "bool", /*default=*/"false",
            "Pass fir.array<T> result as fir.box<fir.array<T>> argument instead"
            " of fir.ref<fir.array<T>>.">
   ];
 }
 
+def AbstractResultOnFuncOpt : AbstractResultOptBase<"func", "mlir::func::FuncOp"> {
+  let constructor = "::fir::createAbstractResultOnFuncOptPass()";
+}
+
 def AffineDialectPromotion : Pass<"promote-to-affine", "::mlir::func::FuncOp"> {
   let summary = "Promotes `fir.{do_loop,if}` to `affine.{for,if}`.";
   let description = [{
     Convert fir operations which satisfy affine constraints to the affine
     dialect.
 
     `fir.do_loop` will be converted to `affine.for` if the loops inside the body
     can be converted and the indices for memory loads and stores satisfy

flang/include/flang/Tools/CLOptions.inc

 DisableOption(CodeGenRewrite, "codegen-rewrite", "rewrite FIR for codegen");
 DisableOption(TargetRewrite, "target-rewrite", "rewrite FIR for target");
 DisableOption(FirToLlvmIr, "fir-to-llvmir", "FIR to LLVM-IR dialect");
 DisableOption(LlvmIrToLlvm, "llvm", "conversion to LLVM");
 DisableOption(BoxedProcedureRewrite, "boxed-procedure-rewrite",
     "rewrite boxed procedures");
 #endif
 
 DisableOption(ExternalNameConversion, "external-name-interop", "convert names with external convention");
 
 /// Generic for adding a pass to the pass manager if it is not disabled.
 template <typename F>
 void addPassConditionally(
     mlir::PassManager &pm, llvm::cl::opt<bool> &disabled, F ctor) {
   if (!disabled)
     pm.addPass(ctor());
 }
inline void createDefaultFIROptimizerPassPipeline(
   pm.addPass(mlir::createCanonicalizerPass(config));
   pm.addPass(fir::createSimplifyRegionLitePass());
   pm.addPass(mlir::createCSEPass());
 }
 
 #if !defined(FLANG_EXCLUDE_CODEGEN)
 inline void createDefaultFIRCodeGenPassPipeline(mlir::PassManager &pm) {
   fir::addBoxedProcedurePass(pm);
-  pm.addNestedPass<mlir::func::FuncOp>(fir::createAbstractResultOptPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      fir::createAbstractResultOnFuncOptPass());
   fir::addCodeGenRewritePass(pm);
   fir::addTargetRewritePass(pm);
   fir::addExternalNameConversionPass(pm);
   fir::addFIRToLLVMPass(pm);
 }
 
 /// Create a pass pipeline for lowering from MLIR to LLVM IR
 ///

flang/lib/Optimizer/Transforms/AbstractResult.cpp

matchAndRewrite(fir::AddrOfOp addrOf,
     rewriter.replaceOpWithNewOp<fir::ConvertOp>(addrOf, oldFuncTy, newAddrOf);
     return mlir::success();
   }
 
 private:
   bool shouldBoxResult;
 };
 
-class AbstractResultOpt : public fir::AbstractResultOptBase<AbstractResultOpt> {
+/// @brief Base CRTP class for AbstractResult pass family.
+/// Contains common logic for abstract result conversion in a reusable fashion.
+/// @tparam Pass target class that implements operation-specific logic.
+/// @tparam PassBase base class template for the pass generated by TableGen.
+/// The `Pass` class must define runOnSpecificOperation(OpTy, bool,
+/// mlir::RewritePatternSet&, mlir::ConversionTarget&) member function.
+/// This function should implement operation-specific functionality.
+template <typename Pass, template <typename> class PassBase>
+class AbstractResultOptTemplate : public PassBase<Pass> {
 public:
   void runOnOperation() override {
-    auto *context = &getContext();
-    auto func = getOperation();
-    auto loc = func.getLoc();
+    auto *context = &this->getContext();
+    auto op = this->getOperation();
+
     mlir::RewritePatternSet patterns(context);
     mlir::ConversionTarget target = *context;
-    const bool shouldBoxResult = passResultAsBox.getValue();
-
-    // Convert function type itself if it has an abstract result
-    auto funcTy = func.getFunctionType().cast<mlir::FunctionType>();
-    if (hasAbstractResult(funcTy)) {
-      func.setType(getNewFunctionType(funcTy, shouldBoxResult));
-      unsigned zero = 0;
-      if (!func.empty()) {
-        // Insert new argument
-        mlir::OpBuilder rewriter(context);
-        auto resultType = funcTy.getResult(0);
-        auto argTy = getResultArgumentType(resultType, shouldBoxResult);
-        mlir::Value newArg = func.front().insertArgument(zero, argTy, loc);
-        if (mustEmboxResult(resultType, shouldBoxResult)) {
-          auto bufferType = fir::ReferenceType::get(resultType);
-          rewriter.setInsertionPointToStart(&func.front());
-          newArg = rewriter.create<fir::BoxAddrOp>(loc, bufferType, newArg);
-        }
-        patterns.insert<ReturnOpConversion>(context, newArg);
-        target.addDynamicallyLegalOp<mlir::func::ReturnOp>(
-            [](mlir::func::ReturnOp ret) { return ret.operands().empty(); });
-      }
-    }
+    const bool shouldBoxResult = this->passResultAsBox.getValue();
 
-    if (func.empty())
-      return;
+    auto &self = static_cast<Pass &>(*this);
+    self.runOnSpecificOperation(op, shouldBoxResult, patterns, target);
 
     // Convert the calls and, if needed,  the ReturnOp in the function body.
     target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithmeticDialect,
                            mlir::func::FuncDialect>();
     target.addIllegalOp<fir::SaveResultOp>();
     target.addDynamicallyLegalOp<fir::CallOp>([](fir::CallOp call) {
       return !hasAbstractResult(call.getFunctionType());
     });
target.addDynamicallyLegalOp<fir::DispatchOp>([](fir::DispatchOp dispatch) {
       }
       return true;
     });
 
     patterns.insert<CallOpConversion>(context, shouldBoxResult);
     patterns.insert<SaveResultOpConversion>(context);
     patterns.insert<AddrOfOpConversion>(context, shouldBoxResult);
     if (mlir::failed(
-            mlir::applyPartialConversion(func, target, std::move(patterns)))) {
-      mlir::emitError(func.getLoc(), "error in converting abstract results\n");
-      signalPassFailure();
+            mlir::applyPartialConversion(op, target, std::move(patterns)))) {
+      mlir::emitError(op.getLoc(), "error in converting abstract results\n");
+      this->signalPassFailure();
+    }
+  }
+};
+
+class AbstractResultOnFuncOpt
+    : public AbstractResultOptTemplate<AbstractResultOnFuncOpt,
+                                       fir::AbstractResultOnFuncOptBase> {
+public:
+  void runOnSpecificOperation(mlir::func::FuncOp func, bool shouldBoxResult,
+                              mlir::RewritePatternSet &patterns,
+                              mlir::ConversionTarget &target) {
+    auto loc = func.getLoc();
+    auto *context = &getContext();
+    // Convert function type itself if it has an abstract result.
+    auto funcTy = func.getFunctionType().cast<mlir::FunctionType>();
+    if (hasAbstractResult(funcTy)) {
+      func.setType(getNewFunctionType(funcTy, shouldBoxResult));
+      if (!func.empty()) {
+        // Insert new argument.
+        mlir::OpBuilder rewriter(context);
+        auto resultType = funcTy.getResult(0);
+        auto argTy = getResultArgumentType(resultType, shouldBoxResult);
+        mlir::Value newArg = func.front().insertArgument(0u, argTy, loc);
+        if (mustEmboxResult(resultType, shouldBoxResult)) {
+          auto bufferType = fir::ReferenceType::get(resultType);
+          rewriter.setInsertionPointToStart(&func.front());
+          newArg = rewriter.create<fir::BoxAddrOp>(loc, bufferType, newArg);
+        }
+        patterns.insert<ReturnOpConversion>(context, newArg);
+        target.addDynamicallyLegalOp<mlir::func::ReturnOp>(
+            [](mlir::func::ReturnOp ret) { return ret.operands().empty(); });
+      }
     }
   }
 };
 } // end anonymous namespace
 } // namespace fir
 
-std::unique_ptr<mlir::Pass> fir::createAbstractResultOptPass() {
-  return std::make_unique<AbstractResultOpt>();
+std::unique_ptr<mlir::Pass> fir::createAbstractResultOnFuncOptPass() {
+  return std::make_unique<AbstractResultOnFuncOpt>();
 }

flang/test/Driver/mlir-pass-pipeline.f90

 ! ALL-NEXT: Canonicalizer
 ! ALL-NEXT: SimplifyRegionLite
 ! ALL-NEXT: CSE
 ! ALL-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 ! ALL-NEXT: BoxedProcedurePass
 
 ! ALL-NEXT: 'func.func' Pipeline
-! ALL-NEXT:   AbstractResultOpt
+! ALL-NEXT:   AbstractResultOnFuncOpt
 
 ! ALL-NEXT: CodeGenRewrite
 ! ALL-NEXT:   (S) 0 num-dce'd - Number of operations eliminated
 ! ALL-NEXT: TargetRewrite
 ! ALL-NEXT: ExternalNameConversion
 ! ALL-NEXT: FIRToLLVMLowering
 ! ALL-NOT: LLVMIRLoweringPass

flang/test/Fir/abstract-results.fir

 // Test rewrite of functions that return fir.array<>, fir.type<>, fir.box<> to
 // functions that take an additional argument for the result.
 
-// RUN: fir-opt %s --abstract-result-opt | FileCheck %s
-// RUN: fir-opt %s --abstract-result-opt=abstract-result-as-box | FileCheck %s --check-prefix=CHECK-BOX
+// RUN: fir-opt %s --abstract-result-on-func-opt | FileCheck %s
+// RUN: fir-opt %s --abstract-result-on-func-opt=abstract-result-as-box | FileCheck %s --check-prefix=CHECK-BOX
 
 // ----------------------- Test declaration rewrite ----------------------------
 
 // CHECK-LABEL:  func private @arrayfunc(!fir.ref<!fir.array<?xf32>>, i32)
 // CHECK-BOX-LABEL:  func private @arrayfunc(!fir.box<!fir.array<?xf32>>, i32)
 func.func private @arrayfunc(i32) -> !fir.array<?xf32>
 
 // CHECK-LABEL:  func private @derivedfunc(!fir.ref<!fir.type<t{x:f32}>>, f32)

flang/test/Fir/basic-program.fir

 // PASSES-NEXT: Canonicalizer
 // PASSES-NEXT: SimplifyRegionLite
 // PASSES-NEXT: CSE
 // PASSES-NEXT:   (S) 0 num-cse'd - Number of operations CSE'd
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations DCE'd
 // PASSES-NEXT: BoxedProcedurePass
 
 // PASSES-NEXT: 'func.func' Pipeline
-// PASSES-NEXT:   AbstractResultOpt
+// PASSES-NEXT:   AbstractResultOnFuncOpt
 
 // PASSES-NEXT: CodeGenRewrite
 // PASSES-NEXT:   (S) 0 num-dce'd - Number of operations eliminated
 // PASSES-NEXT: TargetRewrite
 // PASSES-NEXT: FIRToLLVMLowering
 // PASSES-NEXT: LLVMIRLoweringPass