[flang][OpenMP][NFC] Further refactoring for genOpWithBody &

[ergawy]

createBodyOfOp

This refactors the arguments to the above functions in 2 ways:

• Combines the 2 structs of arguments into one since they were almost identical.
• Replaces the args argument with a callback to a rebion-body generation function. This is a preparation for delayed privatization as we will need different callbacks for ws loops and parallel ops with delayed privatization.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

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

Author: Kareem Ergawy (ergawy)

Changes

createBodyOfOp

This refactors the arguments to the above functions in 2 ways:

• Combines the 2 structs of arguments into one since they were almost identical.
• Replaces the args argument with a callback to a rebion-body generation function. This is a preparation for delayed privatization as we will need different callbacks for ws loops and parallel ops with delayed privatization.

---

Patch is 22.27 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/80839.diff

1 Files Affected:

• (modified) flang/lib/Lower/OpenMP.cpp (+168-101)

diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index dad88fc1d764c..9ebb50c46bfde 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2250,31 +2250,68 @@ createAndSetPrivatizedLoopVar(Fortran::lower::AbstractConverter &converter,
   return storeOp;
 }
 
-struct CreateBodyOfOpInfo {
+struct OpWithBodyGenInfo {
+  /// A type for a code-gen callback function. This takes as argument the op for
+  /// which the code is being generated and returns the arguments of the op's
+  /// region.
+  using GenOMPRegionEntryCBFn =
+      std::function<llvm::SmallVector<const Fortran::semantics::Symbol *>(
+          mlir::Operation *)>;
+
+  OpWithBodyGenInfo(Fortran::lower::AbstractConverter &converter,
+                    mlir::Location loc, Fortran::lower::pft::Evaluation &eval)
+      : converter(converter), loc(loc), eval(eval) {}
+
+  OpWithBodyGenInfo &setGenNested(bool value) {
+    genNested = value;
+    return *this;
+  }
+
+  OpWithBodyGenInfo &setOuterCombined(bool value) {
+    outerCombined = value;
+    return *this;
+  }
+
+  OpWithBodyGenInfo &setClauses(const Fortran::parser::OmpClauseList *value) {
+    clauses = value;
+    return *this;
+  }
+
+  OpWithBodyGenInfo &setDataSharingProcessor(DataSharingProcessor *value) {
+    dsp = value;
+    return *this;
+  }
+
+  OpWithBodyGenInfo &setGenRegionEntryCb(GenOMPRegionEntryCBFn value) {
+    genRegionEntryCB = value;
+    return *this;
+  }
+
+  /// [inout] converter to use for the clauses.
   Fortran::lower::AbstractConverter &converter;
-  mlir::Location &loc;
+  /// [in] location in source code.
+  mlir::Location loc;
+  /// [in] current PFT node/evaluation.
   Fortran::lower::pft::Evaluation &eval;
+  /// [in] whether to generate FIR for nested evaluations
   bool genNested = true;
-  const Fortran::parser::OmpClauseList *clauses = nullptr;
-  const llvm::SmallVector<const Fortran::semantics::Symbol *> &args = {};
+  /// [in] is this an outer operation - prevents privatization.
   bool outerCombined = false;
+  /// [in] list of clauses to process.
+  const Fortran::parser::OmpClauseList *clauses = nullptr;
+  /// [in] if provided, processes the construct's data-sharing attributes.
   DataSharingProcessor *dsp = nullptr;
+  /// [in] if provided, emits the op's region entry. Otherwise, an emtpy block
+  /// is created in the region.
+  GenOMPRegionEntryCBFn genRegionEntryCB = nullptr;
 };
 
 /// Create the body (block) for an OpenMP Operation.
 ///
-/// \param [in]    op - the operation the body belongs to.
-/// \param [inout] converter - converter to use for the clauses.
-/// \param [in]    loc - location in source code.
-/// \param [in]    eval - current PFT node/evaluation.
-/// \param [in]    genNested - whether to generate FIR for nested evaluations
-/// \oaran [in]    clauses - list of clauses to process.
-/// \param [in]    args - block arguments (induction variable[s]) for the
-////                      region.
-/// \param [in]    outerCombined - is this an outer operation - prevents
-///                                privatization.
+/// \param [in]   op - the operation the body belongs to.
+/// \param [in] info - options controlling code-gen for the construction.
 template <typename Op>
-static void createBodyOfOp(Op &op, CreateBodyOfOpInfo info) {
+static void createBodyOfOp(Op &op, OpWithBodyGenInfo info) {
   fir::FirOpBuilder &firOpBuilder = info.converter.getFirOpBuilder();
 
   auto insertMarker = [](fir::FirOpBuilder &builder) {
@@ -2287,28 +2324,15 @@ static void createBodyOfOp(Op &op, CreateBodyOfOpInfo info) {
   // 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 (info.args.size()) {
-    std::size_t loopVarTypeSize = 0;
-    for (const Fortran::semantics::Symbol *arg : info.args)
-      loopVarTypeSize = std::max(loopVarTypeSize, arg->GetUltimate().size());
-    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(info.args)) {
-      mlir::Value indexVal =
-          fir::getBase(op.getRegion().front().getArgument(argIndex));
-      storeOp = createAndSetPrivatizedLoopVar(info.converter, info.loc,
-                                              indexVal, argSymbol);
+  auto regionArgs =
+      [&]() -> llvm::SmallVector<const Fortran::semantics::Symbol *> {
+    if (info.genRegionEntryCB != nullptr) {
+      return info.genRegionEntryCB(op);
     }
-    firOpBuilder.setInsertionPointAfter(storeOp);
-  } else {
-    firOpBuilder.createBlock(&op.getRegion());
-  }
 
+    firOpBuilder.createBlock(&op.getRegion());
+    return {};
+  }();
   // Mark the earliest insertion point.
   mlir::Operation *marker = insertMarker(firOpBuilder);
 
@@ -2363,7 +2387,8 @@ static void createBodyOfOp(Op &op, CreateBodyOfOpInfo info) {
     // is more complicated, especially with unstructured control flow, there
     // may be multiple blocks, and some of them may have non-OMP terminators
     // resulting from lowering of the code contained within the operation.
-    // All the remaining blocks are potential exit points from the op's region.
+    // All the remaining blocks are potential exit points from the op's
+    // region.
     //
     // Explicit control flow cannot exit any OpenMP region (other than via
     // STOP), and that is enforced by semantic checks prior to lowering. STOP
@@ -2405,8 +2430,8 @@ static void createBodyOfOp(Op &op, CreateBodyOfOpInfo info) {
         assert(tempDsp.has_value());
         tempDsp->processStep2(op, isLoop);
       } else {
-        if (isLoop && info.args.size() > 0)
-          info.dsp->setLoopIV(info.converter.getSymbolAddress(*info.args[0]));
+        if (isLoop && regionArgs.size() > 0)
+          info.dsp->setLoopIV(info.converter.getSymbolAddress(*regionArgs[0]));
         info.dsp->processStep2(op, isLoop);
       }
     }
@@ -2481,24 +2506,11 @@ 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(GenOpWithBodyInfo info, Args &&...args) {
+static OpTy genOpWithBody(OpWithBodyGenInfo info, Args &&...args) {
   auto op = info.converter.getFirOpBuilder().create<OpTy>(
-      info.currentLocation, std::forward<Args>(args)...);
-  createBodyOfOp<OpTy>(
-      op, {info.converter, info.currentLocation, info.eval, info.genNested,
-           info.clauseList,
-           /*args*/ llvm::SmallVector<const Fortran::semantics::Symbol *>{},
-           info.outerCombined});
+      info.loc, std::forward<Args>(args)...);
+  createBodyOfOp<OpTy>(op, info);
   return op;
 }
 
@@ -2507,7 +2519,8 @@ genMasterOp(Fortran::lower::AbstractConverter &converter,
             Fortran::lower::pft::Evaluation &eval, bool genNested,
             mlir::Location currentLocation) {
   return genOpWithBody<mlir::omp::MasterOp>(
-      {converter, eval, genNested, currentLocation},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested),
       /*resultTypes=*/mlir::TypeRange());
 }
 
@@ -2516,7 +2529,8 @@ genOrderedRegionOp(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval, bool genNested,
                    mlir::Location currentLocation) {
   return genOpWithBody<mlir::omp::OrderedRegionOp>(
-      {converter, eval, genNested, currentLocation},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested),
       /*simd=*/false);
 }
 
@@ -2544,7 +2558,10 @@ genParallelOp(Fortran::lower::AbstractConverter &converter,
     cp.processReduction(currentLocation, reductionVars, reductionDeclSymbols);
 
   return genOpWithBody<mlir::omp::ParallelOp>(
-      {converter, eval, genNested, currentLocation, outerCombined, &clauseList},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setOuterCombined(outerCombined)
+          .setClauses(&clauseList),
       /*resultTypes=*/mlir::TypeRange(), ifClauseOperand,
       numThreadsClauseOperand, allocateOperands, allocatorOperands,
       reductionVars,
@@ -2563,8 +2580,9 @@ 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});
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setClauses(&sectionsClauseList));
 }
 
 static mlir::omp::SingleOp
@@ -2584,8 +2602,9 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
   ClauseProcessor(converter, endClauseList).processNowait(nowaitAttr);
 
   return genOpWithBody<mlir::omp::SingleOp>(
-      {converter, eval, genNested, currentLocation,
-       /*outerCombined=*/false, &beginClauseList},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setClauses(&beginClauseList),
       allocateOperands, allocatorOperands, nowaitAttr);
 }
 
@@ -2617,8 +2636,9 @@ genTaskOp(Fortran::lower::AbstractConverter &converter,
       currentLocation, llvm::omp::Directive::OMPD_task);
 
   return genOpWithBody<mlir::omp::TaskOp>(
-      {converter, eval, genNested, currentLocation,
-       /*outerCombined=*/false, &clauseList},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setClauses(&clauseList),
       ifClauseOperand, finalClauseOperand, untiedAttr, mergeableAttr,
       /*in_reduction_vars=*/mlir::ValueRange(),
       /*in_reductions=*/nullptr, priorityClauseOperand,
@@ -2640,8 +2660,9 @@ 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},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setClauses(&clauseList),
       /*task_reduction_vars=*/mlir::ValueRange(),
       /*task_reductions=*/nullptr, allocateOperands, allocatorOperands);
 }
@@ -2732,8 +2753,9 @@ genEnterExitUpdateDataOp(Fortran::lower::AbstractConverter &converter,
                                    deviceOperand, nowaitAttr, mapOperands);
 }
 
-// This functions creates a block for the body of the targetOp's region. It adds
-// all the symbols present in mapSymbols as block arguments to this block.
+// This functions creates a block for the body of the targetOp's region. It
+// adds all the symbols present in mapSymbols as block arguments to this
+// block.
 static void genBodyOfTargetOp(
     Fortran::lower::AbstractConverter &converter,
     Fortran::lower::pft::Evaluation &eval, bool genNested,
@@ -2750,7 +2772,8 @@ static void genBodyOfTargetOp(
   auto *regionBlock =
       firOpBuilder.createBlock(&region, {}, mapSymTypes, mapSymLocs);
 
-  // Clones the `bounds` placing them inside the target region and returns them.
+  // Clones the `bounds` placing them inside the target region and returns
+  // them.
   auto cloneBound = [&](mlir::Value bound) {
     if (mlir::isMemoryEffectFree(bound.getDefiningOp())) {
       mlir::Operation *clonedOp = bound.getDefiningOp()->clone();
@@ -2812,10 +2835,11 @@ static void genBodyOfTargetOp(
         });
   }
 
-  // Check if cloning the bounds introduced any dependency on the outer region.
-  // If so, then either clone them as well if they are MemoryEffectFree, or else
-  // copy them to a new temporary and add them to the map and block_argument
-  // lists and replace their uses with the new temporary.
+  // Check if cloning the bounds introduced any dependency on the outer
+  // region. If so, then either clone them as well if they are
+  // MemoryEffectFree, or else copy them to a new temporary and add them to
+  // the map and block_argument lists and replace their uses with the new
+  // temporary.
   llvm::SetVector<mlir::Value> valuesDefinedAbove;
   mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
   while (!valuesDefinedAbove.empty()) {
@@ -3024,7 +3048,10 @@ genTeamsOp(Fortran::lower::AbstractConverter &converter,
       currentLocation, llvm::omp::Directive::OMPD_teams);
 
   return genOpWithBody<mlir::omp::TeamsOp>(
-      {converter, eval, genNested, currentLocation, outerCombined, &clauseList},
+      OpWithBodyGenInfo(converter, currentLocation, eval)
+          .setGenNested(genNested)
+          .setOuterCombined(outerCombined)
+          .setClauses(&clauseList),
       /*num_teams_lower=*/nullptr, numTeamsClauseOperand, ifClauseOperand,
       threadLimitClauseOperand, allocateOperands, allocatorOperands,
       reductionVars,
@@ -3221,6 +3248,33 @@ static void convertLoopBounds(Fortran::lower::AbstractConverter &converter,
   }
 }
 
+static llvm::SmallVector<const Fortran::semantics::Symbol *> genCodeForIterVar(
+    mlir::Operation *op, Fortran::lower::AbstractConverter &converter,
+    mlir::Location &loc,
+    const llvm::SmallVector<const Fortran::semantics::Symbol *> &args) {
+  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+  auto &region = op->getRegion(0);
+
+  std::size_t loopVarTypeSize = 0;
+  for (const Fortran::semantics::Symbol *arg : 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);
+  firOpBuilder.createBlock(&region, {}, 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)) {
+    mlir::Value indexVal = fir::getBase(region.front().getArgument(argIndex));
+    storeOp =
+        createAndSetPrivatizedLoopVar(converter, loc, indexVal, argSymbol);
+  }
+  firOpBuilder.setInsertionPointAfter(storeOp);
+
+  return args;
+}
+
 static void
 createSimdLoop(Fortran::lower::AbstractConverter &converter,
                Fortran::lower::pft::Evaluation &eval,
@@ -3268,10 +3322,16 @@ createSimdLoop(Fortran::lower::AbstractConverter &converter,
 
   auto *nestedEval = getCollapsedLoopEval(
       eval, Fortran::lower::getCollapseValue(loopOpClauseList));
+
+  auto ivCallback = [&](mlir::Operation *op) {
+    return genCodeForIterVar(op, converter, loc, iv);
+  };
+
   createBodyOfOp<mlir::omp::SimdLoopOp>(
-      simdLoopOp, {converter, loc, *nestedEval,
-                   /*genNested=*/true, &loopOpClauseList, iv,
-                   /*outerCombined=*/false, &dsp});
+      simdLoopOp, OpWithBodyGenInfo(converter, loc, *nestedEval)
+                      .setClauses(&loopOpClauseList)
+                      .setDataSharingProcessor(&dsp)
+                      .setGenRegionEntryCb(ivCallback));
 }
 
 static void createWsLoop(Fortran::lower::AbstractConverter &converter,
@@ -3344,10 +3404,16 @@ 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,
-                                       /*outerCombined=*/false, &dsp});
+
+  auto ivCallback = [&](mlir::Operation *op) {
+    return genCodeForIterVar(op, converter, loc, iv);
+  };
+
+  createBodyOfOp<mlir::omp::WsLoopOp>(
+      wsLoopOp, OpWithBodyGenInfo(converter, loc, *nestedEval)
+                    .setClauses(&beginClauseList)
+                    .setDataSharingProcessor(&dsp)
+                    .setGenRegionEntryCb(ivCallback));
 }
 
 static void createSimdWsLoop(
@@ -3366,8 +3432,8 @@ static void createSimdWsLoop(
   // OpenMP standard does not specify the length of vector instructions.
   // Currently we safely assume that for !$omp do simd pragma the SIMD length
   // is equal to 1 (i.e. we generate standard workshare loop).
-  // When support for vectorization is enabled, then we need to add handling of
-  // if clause. Currently if clause can be skipped because we always assume
+  // When support for vectorization is enabled, then we need to add handling
+  // of if clause. Currently if clause can be skipped because we always assume
   // SIMD length = 1.
   createWsLoop(converter, eval, ompDirective, beginClauseList, endClauseList,
                loc);
@@ -3628,8 +3694,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,
         currentLocation, mlir::FlatSymbolRefAttr::get(firOpBuilder.getContext(),
                                                       global.getSymName()));
   }();
-  createBodyOfOp<mlir::omp::CriticalOp>(criticalOp,
-                                        {converter, currentLocation, eval});
+  createBodyOfOp<mlir::omp::CriticalOp>(
+      criticalOp, OpWithBodyGenInfo(converter, currentLocation, eval));
 }
 
 static void
@@ -3671,11 +3737,11 @@ genOMP(Fortran::lower::AbstractConverter &converter,
   }
 
   // SECTIONS construct
-  genOpWithBody<mlir::omp::SectionsOp>({converter, eval,
-                                        /*genNested=*/false, currentLocation},
-                                       /*reduction_vars=*/mlir::ValueRange(),
-                                       /*reductions=*/nullptr, allocateOperands,
-                                       allocatorOperands, nowaitClauseOperand);
+  genOpWithBody<mlir::omp::SectionsOp>(
+      OpWithBodyGenInfo(converter, currentLocation, eval).setGenNested(false),
+      /*reduction_vars=*/mlir::ValueRange(),
+      /*reductions=*/nullptr, allocateOperands, allocatorOperands,
+      nowaitClauseOperand);
 
   const auto &sectionBlocks =
       std::get<Fortran::parser::OmpSectionBlocks>(sectionsConstruct.t);
@@ -3775,8 +3841,8 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
 
     // The function or global already has a declare target applied to it, very
     // likely through implicit capture (usage in another declare target
-    // function/subroutine). It should be marked as any if it has been assigned
-    // both host and nohost, else we skip, as there is no change
+    // function/subroutine). It should be marked as any if it has been
+    // assigned both host and nohost, else we skip, as there is no change
     if (declareTargetOp.isDeclareTarget()) {
       if (declareTargetOp.getDeclareTargetDeviceType() != deviceType)
         declareTargetOp.setDeclareTarget(
@@ -3955,8 +4021,8 @@ void Fortran::lower::genThreadprivateOp(
           Fortran::semantics::FindCommonBlockContaining(sym.GetUltimate())) {
     mlir::Value commonValue = converter.getSymbolAddress(*common);
     if (mlir::isa<mlir::omp::ThreadprivateOp>(commonValue.getDefiningOp())) {
-      // Generate ThreadprivateOp for a common block instead of its members and
-      // only do it once for a common block.
+      // Generate ThreadprivateOp for a common block instead of its members
+      // and only do it once for a common block.
       return;
     }
     // Generate ThreadprivateOp and rebind the common block.
@@ -3969,8 +4035,8 @@ void Fortran::l...
[truncated]

[skatrak]

Thanks, this mostly LGTM. I just have a few nits regarding function names and passing the new structure by copy to functions. If there's a reason for the second I'd be curious to know why this is better in this case.

I see a few comment blocks that have been reformatted which doesn't seem like they had to. It adds a bit of noise to the diff, but not sure if that's something that needs addressing.

[ergawy]

If there's a reason for the second I'd be curious to know why this is better in this case.

No particular reason. Passing them by ref now.

I see a few comment blocks that have been reformatted which doesn't seem like they had to. It adds a bit of noise to the diff, but not sure if that's something that needs addressing.

Indeed, but I just ran clang format as I usually do. Is it possible some wrong formatting was merged recently?
Undid the unrelated formatting changes. I have no idea how they happened.

[luporl]

LGTM. Thanks for the refactoring!

[skatrak]

LGTM, thanks!