[Mlir-commits] [clang] [llvm] [mlir] [openmp] [OpenMP][offload] Cross-team reductions with variable number of teams (PR #195102)

[Robert Imschweiler]

================
@@ -187,130 +248,57 @@ int32_t __kmpc_nvptx_teams_reduce_nowait_v2(
     ThreadId = 0;
   }
 
-  uint32_t &IterCnt = state::getKernelLaunchEnvironment().ReductionIterCnt;
-  uint32_t &Cnt = state::getKernelLaunchEnvironment().ReductionCnt;
-
   // In non-generic mode all workers participate in the teams reduction.
   // In generic mode only the team master participates in the teams
   // reduction because the workers are waiting for parallel work.
   uint32_t NumThreads = omp_get_num_threads();
   uint32_t TeamId = omp_get_team_num();
   uint32_t NumTeams = omp_get_num_teams();
-  [[clang::loader_uninitialized]] static Local<unsigned> Bound;
-  [[clang::loader_uninitialized]] static Local<unsigned> ChunkTeamCount;
-
-  // Block progress for teams greater than the current upper
-  // limit. We always only allow a number of teams less or equal
-  // to the number of slots in the buffer.
-  bool IsMaster = (ThreadId == 0);
-  while (IsMaster) {
-    Bound = atomic::load(&IterCnt, atomic::acquire);
-    if (TeamId < Bound + num_of_records)
-      break;
-  }
 
-  if (IsMaster) {
-    int ModBockId = TeamId % num_of_records;
-    if (TeamId < num_of_records) {
-      lgcpyFct(GlobalBuffer, ModBockId, reduce_data);
-    } else
-      lgredFct(GlobalBuffer, ModBockId, reduce_data);
-
-    // Propagate the memory writes above to the world.
-    fence::kernel(atomic::release);
-
-    // Increment team counter.
-    // This counter is incremented by all teams in the current
-    // num_of_records chunk.
-    ChunkTeamCount = atomic::inc(&Cnt, num_of_records - 1u, atomic::seq_cst,
-                                 atomic::MemScopeTy::device);
+  // Fast path for single-team kernels: no cross-team work required,
+  // the team-local reduction already produced the final result.
+  if (NumTeams <= 1)
+    return ThreadId == 0;
+
+  uint32_t &TeamsDone = state::getKernelLaunchEnvironment().ReductionTeamsDone;
+  void *GlobalBuffer = state::getKernelLaunchEnvironment().ReductionBuffer;
+  [[clang::loader_uninitialized]] static Local<uint32_t> TeamsDoneResult;
+
+  // Save the team's reduced value in the global buffer and atomically
+  // increment the teams-done counter.
+  if (ThreadId == 0) {
+    lgcpyFct(GlobalBuffer, TeamId, reduce_data);
+    TeamsDoneResult = atomic::inc(&TeamsDone, NumTeams - 1u, atomic::acq_rel,
+                                  atomic::MemScopeTy::device);
----------------
ro-i wrote:

TeamsDone is a reference to ReductionTeamsDone, so to the counter variable in global memory. That's what the atomic inc is working on. TeamsDoneResult is in teams/block-local storage and holds the return value of the atomic inc. (Note that ThreadId is the thread id on block-level, not warp-level.)
The following `synchronize::threadsAligned(atomic::acq_rel)` (in SPMD mode) makes sure that all threads in the team/block see the correct value for TeamsDoneResult.
After that, it's safe to check if `TeamsDoneResult != NumTeams - 1u` and thus determine if we're a thread in the last team ("last" as-in "finished last", independent of TeamId). Note that the atomic inc returns the value from *before* the inc, that's why we compare to `NumTeams - 1`.

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