[flang] Make `AbstractResultsOpt` run on `ModuleOp`

[unterumarmung]

Summary

This change introduce refactoring for the AbstractResultsOpt pass. Now this pass run on mlir::ModuleOp instead of mlir::func::FuncOp. this allows further changes to be made more easily - like converting "abstract results" in fir::GlobalOp.

• mlir::func::FuncOp-related conversions were decoupled into FuncOpConversion
• ReturnOpConversion was inlined to FuncOpConversion

[unterumarmung]

Added commit from D129485

[unterumarmung]

Make the patch work & rebase

[unterumarmung]

Commit "[flang][NFC] Drop AbstractResultOptions structure" to be reviewed in D129485

[clementval]

In D129486#3645463, @unterumarmung wrote:

Commit "[flang][NFC] Drop AbstractResultOptions structure" to be reviewed in D129485

You should make this patch dependent on D129485 so you do not have to include the change here. A patch should show a single commit.

flang/lib/Optimizer/Dialect/FIRType.cpp

940 ↗

(On Diff #443954)

[unterumarmung]

Drop commit from D129485

[unterumarmung]

In D129486#3646132, @clementval wrote:

In D129486#3645463, @unterumarmung wrote:

Commit "[flang][NFC] Drop AbstractResultOptions structure" to be reviewed in D129485

You should make this patch dependent on D129485 so you do not have to include the change here. A patch should show a single commit.

Done, thank you!

[clementval]

In fact I think it will be better to keep the pass running on function so we can keep the multithreading from MLIR. Then for the GlobalOp we can just have a separate conversion pattern that reuse the code from here.

[unterumarmung]

In D129486#3647471, @clementval wrote:

In fact I think it will be better to keep the pass running on function so we can keep the multithreading from MLIR. Then for the GlobalOp we can just have a separate conversion pattern that reuse the code from here.

I didn't quite get how you want to reuse the conversion, but I can propose the following "architecture":

template <typename Op>
class AbstractResultBase { ... };

contains common logic of both FuncOp and GlobalOp

class AbstractResultInFuncOpt : AbstractResultBase<FuncOp>, /* Base from TableGen-generated code */ { ... };

contains rewriting of signature and adding ReturnOp conversion pattern to the mlir::RewritePatternSet patterns

class AbstractResultInGlobalOpt : AbstractResultBase<GlobalOp>, /* Base from TableGen-generated code */ { ... };

probably will be empty. Or is there some GlobalOp-specific quirks that have to be handled?

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"> {
+def AbstractResultOpt : Pass<"abstract-result-opt", "mlir::ModuleOp"> {
   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 = [

flang/include/flang/Tools/CLOptions.inc

inline void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm) {
   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.addPass(fir::createAbstractResultOptPass());
   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

 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #define DEBUG_TYPE "flang-abstract-result-opt"
 
 namespace fir {
 namespace {
 
-struct AbstractResultOptions {
-  // Always pass result as a fir.box argument.
-  bool boxResult = false;
-  // New function block argument for the result if the current FuncOp had
-  // an abstract result.
-  mlir::Value newArg;
-};
-
 static mlir::Type getResultArgumentType(mlir::Type resultType,
-                                        const AbstractResultOptions &options) {
+                                        bool shouldBoxResult) {
   return llvm::TypeSwitch<mlir::Type, mlir::Type>(resultType)
       .Case<fir::SequenceType, fir::RecordType>(
           [&](mlir::Type type) -> mlir::Type {
-            if (options.boxResult)
+            if (shouldBoxResult)
               return fir::BoxType::get(type);
             return fir::ReferenceType::get(type);
           })
       .Case<fir::BoxType>([](mlir::Type type) -> mlir::Type {
         return fir::ReferenceType::get(type);
       })
       .Default([](mlir::Type) -> mlir::Type {
         llvm_unreachable("bad abstract result type");
       });
 }
 
-static mlir::FunctionType
-getNewFunctionType(mlir::FunctionType funcTy,
-                   const AbstractResultOptions &options) {
+static mlir::FunctionType getNewFunctionType(mlir::FunctionType funcTy,
+                                             bool shouldBoxResult) {
   auto resultType = funcTy.getResult(0);
-  auto argTy = getResultArgumentType(resultType, options);
+  auto argTy = getResultArgumentType(resultType, shouldBoxResult);
   llvm::SmallVector<mlir::Type> newInputTypes = {argTy};
   newInputTypes.append(funcTy.getInputs().begin(), funcTy.getInputs().end());
   return mlir::FunctionType::get(funcTy.getContext(), newInputTypes,
                                  /*resultTypes=*/{});
 }
 
-static bool mustEmboxResult(mlir::Type resultType,
-                            const AbstractResultOptions &options) {
+static bool mustEmboxResult(mlir::Type resultType, bool shouldBoxResult) {
   return resultType.isa<fir::SequenceType, fir::RecordType>() &&
-         options.boxResult;
+         shouldBoxResult;
 }
 
 class CallOpConversion : public mlir::OpRewritePattern<fir::CallOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
-  CallOpConversion(mlir::MLIRContext *context, const AbstractResultOptions &opt)
-      : OpRewritePattern(context), options{opt} {}
+  CallOpConversion(mlir::MLIRContext *context, bool shouldBoxResult)
+      : OpRewritePattern(context), shouldBoxResult{shouldBoxResult} {}
   mlir::LogicalResult
   matchAndRewrite(fir::CallOp callOp,
                   mlir::PatternRewriter &rewriter) const override {
     auto loc = callOp.getLoc();
     auto result = callOp->getResult(0);
     if (!result.hasOneUse()) {
       mlir::emitError(loc,
                       "calls with abstract result must have exactly one user");
       return mlir::failure();
     }
     auto saveResult =
         mlir::dyn_cast<fir::SaveResultOp>(result.use_begin().getUser());
     if (!saveResult) {
       mlir::emitError(
           loc, "calls with abstract result must be used in fir.save_result");
       return mlir::failure();
     }
-    auto argType = getResultArgumentType(result.getType(), options);
+    auto argType = getResultArgumentType(result.getType(), shouldBoxResult);
     auto buffer = saveResult.getMemref();
     mlir::Value arg = buffer;
-    if (mustEmboxResult(result.getType(), options))
+    if (mustEmboxResult(result.getType(), shouldBoxResult))
       arg = rewriter.create<fir::EmboxOp>(
           loc, argType, buffer, saveResult.getShape(), /*slice*/ mlir::Value{},
           saveResult.getTypeparams());
 
     llvm::SmallVector<mlir::Type> newResultTypes;
     if (callOp.getCallee()) {
       llvm::SmallVector<mlir::Value> newOperands = {arg};
       newOperands.append(callOp.getOperands().begin(),
if (callOp.getCallee()) {
                                    newOperands);
     }
     callOp->dropAllReferences();
     rewriter.eraseOp(callOp);
     return mlir::success();
   }
 
 private:
-  const AbstractResultOptions &options;
+  bool shouldBoxResult;
 };
 
 class SaveResultOpConversion
     : public mlir::OpRewritePattern<fir::SaveResultOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
   SaveResultOpConversion(mlir::MLIRContext *context)
       : OpRewritePattern(context) {}
   mlir::LogicalResult
   matchAndRewrite(fir::SaveResultOp op,
                   mlir::PatternRewriter &rewriter) const override {
     rewriter.eraseOp(op);
     return mlir::success();
   }
 };
 
 class ReturnOpConversion : public mlir::OpRewritePattern<mlir::func::ReturnOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
-  ReturnOpConversion(mlir::MLIRContext *context,
-                     const AbstractResultOptions &opt)
-      : OpRewritePattern(context), options{opt} {}
+  ReturnOpConversion(mlir::MLIRContext *context, mlir::Value newArg)
+      : OpRewritePattern(context), newArg{newArg} {}
   mlir::LogicalResult
   matchAndRewrite(mlir::func::ReturnOp ret,
                   mlir::PatternRewriter &rewriter) const override {
     rewriter.setInsertionPoint(ret);
     auto returnedValue = ret.getOperand(0);
     bool replacedStorage = false;
     if (auto *op = returnedValue.getDefiningOp())
       if (auto load = mlir::dyn_cast<fir::LoadOp>(op)) {
         auto resultStorage = load.getMemref();
-        load.getMemref().replaceAllUsesWith(options.newArg);
+        load.getMemref().replaceAllUsesWith(newArg);
         replacedStorage = true;
         if (auto *alloc = resultStorage.getDefiningOp())
           if (alloc->use_empty())
             rewriter.eraseOp(alloc);
       }
     // The result storage may have been optimized out by a memory to
     // register pass, this is possible for fir.box results, or fir.record
     // with no length parameters. Simply store the result in the result storage.
     // at the return point.
     if (!replacedStorage)
-      rewriter.create<fir::StoreOp>(ret.getLoc(), returnedValue,
-                                    options.newArg);
+      rewriter.create<fir::StoreOp>(ret.getLoc(), returnedValue, newArg);
     rewriter.replaceOpWithNewOp<mlir::func::ReturnOp>(ret);
     return mlir::success();
   }
 
 private:
-  const AbstractResultOptions &options;
+  mlir::Value newArg;
 };
 
 class AddrOfOpConversion : public mlir::OpRewritePattern<fir::AddrOfOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
-  AddrOfOpConversion(mlir::MLIRContext *context,
-                     const AbstractResultOptions &opt)
-      : OpRewritePattern(context), options{opt} {}
+  AddrOfOpConversion(mlir::MLIRContext *context, bool shouldBoxResult)
+      : OpRewritePattern(context), shouldBoxResult{shouldBoxResult} {}
   mlir::LogicalResult
   matchAndRewrite(fir::AddrOfOp addrOf,
                   mlir::PatternRewriter &rewriter) const override {
     auto oldFuncTy = addrOf.getType().cast<mlir::FunctionType>();
-    auto newFuncTy = getNewFunctionType(oldFuncTy, options);
+    auto newFuncTy = getNewFunctionType(oldFuncTy, shouldBoxResult);
     auto newAddrOf = rewriter.create<fir::AddrOfOp>(addrOf.getLoc(), newFuncTy,
                                                     addrOf.getSymbol());
     // Rather than converting all op a function pointer might transit through
     // (e.g calls, stores, loads, converts...), cast new type to the abstract
     // type. A conversion will be added when calling indirect calls of abstract
     // types.
     rewriter.replaceOpWithNewOp<fir::ConvertOp>(addrOf, oldFuncTy, newAddrOf);
     return mlir::success();
   }
 
 private:
-  const AbstractResultOptions &options;
+  bool shouldBoxResult;
 };
 
-class AbstractResultOpt : public fir::AbstractResultOptBase<AbstractResultOpt> {
+class FuncOpConversion : public mlir::OpRewritePattern<mlir::func::FuncOp> {
 public:
-  void runOnOperation() override {
-    auto *context = &getContext();
-    auto func = getOperation();
+  using OpRewritePattern::OpRewritePattern;
+  FuncOpConversion(mlir::MLIRContext *context, bool shouldBoxResult)
+      : OpRewritePattern(context), shouldBoxResult{shouldBoxResult} {}
+  mlir::LogicalResult
+  matchAndRewrite(mlir::func::FuncOp func,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto *context = getContext();
     auto loc = func.getLoc();
     mlir::RewritePatternSet patterns(context);
     mlir::ConversionTarget target = *context;
-    AbstractResultOptions options{passResultAsBox.getValue(),
-                                  /*newArg=*/{}};
 
     // Convert function type itself if it has an abstract result
     auto funcTy = func.getFunctionType().cast<mlir::FunctionType>();
     if (hasAbstractResult(funcTy)) {
-      func.setType(getNewFunctionType(funcTy, options));
-      unsigned zero = 0;
-      if (!func.empty()) {
-        // Insert new argument
-        mlir::OpBuilder rewriter(context);
-        auto resultType = funcTy.getResult(0);
-        auto argTy = getResultArgumentType(resultType, options);
-        options.newArg = func.front().insertArgument(zero, argTy, loc);
-        if (mustEmboxResult(resultType, options)) {
-          auto bufferType = fir::ReferenceType::get(resultType);
-          rewriter.setInsertionPointToStart(&func.front());
-          options.newArg =
-              rewriter.create<fir::BoxAddrOp>(loc, bufferType, options.newArg);
-        }
-        patterns.insert<ReturnOpConversion>(context, options);
-        target.addDynamicallyLegalOp<mlir::func::ReturnOp>(
-            [](mlir::func::ReturnOp ret) { return ret.operands().empty(); });
-      }
-    }
+      func.setType(getNewFunctionType(funcTy, shouldBoxResult));
+      if (func.empty()) {
+        // If a function has no body, its conversion is done
+        return mlir::success();
+      }
+      auto argTy = getResultArgumentType(resultType, shouldBoxResult);
+      mlir::Value newArg = func.front().insertArgument(0u, argTy, loc);
+      if (mustEmboxResult(resultType, shouldBoxResult)) {
+        newArg = rewriter.create<fir::BoxAddrOp>(loc, bufferType, newArg);
+      patterns.insert<ReturnOpConversion>(context, newArg);
+
+    return mlir::applyPartialConversion(func, target, std::move(patterns));
 
-    if (func.empty())
-      return;
+private:
+  bool shouldBoxResult;
+};
+
+class AbstractResultOpt : public fir::AbstractResultOptBase<AbstractResultOpt> {
+public:
+  void runOnOperation() override {
+    auto *context = &getContext();
+    auto module = getOperation();
+    mlir::RewritePatternSet patterns(context);
+    mlir::ConversionTarget target = *context;
+    const bool shouldBoxResult = passResultAsBox.getValue();
 
     // 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::AddrOfOp>([](fir::AddrOfOp addrOf) {
       if (auto funTy = addrOf.getType().dyn_cast<mlir::FunctionType>())
         return !hasAbstractResult(funTy);
       return true;
     });
     target.addDynamicallyLegalOp<fir::DispatchOp>([](fir::DispatchOp dispatch) {
       if (dispatch->getNumResults() != 1)
         return true;
       auto resultType = dispatch->getResult(0).getType();
       if (resultType.isa<fir::SequenceType, fir::BoxType, fir::RecordType>()) {
         TODO(dispatch.getLoc(), "dispatchOp with abstract results");
         return false;
       }
       return true;
     });
 
-    patterns.insert<CallOpConversion>(context, options);
+    patterns.insert<CallOpConversion>(context, shouldBoxResult);
+    patterns.insert<FuncOpConversion>(context, shouldBoxResult);
     patterns.insert<SaveResultOpConversion>(context);
-    patterns.insert<AddrOfOpConversion>(context, options);
-    if (mlir::failed(
-            mlir::applyPartialConversion(func, target, std::move(patterns)))) {
-      mlir::emitError(func.getLoc(), "error in converting abstract results\n");
+    patterns.insert<AddrOfOpConversion>(context, shouldBoxResult);
+    if (mlir::failed(mlir::applyPartialConversion(module, target,
+                                                  std::move(patterns)))) {
+      mlir::emitError(module.getLoc(),
+                      "error in converting abstract results\n");
       signalPassFailure();
     }
   }
 };
 } // end anonymous namespace
 } // namespace fir
 
 std::unique_ptr<mlir::Pass> fir::createAbstractResultOptPass() {
   return std::make_unique<AbstractResultOpt>();
 }

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

 
 ! CHECK-LABEL: 'func.func' Pipeline
 ! CHECK: CFGConversion
 ! CHECK: SCFToControlFlow
 ! CHECK: Canonicalizer
 ! CHECK: SimplifyRegionLite
 ! CHECK: CSE
 ! CHECK: BoxedProcedurePass
-
-! CHECK-LABEL: 'func.func' Pipeline
 ! CHECK: AbstractResultOpt
 ! CHECK: CodeGenRewrite
 ! CHECK: TargetRewrite
 ! CHECK: ExternalNameConversion
 ! CHECK: FIRToLLVMLowering
 ! CHECK-NOT: LLVMIRLoweringPass

flang/test/Fir/basic-program.fir

 
 // PASSES-LABEL: 'func.func' Pipeline
 // PASSES: CFGConversion
 // PASSES: SCFToControlFlow
 // PASSES: Canonicalizer
 // PASSES: SimplifyRegionLite
 // PASSES: CSE
 // PASSES: BoxedProcedurePass
-
-// PASSES-LABEL: 'func.func' Pipeline
 // PASSES: AbstractResultOpt
 // PASSES: CodeGenRewrite
 // PASSES: TargetRewrite
 // PASSES: FIRToLLVMLowering
 // PASSES: LLVMIRLoweringPass