[flang-commits] [flang] [Flang][OpenMP] Support conditional lastprivate on GPU targets (PR #224380)

[Sairudra More via flang-commits]

================
@@ -4982,6 +5097,90 @@ getOrCreateConditionalLpGlobal(lower::AbstractConverter &converter,
                                global.getSymbol());
 }
 
+/// Return the innermost enclosing omp.target, or a null op if the current
+/// insertion point is not inside a target region (i.e. host lowering).
+static mlir::omp::TargetOp findEnclosingTargetOp(fir::FirOpBuilder &builder) {
+  for (auto *op = builder.getInsertionBlock()->getParentOp(); op;
+       op = op->getParentOp()) {
+    if (auto targetOp = mlir::dyn_cast<mlir::omp::TargetOp>(op))
+      return targetOp;
+  }
+  return {};
+}
+
+/// Return the innermost enclosing omp.teams, or a null op if there is none.
+static mlir::omp::TeamsOp findEnclosingTeamsOp(fir::FirOpBuilder &builder) {
+  for (auto *op = builder.getInsertionBlock()->getParentOp(); op;
+       op = op->getParentOp()) {
+    if (auto teamsOp = mlir::dyn_cast<mlir::omp::TeamsOp>(op))
+      return teamsOp;
+  }
+  return {};
+}
+
+/// Host placement of the conditional-lastprivate reduction struct for a
+/// worksharing loop.  Create the struct alloca OUTSIDE the parent omp.parallel
+/// (if any), so the reduction result persists after the parallel region ends.
+/// In the orphaned case (no enclosing ParallelOp), use a module-scope global
+/// so that all threads share one reduction target.
+static void placeConditionalLpStructOnHost(lower::AbstractConverter &converter,
+                                           mlir::Location loc,
+                                           fir::RecordType lpType,
+                                           mlir::Value &lpAlloca) {
+  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+  mlir::omp::ParallelOp enclosingParallel = findEnclosingParallelOp(builder);
+  bool isOrphaned = !enclosingParallel;
+
+  // Guard against nested parallelism in the orphaned case.
+  // Emit this BEFORE touching the global to avoid racing on it.
+  if (isOrphaned)
+    emitNestedParallelGuardForCondLp(converter, loc);
+
+  if (enclosingParallel) {
+    mlir::OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPoint(enclosingParallel);
+    lpAlloca = builder.createTemporary(loc, lpType);
+    // Index fields are initialised to -1 so the combiner's "sequentially
+    // last" comparison treats them as "no iteration has written yet"
+    // (any real canonical loop IV >= 0 beats -1).
+    initConditionalLpStructDefault(builder, loc, lpType, lpAlloca);
+  } else {
+    lpAlloca = getOrCreateConditionalLpGlobal(converter, loc, lpType);
+    // The global is shared across all threads. Use omp.single (which
+    // has an implicit barrier at exit) so that exactly one thread
+    // initialises and all threads wait before entering the construct.
+    mlir::omp::SingleOperands initSingleOps;
+    auto singleOp = mlir::omp::SingleOp::create(builder, loc, initSingleOps);
+    mlir::Block *singleBlock = builder.createBlock(&singleOp.getRegion());
+    builder.setInsertionPointToStart(singleBlock);
+    initConditionalLpStructDefault(builder, loc, lpType, lpAlloca);
+    mlir::omp::TerminatorOp::create(builder, loc);
+    builder.setInsertionPointAfter(singleOp);
+  }
+}
+
+/// Device (GPU offload) placement of the conditional-lastprivate reduction
+/// struct for a worksharing loop nested inside an omp.target region.
+///
+/// Each team runs an independent wsloop reduction, so the struct is allocated
+/// at the start of the enclosing omp.teams body (or the omp.target body when
+/// there is no teams) to give every team its own reduction target rather than
+/// racing on a single shared alloca.  The copy-back is likewise deferred until
+/// after the parallel region (see genStandaloneDo).
+static void placeConditionalLpStructInTarget(fir::FirOpBuilder &builder,
+                                             mlir::Location loc,
+                                             fir::RecordType lpType,
+                                             mlir::omp::TargetOp targetOp,
+                                             mlir::Value &lpAlloca) {
+  mlir::OpBuilder::InsertionGuard guard(builder);
+  if (mlir::omp::TeamsOp teamsOp = findEnclosingTeamsOp(builder))
+    builder.setInsertionPointToStart(&teamsOp.getRegion().front());
+  else
+    builder.setInsertionPointToStart(&targetOp.getRegion().front());
+  lpAlloca = builder.createTemporary(loc, lpType);
+  initConditionalLpStructDefault(builder, loc, lpType, lpAlloca);
----------------
Saieiei wrote:

Could we reset this state for each dynamic encounter of the worksharing loop?

On b37b2eb8e401, I can reproduce this with a `target if(.false.)` containing `do j = 1, 2` around the same `parallel do lastprivate(conditional:x)`, assigning `x = j` at the same final iteration each time. The result is `1, 1` instead of `1, 2`.

It looks like the accumulator initialized here survives both encounters, so the equal-index winner from the second encounter does not replace the first. A repeated-encounter regression would cover this.


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