[Openmp-commits] [openmp] 9c317bb - [OpenMP][DeviceRTL] Use the actual block size in the SPMD no-loop distribute path (#214073)

[via Openmp-commits]

Author: Spencer Bryngelson
Date: 2026-08-11T10:46:23-04:00
New Revision: 9c317bbbdbeaaf2b2e039ece201baba8f21f9d2a

URL: https://github.com/llvm/llvm-project/commit/9c317bbbdbeaaf2b2e039ece201baba8f21f9d2a
DIFF: https://github.com/llvm/llvm-project/commit/9c317bbbdbeaaf2b2e039ece201baba8f21f9d2a.diff

LOG: [OpenMP][DeviceRTL] Use the actual block size in the SPMD no-loop distribute path (#214073)

In the SPMD no-loop path the iteration index is `BId * NumThreads +
TId`, so `NumThreads` has to
be the size of the block the threads are actually in. `DistributeFor`
uses the caller-supplied
value as-is. If that value is larger than the block size, the index
strides past the end of each
block and the iterations in between are never run.

Fixes #198621.

Reproduced on gfx90a, `num_teams(4)`, 128 iterations, varying the value
passed as `NumThreads`
against the threads the block actually has:

| NumThreads passed | threads in block | iterations not run |
|---|---|---|
| 256 | 32 | 96 of 128 |
| 256 | 64 | 64 of 128 |
| 128 | 32 | 96 of 128 |
| 64 | 64 | 0 |

Exactly as many iterations run as there are threads in the block, and
the rest are dropped. With
this patch all four cases run all 128.

AMD carries the same fix downstream as ROCm/llvm-project#3058.

### Testing

The new test calls `__kmpc_distribute_for_static_loop_4u` directly. That
is unusual, and the reason
is that no C or C++ construct reaches it: clang lowers worksharing loops
to
`__kmpc_for_static_init_4`, and only flang emits
`__kmpc_distribute_for_static_loop_*`. Verified on
device IR from both frontends. Calling the entry directly keeps the test
independent of the Fortran
runtime, which upstream does not ship for the device.

Verified the test fails without the patch (`unwritten: 96`) and passes
with it (`unwritten: 0`).

Also ran 66 Fortran offload probes covering collapse, reductions,
private and firstprivate arrays,
nested parallel, simd, allocatable, atomic and complex arithmetic, in
no-loop mode with
`exec_mode` confirmed to be 6 rather than assumed: all pass before and
after, so the clamp does not
disturb the cases where the caller's value was already correct.

---

Parts of this change were written or audited with Claude Code. I have
reviewed all of it and take
full responsibility for the contribution. See
`llvm/docs/AIToolPolicy.md`.

Added: 
    offload/test/offloading/distribute_for_no_loop_num_threads.c

Modified: 
    openmp/device/src/Workshare.cpp

Removed: 
    


################################################################################
diff  --git a/offload/test/offloading/distribute_for_no_loop_num_threads.c b/offload/test/offloading/distribute_for_no_loop_num_threads.c
new file mode 100644
index 0000000000000..4ad4012efb81d
--- /dev/null
+++ b/offload/test/offloading/distribute_for_no_loop_num_threads.c
@@ -0,0 +1,71 @@
+// SPMD no-loop distribute: the index is BId * NumThreads + TId, so a NumThreads
+// larger than the real block size drops iterations.
+//
+// The entry is called directly because no C/C++ construct reaches it: clang
+// emits __kmpc_for_static_init_4, only flang emits this one.
+//
+// RUN: %libomptarget-compile-run-and-check-generic
+// REQUIRES: gpu
+
+#include <omp.h>
+#include <stdio.h>
+
+#define N 128
+#define NUM_THREADS 256
+#define THREAD_LIMIT 32
+
+struct Args {
+  double *Out;
+};
+
+#pragma omp begin declare target
+extern void __kmpc_distribute_for_static_loop_4u(
+    void *Loc, void (*Fn)(unsigned, void *), void *Arg, unsigned NumIters,
+    unsigned NumThreads, unsigned BlockChunk, unsigned ThreadChunk,
+    unsigned char OneIterationPerThread);
+
+__attribute__((noinline)) static void body(unsigned I, void *A) {
+  ((struct Args *)A)->Out[I] = (double)(I + 1);
+}
+#pragma omp end declare target
+
+// For the host fallback copy only; the device uses the runtime's definition.
+#ifndef __AMDGCN__
+void __kmpc_distribute_for_static_loop_4u(
+    void *Loc, void (*Fn)(unsigned, void *), void *Arg, unsigned NumIters,
+    unsigned NumThreads, unsigned BlockChunk, unsigned ThreadChunk,
+    unsigned char OneIterationPerThread) {
+  for (unsigned I = 0; I < NumIters; ++I)
+    Fn(I, Arg);
+}
+#endif
+
+int main(void) {
+  static double Out[N];
+  for (int I = 0; I < N; ++I)
+    Out[I] = -1.0;
+
+  // NUM_THREADS deliberately exceeds the block's actual thread count.
+#pragma omp target teams map(tofrom : Out[0 : N]) num_teams(4)                 \
+    thread_limit(THREAD_LIMIT)
+  {
+#pragma omp parallel
+    {
+      struct Args A;
+      A.Out = Out;
+      __kmpc_distribute_for_static_loop_4u(0, body, &A, N, NUM_THREADS,
+                                           /*BlockChunk=*/0,
+                                           /*ThreadChunk=*/0,
+                                           /*OneIterationPerThread=*/1);
+    }
+  }
+
+  int Unwritten = 0;
+  for (int I = 0; I < N; ++I)
+    if (Out[I] == -1.0)
+      ++Unwritten;
+
+  // CHECK: unwritten: 0
+  printf("unwritten: %d\n", Unwritten);
+  return 0;
+}

diff  --git a/openmp/device/src/Workshare.cpp b/openmp/device/src/Workshare.cpp
index 6e6440b690db0..a16ef93707989 100644
--- a/openmp/device/src/Workshare.cpp
+++ b/openmp/device/src/Workshare.cpp
@@ -895,6 +895,11 @@ template <typename Ty> class StaticLoopChunker {
     Ty NumBlocks = mapping::getNumberOfBlocksInKernel();
     Ty BId = mapping::getBlockIdInKernel();
 
+    // The index is BId * NumThreads + TId, so NumThreads must be the real
+    // block size; a larger value strides past each block and drops iterations.
+    if (OneIterationPerThread)
+      NumThreads = static_cast<Ty>(mapping::getMaxTeamThreads());
+
     // If the block chunk is not specified we pick a default now.
     if (BlockChunk == 0)
       BlockChunk = NumThreads;