[flang][OpenMP][NFC] Outline genOpWithBody & createBodyOfOp args

[ergawy]

This PR outlines the arguments of the open CodeGen functions into 2 separate structs. This was, in part, motivated by the delayed privatization WIP #79862 where we had to extend the signatures of both functions containing quite a bit of default values (nullptr, false). This PR does not add any new arguments yet though, just outlines the existing ones.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-flang-fir-hlfir

Author: Kareem Ergawy (ergawy)

Changes

This PR outlines the arguments of the open CodeGen functions into 2 separate structs. This was, in part, motivated by the delayed privatization WIP #79862 where we had to extend the signatures of both functions containing quite a bit of default values (nullptr, false). This PR does not add any new arguments yet though, just outlines the existing ones.

---

Full diff: https://github.com/llvm/llvm-project/pull/80817.diff

1 Files Affected:

• (modified) flang/lib/Lower/OpenMP.cpp (+123-77)

diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 0a68aba162618b..97fffede563643 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2250,6 +2250,17 @@ createAndSetPrivatizedLoopVar(Fortran::lower::AbstractConverter &converter,
   return storeOp;
 }
 
+struct CreateBodyOfOpInfo {
+  Fortran::lower::AbstractConverter &converter;
+  mlir::Location &loc;
+  Fortran::lower::pft::Evaluation &eval;
+  bool genNested = true;
+  const Fortran::parser::OmpClauseList *clauses = nullptr;
+  const llvm::SmallVector<const Fortran::semantics::Symbol *> &args = {};
+  bool outerCombined = false;
+  DataSharingProcessor *dsp = nullptr;
+};
+
 /// Create the body (block) for an OpenMP Operation.
 ///
 /// \param [in]    op - the operation the body belongs to.
@@ -2263,13 +2274,8 @@ createAndSetPrivatizedLoopVar(Fortran::lower::AbstractConverter &converter,
 /// \param [in]    outerCombined - is this an outer operation - prevents
 ///                                privatization.
 template <typename Op>
-static void createBodyOfOp(
-    Op &op, Fortran::lower::AbstractConverter &converter, mlir::Location &loc,
-    Fortran::lower::pft::Evaluation &eval, bool genNested,
-    const Fortran::parser::OmpClauseList *clauses = nullptr,
-    const llvm::SmallVector<const Fortran::semantics::Symbol *> &args = {},
-    bool outerCombined = false, DataSharingProcessor *dsp = nullptr) {
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+static void createBodyOfOp(Op &op, CreateBodyOfOpInfo info) {
+  fir::FirOpBuilder &firOpBuilder = info.converter.getFirOpBuilder();
 
   auto insertMarker = [](fir::FirOpBuilder &builder) {
     mlir::Value undef = builder.create<fir::UndefOp>(builder.getUnknownLoc(),
@@ -2281,22 +2287,22 @@ static void createBodyOfOp(
   // argument. Also update the symbol's address with the mlir argument value.
   // e.g. For loops the argument is the induction variable. And all further
   // uses of the induction variable should use this mlir value.
-  if (args.size()) {
+  if (info.args.size()) {
     std::size_t loopVarTypeSize = 0;
-    for (const Fortran::semantics::Symbol *arg : args)
+    for (const Fortran::semantics::Symbol *arg : info.args)
       loopVarTypeSize = std::max(loopVarTypeSize, arg->GetUltimate().size());
-    mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
-    llvm::SmallVector<mlir::Type> tiv(args.size(), loopVarType);
-    llvm::SmallVector<mlir::Location> locs(args.size(), loc);
+    mlir::Type loopVarType = getLoopVarType(info.converter, loopVarTypeSize);
+    llvm::SmallVector<mlir::Type> tiv(info.args.size(), loopVarType);
+    llvm::SmallVector<mlir::Location> locs(info.args.size(), info.loc);
     firOpBuilder.createBlock(&op.getRegion(), {}, tiv, locs);
     // The argument is not currently in memory, so make a temporary for the
     // argument, and store it there, then bind that location to the argument.
     mlir::Operation *storeOp = nullptr;
-    for (auto [argIndex, argSymbol] : llvm::enumerate(args)) {
+    for (auto [argIndex, argSymbol] : llvm::enumerate(info.args)) {
       mlir::Value indexVal =
           fir::getBase(op.getRegion().front().getArgument(argIndex));
-      storeOp =
-          createAndSetPrivatizedLoopVar(converter, loc, indexVal, argSymbol);
+      storeOp = createAndSetPrivatizedLoopVar(info.converter, info.loc,
+                                              indexVal, argSymbol);
     }
     firOpBuilder.setInsertionPointAfter(storeOp);
   } else {
@@ -2308,44 +2314,44 @@ static void createBodyOfOp(
 
   // If it is an unstructured region and is not the outer region of a combined
   // construct, create empty blocks for all evaluations.
-  if (eval.lowerAsUnstructured() && !outerCombined)
+  if (info.eval.lowerAsUnstructured() && !info.outerCombined)
     Fortran::lower::createEmptyRegionBlocks<mlir::omp::TerminatorOp,
                                             mlir::omp::YieldOp>(
-        firOpBuilder, eval.getNestedEvaluations());
+        firOpBuilder, info.eval.getNestedEvaluations());
 
   // Start with privatization, so that the lowering of the nested
   // code will use the right symbols.
   constexpr bool isLoop = std::is_same_v<Op, mlir::omp::WsLoopOp> ||
                           std::is_same_v<Op, mlir::omp::SimdLoopOp>;
-  bool privatize = clauses && !outerCombined;
+  bool privatize = info.clauses && !info.outerCombined;
 
   firOpBuilder.setInsertionPoint(marker);
   std::optional<DataSharingProcessor> tempDsp;
   if (privatize) {
-    if (!dsp) {
-      tempDsp.emplace(converter, *clauses, eval);
+    if (!info.dsp) {
+      tempDsp.emplace(info.converter, *info.clauses, info.eval);
       tempDsp->processStep1();
     }
   }
 
   if constexpr (std::is_same_v<Op, mlir::omp::ParallelOp>) {
-    threadPrivatizeVars(converter, eval);
-    if (clauses) {
+    threadPrivatizeVars(info.converter, info.eval);
+    if (info.clauses) {
       firOpBuilder.setInsertionPoint(marker);
-      ClauseProcessor(converter, *clauses).processCopyin();
+      ClauseProcessor(info.converter, *info.clauses).processCopyin();
     }
   }
 
-  if (genNested) {
+  if (info.genNested) {
     // genFIR(Evaluation&) tries to patch up unterminated blocks, causing
     // a lot of complications for our approach if the terminator generation
     // is delayed past this point. Insert a temporary terminator here, then
     // delete it.
     firOpBuilder.setInsertionPointToEnd(&op.getRegion().back());
-    auto *temp = Fortran::lower::genOpenMPTerminator(firOpBuilder,
-                                                     op.getOperation(), loc);
+    auto *temp = Fortran::lower::genOpenMPTerminator(
+        firOpBuilder, op.getOperation(), info.loc);
     firOpBuilder.setInsertionPointAfter(marker);
-    genNestedEvaluations(converter, eval);
+    genNestedEvaluations(info.converter, info.eval);
     temp->erase();
   }
 
@@ -2380,28 +2386,28 @@ static void createBodyOfOp(
     mlir::Block *exit = firOpBuilder.createBlock(&region);
     for (mlir::Block *b : exits) {
       firOpBuilder.setInsertionPointToEnd(b);
-      firOpBuilder.create<mlir::cf::BranchOp>(loc, exit);
+      firOpBuilder.create<mlir::cf::BranchOp>(info.loc, exit);
     }
     return exit;
   };
 
   if (auto *exitBlock = getUniqueExit(op.getRegion())) {
     firOpBuilder.setInsertionPointToEnd(exitBlock);
-    auto *term = Fortran::lower::genOpenMPTerminator(firOpBuilder,
-                                                     op.getOperation(), loc);
+    auto *term = Fortran::lower::genOpenMPTerminator(
+        firOpBuilder, op.getOperation(), info.loc);
     // Only insert lastprivate code when there actually is an exit block.
     // Such a block may not exist if the nested code produced an infinite
     // loop (this may not make sense in production code, but a user could
     // write that and we should handle it).
     firOpBuilder.setInsertionPoint(term);
     if (privatize) {
-      if (!dsp) {
+      if (!info.dsp) {
         assert(tempDsp.has_value());
         tempDsp->processStep2(op, isLoop);
       } else {
-        if (isLoop && args.size() > 0)
-          dsp->setLoopIV(converter.getSymbolAddress(*args[0]));
-        dsp->processStep2(op, isLoop);
+        if (isLoop && info.args.size() > 0)
+          info.dsp->setLoopIV(info.converter.getSymbolAddress(*info.args[0]));
+        info.dsp->processStep2(op, isLoop);
       }
     }
   }
@@ -2475,17 +2481,25 @@ static void genBodyOfTargetDataOp(
     genNestedEvaluations(converter, eval);
 }
 
+struct GenOpWithBodyInfo {
+  Fortran::lower::AbstractConverter &converter;
+  Fortran::lower::pft::Evaluation &eval;
+  bool genNested = false;
+  mlir::Location currentLocation;
+  bool outerCombined = false;
+  const Fortran::parser::OmpClauseList *clauseList = nullptr;
+};
+
 template <typename OpTy, typename... Args>
-static OpTy genOpWithBody(Fortran::lower::AbstractConverter &converter,
-                          Fortran::lower::pft::Evaluation &eval, bool genNested,
-                          mlir::Location currentLocation, bool outerCombined,
-                          const Fortran::parser::OmpClauseList *clauseList,
-                          Args &&...args) {
-  auto op = converter.getFirOpBuilder().create<OpTy>(
-      currentLocation, std::forward<Args>(args)...);
-  createBodyOfOp<OpTy>(op, converter, currentLocation, eval, genNested,
-                       clauseList,
-                       /*args=*/{}, outerCombined);
+static OpTy genOpWithBody(GenOpWithBodyInfo info, Args &&...args) {
+  auto op = info.converter.getFirOpBuilder().create<OpTy>(
+      info.currentLocation, std::forward<Args>(args)...);
+  createBodyOfOp<OpTy>(op, {.converter = info.converter,
+                            .loc = info.currentLocation,
+                            .eval = info.eval,
+                            .genNested = info.genNested,
+                            .clauses = info.clauseList,
+                            .outerCombined = info.outerCombined});
   return op;
 }
 
@@ -2493,11 +2507,12 @@ static mlir::omp::MasterOp
 genMasterOp(Fortran::lower::AbstractConverter &converter,
             Fortran::lower::pft::Evaluation &eval, bool genNested,
             mlir::Location currentLocation) {
-  return genOpWithBody<mlir::omp::MasterOp>(converter, eval, genNested,
-                                            currentLocation,
-                                            /*outerCombined=*/false,
-                                            /*clauseList=*/nullptr,
-                                            /*resultTypes=*/mlir::TypeRange());
+  return genOpWithBody<mlir::omp::MasterOp>(
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation},
+      /*resultTypes=*/mlir::TypeRange());
 }
 
 static mlir::omp::OrderedRegionOp
@@ -2505,9 +2520,11 @@ genOrderedRegionOp(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval, bool genNested,
                    mlir::Location currentLocation) {
   return genOpWithBody<mlir::omp::OrderedRegionOp>(
-      converter, eval, genNested, currentLocation,
-      /*outerCombined=*/false,
-      /*clauseList=*/nullptr, /*simd=*/false);
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation},
+      /*simd=*/false);
 }
 
 static mlir::omp::ParallelOp
@@ -2534,7 +2551,12 @@ genParallelOp(Fortran::lower::AbstractConverter &converter,
     cp.processReduction(currentLocation, reductionVars, reductionDeclSymbols);
 
   return genOpWithBody<mlir::omp::ParallelOp>(
-      converter, eval, genNested, currentLocation, outerCombined, &clauseList,
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation,
+       .outerCombined = outerCombined,
+       .clauseList = &clauseList},
       /*resultTypes=*/mlir::TypeRange(), ifClauseOperand,
       numThreadsClauseOperand, allocateOperands, allocatorOperands,
       reductionVars,
@@ -2553,8 +2575,11 @@ genSectionOp(Fortran::lower::AbstractConverter &converter,
   // Currently only private/firstprivate clause is handled, and
   // all privatization is done within `omp.section` operations.
   return genOpWithBody<mlir::omp::SectionOp>(
-      converter, eval, genNested, currentLocation,
-      /*outerCombined=*/false, &sectionsClauseList);
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation,
+       .clauseList = &sectionsClauseList});
 }
 
 static mlir::omp::SingleOp
@@ -2573,10 +2598,13 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
 
   ClauseProcessor(converter, endClauseList).processNowait(nowaitAttr);
 
-  return genOpWithBody<mlir::omp::SingleOp>(
-      converter, eval, genNested, currentLocation,
-      /*outerCombined=*/false, &beginClauseList, allocateOperands,
-      allocatorOperands, nowaitAttr);
+  return genOpWithBody<mlir::omp::SingleOp>({.converter = converter,
+                                             .eval = eval,
+                                             .genNested = genNested,
+                                             .currentLocation = currentLocation,
+                                             .clauseList = &beginClauseList},
+                                            allocateOperands, allocatorOperands,
+                                            nowaitAttr);
 }
 
 static mlir::omp::TaskOp
@@ -2607,9 +2635,12 @@ genTaskOp(Fortran::lower::AbstractConverter &converter,
       currentLocation, llvm::omp::Directive::OMPD_task);
 
   return genOpWithBody<mlir::omp::TaskOp>(
-      converter, eval, genNested, currentLocation,
-      /*outerCombined=*/false, &clauseList, ifClauseOperand, finalClauseOperand,
-      untiedAttr, mergeableAttr,
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation,
+       .clauseList = &clauseList},
+      ifClauseOperand, finalClauseOperand, untiedAttr, mergeableAttr,
       /*in_reduction_vars=*/mlir::ValueRange(),
       /*in_reductions=*/nullptr, priorityClauseOperand,
       dependTypeOperands.empty()
@@ -2630,8 +2661,11 @@ genTaskGroupOp(Fortran::lower::AbstractConverter &converter,
   cp.processTODO<Fortran::parser::OmpClause::TaskReduction>(
       currentLocation, llvm::omp::Directive::OMPD_taskgroup);
   return genOpWithBody<mlir::omp::TaskGroupOp>(
-      converter, eval, genNested, currentLocation,
-      /*outerCombined=*/false, &clauseList,
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation,
+       .clauseList = &clauseList},
       /*task_reduction_vars=*/mlir::ValueRange(),
       /*task_reductions=*/nullptr, allocateOperands, allocatorOperands);
 }
@@ -3014,7 +3048,12 @@ genTeamsOp(Fortran::lower::AbstractConverter &converter,
       currentLocation, llvm::omp::Directive::OMPD_teams);
 
   return genOpWithBody<mlir::omp::TeamsOp>(
-      converter, eval, genNested, currentLocation, outerCombined, &clauseList,
+      {.converter = converter,
+       .eval = eval,
+       .genNested = genNested,
+       .currentLocation = currentLocation,
+       .outerCombined = outerCombined,
+       .clauseList = &clauseList},
       /*num_teams_lower=*/nullptr, numTeamsClauseOperand, ifClauseOperand,
       threadLimitClauseOperand, allocateOperands, allocatorOperands,
       reductionVars,
@@ -3258,9 +3297,13 @@ createSimdLoop(Fortran::lower::AbstractConverter &converter,
 
   auto *nestedEval = getCollapsedLoopEval(
       eval, Fortran::lower::getCollapseValue(loopOpClauseList));
-  createBodyOfOp<mlir::omp::SimdLoopOp>(simdLoopOp, converter, loc, *nestedEval,
-                                        /*genNested=*/true, &loopOpClauseList,
-                                        iv, /*outer=*/false, &dsp);
+  createBodyOfOp<mlir::omp::SimdLoopOp>(simdLoopOp,
+                                        {.converter = converter,
+                                         .loc = loc,
+                                         .eval = *nestedEval,
+                                         .clauses = &loopOpClauseList,
+                                         .args = iv,
+                                         .dsp = &dsp});
 }
 
 static void createWsLoop(Fortran::lower::AbstractConverter &converter,
@@ -3333,9 +3376,12 @@ static void createWsLoop(Fortran::lower::AbstractConverter &converter,
 
   auto *nestedEval = getCollapsedLoopEval(
       eval, Fortran::lower::getCollapseValue(beginClauseList));
-  createBodyOfOp<mlir::omp::WsLoopOp>(wsLoopOp, converter, loc, *nestedEval,
-                                      /*genNested=*/true, &beginClauseList, iv,
-                                      /*outer=*/false, &dsp);
+  createBodyOfOp<mlir::omp::WsLoopOp>(wsLoopOp, {.converter = converter,
+                                                 .loc = loc,
+                                                 .eval = *nestedEval,
+                                                 .clauses = &beginClauseList,
+                                                 .args = iv,
+                                                 .dsp = &dsp});
 }
 
 static void createSimdWsLoop(
@@ -3616,8 +3662,9 @@ genOMP(Fortran::lower::AbstractConverter &converter,
         currentLocation, mlir::FlatSymbolRefAttr::get(firOpBuilder.getContext(),
                                                       global.getSymName()));
   }();
-  createBodyOfOp<mlir::omp::CriticalOp>(criticalOp, converter, currentLocation,
-                                        eval, /*genNested=*/true);
+  createBodyOfOp<mlir::omp::CriticalOp>(
+      criticalOp,
+      {.converter = converter, .loc = currentLocation, .eval = eval});
 }
 
 static void
@@ -3659,10 +3706,9 @@ genOMP(Fortran::lower::AbstractConverter &converter,
   }
 
   // SECTIONS construct
-  genOpWithBody<mlir::omp::SectionsOp>(converter, eval,
-                                       /*genNested=*/false, currentLocation,
-                                       /*outerCombined=*/false,
-                                       /*clauseList=*/nullptr,
+  genOpWithBody<mlir::omp::SectionsOp>({.converter = converter,
+                                        .eval = eval,
+                                        .currentLocation = currentLocation},
                                        /*reduction_vars=*/mlir::ValueRange(),
                                        /*reductions=*/nullptr, allocateOperands,
                                        allocatorOperands, nowaitClauseOperand);

[tblah]

LGTM, nice cleanup!

[ergawy]

Thanks for the approval @tblah!

Since this needs C++20, I opened an RFC to upgrade to 20. I will park this PR until we reach a conclusion.

[skatrak]

Thanks for this work Kareem, the argument lists for these functions were becoming quite large. Would it be possible land this patch without any C++20 features, and then add another PR after there's community consensus to make the switch?

This suggestion is just to unblock this patch in case a decision there takes some time, assuming it's not much work to remove C++20 features.

[tblah]

Thanks for this work Kareem, the argument lists for these functions were becoming quite large. Would it be possible land this patch without any C++20 features, and then add another PR after there's community consensus to make the switch?

+1 we could do something like this for the time being 7006b90

[ergawy]

@luporl suggested this refactoring on #79862 and I think it is indeed a good idea. Thanks Leandro!

Undid the member designator usage.

[ergawy]

+1 we could do something like this for the time being 7006b90

I think the PR now aligns with the linked change. Let me know if I missed something.

[skatrak]

I have some small comments, but they can be addressed as part of another PR since this one has already been merged.

[skatrak]

This structure contains the same fields as CreateBodyOfOpInfo except for the args and dsp fields. Maybe it makes sense to have that structure inherit from this one so we're not using different names/types for the same things (e.g. Location &loc vs Location currentLocation).

[skatrak]

Suggested change

	mlir::Location &loc;
	mlir::Location loc;

[skatrak]

Suggested change

	const llvm::SmallVector<const Fortran::semantics::Symbol *> &args = {};
	const llvm::ArrayRef<const Fortran::semantics::Symbol *> args = {};

[luporl]

While we can't use member designators, what about creating setters for the fields that have a default value?

  CreateBodyOfOpInfo &setGenNested(bool val) { genNested = val; }
  CreateBodyOfOpInfo &setClauses(const Fortran::parser::OmpClauseList *clauses) { clauses = clauses; }
  ...

By returning a reference to the object, it should be possible to chain calls to setters, like below:
{ converter, loc, eval }.setClauses(clauses).setDsp(dsp)

The advantage of this is that only fields with non-default values need to be explicitly set.
So, for instance, if a given function just needs to set dsp, it doesn't need to specify all other fields before it.
This makes it easier to add a new field or change default values.

Another option would be to add constructors that handle the most commonly set fields.
Example:

CreateBodyOfOpInfo(Fortran::lower::AbstractConverter &converter,
  mlir::Location &loc,
  Fortran::lower::pft::Evaluation &eval,
  bool genNested = true,
  const Fortran::parser::OmpClauseList *clauses = nullptr,
  const llvm::SmallVector<const Fortran::semantics::Symbol *> &args = {},
  bool outerCombined = false,
  DataSharingProcessor *dsp = nullptr)
  : loc(loc), eval(eval), genNested(genNested), ... {}

CreateBodyOfOpInfo(Fortran::lower::AbstractConverter &converter,
  mlir::Location &loc,
  Fortran::lower::pft::Evaluation &eval,
  bool genNested,
  bool outerCombined)
  : CreateBodyOfOpInfo(converter, loc, eval, genNested, nullptr, {}, outerCombined) {}

CreateBodyOfOpInfo(Fortran::lower::AbstractConverter &converter,
  mlir::Location &loc,
  Fortran::lower::pft::Evaluation &eval,
  DataSharingProcessor *dsp)
  : CreateBodyOfOpInfo(converter, loc, eval, true, nullptr, {}, false, dsp) {}