[clang][OpenMP] Emit lastprivate final copies in no-loop kernels

[nicebert]

Promotion rejected lastprivate because the no-loop branch leaves the
worksharing path before it privatizes or copies out, so the clause would
have been silently dropped.

Privatize the non-counter variables in the parallel region and copy them
out under the last iteration that applyWorkshareLoop publishes, forcing
the exit barrier the copy reads through. Loop counters stay at the
distribute level and reach EmitOMPSimdFinal.

Split out of #205325.

Depends on: #224040

[llvmorg-github-actions]

@llvm/pr-subscribers-offload

@llvm/pr-subscribers-clang-codegen

Author: Nicole Aschenbrenner (nicebert)

Changes

Promotion rejected lastprivate because the no-loop branch leaves the
worksharing path before it privatizes or copies out, so the clause would
have been silently dropped.

Privatize the non-counter variables in the parallel region and copy them
out under the last iteration that applyWorkshareLoop publishes, forcing
the exit barrier the copy reads through. Loop counters stay at the
distribute level and reach EmitOMPSimdFinal.

Split out of #205325.

Depends on: llvm/llvm-project#224040

---

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

4 Files Affected:

• (modified) clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp (+1-3)
• (modified) clang/lib/CodeGen/CGStmtOpenMP.cpp (+60-8)
• (modified) clang/test/OpenMP/target_no_loop.c (+54-1)
• (modified) offload/test/offloading/target-no-loop.c (+48)

diff --git a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
index 264d78b0f5245..a22ccfcf3ddc2 100644
--- a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
+++ b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
@@ -1158,9 +1158,7 @@ bool CGOpenMPRuntimeGPU::canPromoteToNoLoop(
           DKind == OMPD_target_teams_distribute_parallel_for_simd) &&
          LangOpts.OpenMPTeamSubscription && LangOpts.OpenMPThreadSubscription &&
          !D.hasClausesOfKind<OMPNumTeamsClause>() &&
-         !D.hasClausesOfKind<OMPReductionClause>() &&
-         !D.hasClausesOfKind<OMPLastprivateClause>() &&
-         !D.hasClausesOfKind<OMPLinearClause>();
+         !D.hasClausesOfKind<OMPReductionClause>();
 }
 
 std::pair<llvm::Value *, llvm::Value *>
diff --git a/clang/lib/CodeGen/CGStmtOpenMP.cpp b/clang/lib/CodeGen/CGStmtOpenMP.cpp
index 4d19fa91172a8..4c59c726c66ff 100644
--- a/clang/lib/CodeGen/CGStmtOpenMP.cpp
+++ b/clang/lib/CodeGen/CGStmtOpenMP.cpp
@@ -66,6 +66,29 @@ static bool canEmitGPUFusedDistSchedule(const CodeGenModule &CGM,
          !S.getSingleClause<OMPOrderedClause>();
 }
 
+static bool isLoopCounter(const OMPLoopDirective &S, const Expr *E) {
+  const VarDecl *Canonical =
+      cast<VarDecl>(cast<DeclRefExpr>(E->IgnoreParenImpCasts())->getDecl())
+          ->getCanonicalDecl();
+  for (const Expr *C : S.counters())
+    if (cast<VarDecl>(cast<DeclRefExpr>(C)->getDecl())->getCanonicalDecl() ==
+        Canonical)
+      return true;
+  return false;
+}
+
+static bool
+hasLoopLastprivateVariable(const OMPLoopDirective &S,
+                           llvm::function_ref<bool(const Expr *)> Pred) {
+  for (const auto *C : S.getClausesOfKind<OMPLastprivateClause>()) {
+    for (const Expr *E : C->varlist()) {
+      if (Pred(E))
+        return true;
+    }
+  }
+  return false;
+}
+
 static bool canEmitGPUNoLoopKernel(CodeGenModule &CGM,
                                    const OMPLoopDirective &S) {
   const auto *D = dyn_cast<OMPTargetTeamsDistributeParallelForDirective>(&S);
@@ -3772,6 +3795,10 @@ emitInnerParallelForWhenCombined(CodeGenFunction &CGF,
       CGF.EmitOMPPrivateLoopCounters(S, PrivateScope);
       (void)PrivateScope.Privatize();
 
+      CodeGenFunction::OMPPrivateScope LastprivateScope(CGF);
+      (void)CGF.EmitOMPLastprivateClauseInit(S, LastprivateScope);
+      (void)LastprivateScope.Privatize();
+
       if (isOpenMPTargetExecutionDirective(EKind))
         CGM.getOpenMPRuntime().adjustTargetSpecificDataForLambdas(CGF, S);
 
@@ -3795,17 +3822,33 @@ emitInnerParallelForWhenCombined(CodeGenFunction &CGF,
           CGF.AllocaInsertPt->getParent(), CGF.AllocaInsertPt->getIterator());
       llvm::OpenMPIRBuilder &OMPBuilder =
           CGM.getOpenMPRuntime().getOMPBuilder();
+      bool NeedsLastprivateFinalCopy = hasLoopLastprivateVariable(
+          S, [&S](const Expr *E) { return !isLoopCounter(S, E); });
 
       cantFail(OMPBuilder.applyWorkshareLoop(
           CGF.Builder.getCurrentDebugLocation(), CLI, AllocaIP,
-          /*NeedsBarrier=*/!S.getSingleClause<OMPNowaitClause>(),
+          /*NeedsBarrier=*/!S.getSingleClause<OMPNowaitClause>() ||
+              NeedsLastprivateFinalCopy,
           llvm::omp::OMP_SCHEDULE_Default,
           /*ChunkSize=*/nullptr, /*HasSimdModifier=*/false,
           /*HasMonotonicModifier=*/false, /*HasNonmonotonicModifier=*/false,
           /*HasOrderedClause=*/false,
           llvm::omp::WorksharingLoopType::DistributeForStaticLoop,
           /*NoLoop=*/true, /*HasDistSchedule=*/false,
-          /*DistScheduleChunkSize=*/nullptr));
+          /*DistScheduleChunkSize=*/nullptr,
+          /*NeedsLastIter=*/NeedsLastprivateFinalCopy));
+
+      // Emit final copy for lastprivate variables, excluding counters deferred
+      // to the distribute level.
+      if (NeedsLastprivateFinalCopy) {
+        llvm::Value *LastIter = CLI->getLastIter();
+        assert(LastIter && "workshare loop did not publish last iteration");
+        CGF.EmitOMPLastprivateClauseFinal(
+            S, /*NoFinals=*/true,
+            CGF.Builder.CreateIsNotNull(CGF.Builder.CreateAlignedLoad(
+                CGF.Int32Ty, LastIter, CharUnits::fromQuantity(4),
+                ".omp.is_last")));
+      }
       return;
     }
 
@@ -6364,8 +6407,11 @@ void CodeGenFunction::EmitOMPDistributeLoop(const OMPLoopDirective &S,
           !isOpenMPTeamsDirective(S.getDirectiveKind()))
         EmitOMPReductionClauseInit(S, LoopScope);
 
+      // Defer lastprivate init to the parallel region for no-loop
       const bool NoLoopKernel = canEmitGPUNoLoopKernel(CGM, S);
-      HasLastprivateClause = EmitOMPLastprivateClauseInit(S, LoopScope);
+      if (!NoLoopKernel)
+        HasLastprivateClause = EmitOMPLastprivateClauseInit(S, LoopScope);
+
       EmitOMPPrivateLoopCounters(S, LoopScope);
       (void)LoopScope.Privatize();
       if (isOpenMPTargetExecutionDirective(S.getDirectiveKind()))
@@ -6501,11 +6547,17 @@ void CodeGenFunction::EmitOMPDistributeLoop(const OMPLoopDirective &S,
                                      CodeGenLoop);
         }
       }
-      if (isOpenMPSimdDirective(S.getDirectiveKind())) {
-        EmitOMPSimdFinal(S, [IL, &S](CodeGenFunction &CGF) {
-          return CGF.Builder.CreateIsNotNull(
-              CGF.EmitLoadOfScalar(IL, S.getBeginLoc()));
-        });
+      if (isOpenMPSimdDirective(S.getDirectiveKind()) ||
+          (NoLoopKernel && hasLoopLastprivateVariable(S, [&S](const Expr *E) {
+             return isLoopCounter(S, E);
+           }))) {
+        EmitOMPSimdFinal(
+            S, [IL, &S, NoLoopKernel](CodeGenFunction &CGF) -> llvm::Value * {
+              if (NoLoopKernel)
+                return nullptr;
+              return CGF.Builder.CreateIsNotNull(
+                  CGF.EmitLoadOfScalar(IL, S.getBeginLoc()));
+            });
       }
       if (isOpenMPSimdDirective(S.getDirectiveKind()) &&
           !isOpenMPParallelDirective(S.getDirectiveKind()) &&
diff --git a/clang/test/OpenMP/target_no_loop.c b/clang/test/OpenMP/target_no_loop.c
index 0fe3de1780cce..d754331f3f87b 100644
--- a/clang/test/OpenMP/target_no_loop.c
+++ b/clang/test/OpenMP/target_no_loop.c
@@ -45,12 +45,37 @@ void no_loop_simd(int *array) {
     array[i] = i + 1;
 }
 
+void no_loop_lastprivate_counter(int *array) {
+  int i;
+#pragma omp target teams distribute parallel for lastprivate(i)
+  for (i = 0; i < 1024; ++i)
+    array[i] = i + 1;
+}
+
+void no_loop_lastprivate_scalar(int *array) {
+  int last = 0;
+#pragma omp target teams distribute parallel for lastprivate(last)
+  for (int i = 0; i < 1024; ++i) {
+    array[i] = i + 1;
+    last = i;
+  }
+}
+
 void no_loop_nowait(int *array) {
 #pragma omp target teams distribute parallel for nowait
   for (int i = 0; i < 1024; ++i)
     array[i] = i + 1;
 }
 
+void no_loop_lastprivate_scalar_nowait(int *array) {
+  int last = 0;
+#pragma omp target teams distribute parallel for lastprivate(last) nowait
+  for (int i = 0; i < 1024; ++i) {
+    array[i] = i + 1;
+    last = i;
+  }
+}
+
 // NOLOOP: no_loop_l{{[0-9]+}}_kernel_environment {{.*}} i8 0, i8 1, i8 6
 // NOLOOP: no_loop_simd_l{{[0-9]+}}_kernel_environment {{.*}} i8 0, i8 1, i8 6
 
@@ -71,6 +96,25 @@ void no_loop_nowait(int *array) {
 // NOLOOP: omp_loop.after:
 // NOLOOP-NEXT: ret void
 
+// NOLOOP-LABEL: @__kmpc_parallel_60({{.*}}lastprivate_counter{{.*}})
+// NOLOOP: omp.loop.exit:
+// NOLOOP-NEXT: store i32 1024, ptr %i
+// NOLOOP-NEXT: @__kmpc_free_shared(ptr %i{{.*}})
+// NOLOOP-NEXT: ret void
+// NOLOOP: @__kmpc_distribute_for_static_loop_4u({{.*}}lastprivate_counter{{.*}}, i32 0, i32 0, i8 1)
+// NOLOOP: @__kmpc_barrier
+// NOLOOP: omp_loop.after:
+// NOLOOP-NEXT: ret void
+
+// NOLOOP-LABEL: @__kmpc_parallel_60({{.*}}lastprivate_scalar{{.*}})
+// NOLOOP: omp.loop.exit:
+// NOLOOP-NEXT: @__kmpc_free_shared(ptr %last{{.*}})
+// NOLOOP-NEXT: ret void
+// NOLOOP: @__kmpc_distribute_for_static_loop_4u({{.*}}lastprivate_scalar{{.*}}, i32 0, i32 0, i8 1)
+// NOLOOP: @__kmpc_barrier
+// NOLOOP: store {{.*}}, ptr %last.
+// NOLOOP-NEXT: %.omp.lastprivate.done
+
 // NOLOOP-LABEL: @__kmpc_parallel_60({{.*}}no_loop_nowait{{.*}})
 // NOLOOP: omp.loop.exit:
 // NOLOOP-NEXT: ret void
@@ -80,4 +124,13 @@ void no_loop_nowait(int *array) {
 // NOLOOP: omp_loop.after:
 // NOLOOP-NEXT: ret void
 
-// SPMD-COUNT-3: _kernel_environment {{.*}} i8 0, i8 1, i8 2
+// NOLOOP-LABEL: @__kmpc_parallel_60({{.*}}lastprivate_scalar_nowait{{.*}})
+// NOLOOP: omp.loop.exit:
+// NOLOOP-NEXT: @__kmpc_free_shared(ptr %last{{.*}})
+// NOLOOP-NEXT: ret void
+// NOLOOP: @__kmpc_distribute_for_static_loop_4u({{.*}}lastprivate_scalar_nowait{{.*}}, i32 0, i32 0, i8 1)
+// NOLOOP: @__kmpc_barrier
+// NOLOOP: store {{.*}}, ptr %last.
+// NOLOOP-NEXT: %.omp.lastprivate.done
+
+// SPMD-COUNT-6: _kernel_environment {{.*}} i8 0, i8 1, i8 2
diff --git a/offload/test/offloading/target-no-loop.c b/offload/test/offloading/target-no-loop.c
index 59558c67152e2..ae6c31c966b19 100644
--- a/offload/test/offloading/target-no-loop.c
+++ b/offload/test/offloading/target-no-loop.c
@@ -72,6 +72,48 @@ int main(void) {
   if (red != 1024)
     ++errors;
 
+  // No-loop kernel with a lastprivate variable (final value from thread executing last iteration)
+  for (int i = 0; i < 1024; ++i)
+    array[i] = 1;
+
+  int last = -1;
+#pragma omp target teams distribute parallel for lastprivate(last)
+  for (int i = 0; i < 1024; ++i) {
+    array[i] = i + 1;
+    last = i;
+  }
+  errors += check_errors(array);
+#ifdef ASSERT_LASTPRIVATE
+  if (last != 1023)
+    ++errors;
+#endif
+
+  // No-loop kernel with a lastprivate loop counter (final value from evaluation at distribution level)
+  for (int i = 0; i < 1024; ++i)
+    array[i] = 1;
+  int iv = -1;
+#pragma omp target teams distribute parallel for lastprivate(iv)
+  for (iv = 0; iv < 1024; ++iv)
+    array[iv] = iv + 1;
+  errors += check_errors(array);
+#ifdef ASSERT_LASTPRIVATE
+  if (iv != 1024)
+    ++errors;
+#endif
+
+  // No-loop simd kernel with a lastprivate loop counter
+  for (int i = 0; i < 1024; ++i)
+    array[i] = 1;
+  int simd_iv = -1;
+#pragma omp target teams distribute parallel for simd lastprivate(simd_iv)
+  for (simd_iv = 0; simd_iv < 1024; ++simd_iv)
+    array[simd_iv] = simd_iv + 1;
+  errors += check_errors(array);
+#ifdef ASSERT_LASTPRIVATE
+  if (simd_iv != 1024)
+    ++errors;
+#endif
+
   printf("number of errors: %d\n", errors);
   return 0;
 }
@@ -88,4 +130,10 @@ int main(void) {
 // CHECK:  info: #Args: 2 Teams x Thrds:   16x  16 {{.*}}
 // CHECK:  PluginInterface device {{[0-9]+}} info: Launching kernel {{.*}} SPMD mode
 // CHECK:  info: #Args: 3 Teams x Thrds:   16x  16 {{.*}}
+// CHECK:  PluginInterface device {{[0-9]+}} info: Launching kernel {{.*}} SPMD-No-Loop mode
+// CHECK:  info: #Args: 3 Teams x Thrds:   64x  16 {{.*}}
+// CHECK:  PluginInterface device {{[0-9]+}} info: Launching kernel {{.*}} SPMD-No-Loop mode
+// CHECK:  info: #Args: 3 Teams x Thrds:   64x  16 {{.*}}
+// CHECK:  PluginInterface device {{[0-9]+}} info: Launching kernel {{.*}} SPMD-No-Loop mode
+// CHECK:  info: #Args: 3 Teams x Thrds:   64x  16 {{.*}}
 // CHECK:  number of errors: 0