[flang-commits] [flang] [flang][OpenMP] Lower REDUCTION clause for SECTIONS (PR #97858)

[Sergio Afonso via flang-commits]

================
@@ -1530,11 +1525,52 @@ genSectionsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
   }
 
   // SECTIONS construct.
-  mlir::omp::SectionsOp sectionsOp = genOpWithBody<mlir::omp::SectionsOp>(
-      OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
-                        llvm::omp::Directive::OMPD_sections)
-          .setClauses(&nonDsaClauses),
-      queue, item, clauseOps);
+  auto sectionsOp = builder.create<mlir::omp::SectionsOp>(loc, clauseOps);
+
+  auto reductionCallback = [&](mlir::Operation *op) {
+    genReductionVars(op, converter, loc, reductionSyms, reductionTypes);
+    return reductionSyms;
+  };
+
+  reductionCallback(sectionsOp);
+  // genReductionVars adds a hlfir.declare for the reduction block argument
+  // but only terminators and sectionOps are allowed inside of a SectionsOp
+  llvm::SmallVector<mlir::Operation *> toErase;
+  toErase.reserve(reductionSyms.size());
+  for (auto decl : sectionsOp.getOps<hlfir::DeclareOp>())
+    toErase.push_back(decl);
+  for (mlir::Operation *op : toErase)
+    op->erase();
+
+  mlir::Operation *terminator =
+      lower::genOpenMPTerminator(builder, sectionsOp, loc);
+
+  // Generate nested SECTION constructs.
+  // This is done here rather than in genOMP([...], OpenMPSectionConstruct )
+  // because we need to run genReductionVars on each omp.section so that the
+  // reduction variable gets mapped to the private version
+  for (auto [construct, nestedEval] :
+       llvm::zip(sectionBlocks.v, eval.getNestedEvaluations())) {
+    const auto *sectionConstruct =
+        std::get_if<parser::OpenMPSectionConstruct>(&construct.u);
+    if (!sectionConstruct) {
+      assert(false &&
+             "unexpected construct nested inside of SECTIONS construct");
+      continue;
+    }
+
+    ConstructQueue sectionQueue{buildConstructQueue(
+        converter.getFirOpBuilder().getModule(), semaCtx, nestedEval,
+        sectionConstruct->source, llvm::omp::Directive::OMPD_section, {})};
+
+    builder.setInsertionPoint(terminator);
+    genOpWithBody<mlir::omp::SectionOp>(
+        OpWithBodyGenInfo(converter, symTable, semaCtx, loc, nestedEval,
+                          llvm::omp::Directive::OMPD_section)
+            .setClauses(&sectionQueue.begin()->clauses)
+            .setGenRegionEntryCb(reductionCallback),
----------------
skatrak wrote:

Wouldn't the `omp.sections` operation be the only one in this case that should have reduction information? `omp.section` doesn't have a reduction clause, which matches the spec as far as I can tell.

My understanding is that what we need to achieve here is to generate the `hlfir.declare` operations that would have been created in the parent `omp.sections`, but had to be removed. Maybe the best way to have reduction-related block arguments in `omp.sections` and the corresponding `hlfir.declare` operations in each nested `omp.section` is to avoid using the `genReductionVars` functions and split its contents instead. Something like:
```c++
auto sectionsOp = builder.create<mlir::omp::SectionsOp>(loc, clauseOps);
// Create entry block with reduction variables as arguments.
llvm::SmallVector<mlir::Location> blockArgLocs(reductionSyms.size(), loc);
mlir::Block *entryBlock = firOpBuilder.createBlock(&sectionsOp->getRegion(0), {}, reductionTypes, blockArgLocs);
// No need to remove hlfir.declare operations anymore.
mlir::Operation *terminator = lower::genOpenMPTerminator(builder, sectionsOp, loc);

auto reductionCallback = [&](mlir::Operation *op) {
  // Bind the reduction arguments to their block arguments.
  for (auto [arg, prv] :
       llvm::zip_equal(reductionSyms, entryBlock->getArguments())) {
    converter.bindSymbol(*arg, prv);
  }
  return {};
};

for ... {
  // ...
  genOpWithBody<mlir::omp::SectionOp>(...);
}
```
This way no block arguments are added to the `omp.section` operations and there's no need to remove `hlfir.declare` operations inside of `omp.sections`. Maybe this alternative introduces problems I'm not aware of, so let me know.

https://github.com/llvm/llvm-project/pull/97858