[flang-commits] [flang] [Flang][OpenMP]Support for lowering task_reduction and in_reduction to MLIR (PR #111155)

[Sergio Afonso via flang-commits]

================
@@ -1926,13 +1942,26 @@ genTaskgroupOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
                const ConstructQueue &queue,
                ConstructQueue::const_iterator item) {
   mlir::omp::TaskgroupOperands clauseOps;
-  genTaskgroupClauses(converter, semaCtx, item->clauses, loc, clauseOps);
+  llvm::SmallVector<const semantics::Symbol *> taskReductionSyms;
+  genTaskgroupClauses(converter, semaCtx, item->clauses, loc, clauseOps,
+                      taskReductionSyms);
 
-  return genOpWithBody<mlir::omp::TaskgroupOp>(
+  OpWithBodyGenInfo genInfo =
       OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
                         llvm::omp::Directive::OMPD_taskgroup)
-          .setClauses(&item->clauses),
-      queue, item, clauseOps);
+          .setClauses(&item->clauses);
+
+  auto taskgroupOp =
+      genOpWithBody<mlir::omp::TaskgroupOp>(genInfo, queue, item, clauseOps);
+
+  llvm::SmallVector<mlir::Type> taskReductionTypes;
+  for (const auto &taskreductionVar : clauseOps.taskReductionVars)
+    taskReductionTypes.push_back(taskreductionVar.getType());
+
+  // Add reduction variables as entry block arguments to the taskgroup region
+  llvm::SmallVector<mlir::Location> blockArgLocs(taskReductionSyms.size(), loc);
+  taskgroupOp->getRegion(0).addArguments(taskReductionTypes, blockArgLocs);
----------------
skatrak wrote:

Please replace this with an approach similar to what's done in `genTaskOp`. The callback to generate the entry block should be the one that introduces these arguments and binds them to their corresponding symbol. Otherwise, the outside values will be used instead by nested operations.

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