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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Spencer Bryngelson (sbryngelson)

<details>
<summary>Changes</summary>

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`.


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


2 Files Affected:

- (added) offload/test/offloading/distribute_for_no_loop_num_threads.c (+79) 
- (modified) openmp/device/src/Workshare.cpp (+7) 


``````````diff
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..d0e3afa59bfcf
--- /dev/null
+++ b/offload/test/offloading/distribute_for_no_loop_num_threads.c
@@ -0,0 +1,79 @@
+// Regression test for the SPMD no-loop path of __kmpc_distribute_for_static_loop.
+//
+// In that path the index is computed as BId * NumThreads + TId, so NumThreads
+// has to be the size of the block the threads are actually in. If the caller
+// passes a larger value the index strides past the end of each block and the
+// iterations in between are never run.
+//
+// The entry point is called directly here because 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_*. Calling it directly keeps the test
+// independent of the Fortran runtime.
+//
+// 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
+
+// Only the host fallback copy of the target region needs this; the device uses
+// the definition in the device runtime.
+#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 THREAD_LIMIT, the number of threads the
+  // block actually has.
+#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..22971b557103a 100644
--- a/openmp/device/src/Workshare.cpp
+++ b/openmp/device/src/Workshare.cpp
@@ -895,6 +895,13 @@ template <typename Ty> class StaticLoopChunker {
     Ty NumBlocks = mapping::getNumberOfBlocksInKernel();
     Ty BId = mapping::getBlockIdInKernel();
 
+    // In the one-iteration-per-thread case the index is computed as
+    // BId * NumThreads + TId, so NumThreads has to be the size of the block
+    // the threads are actually in. A caller-supplied value that is larger
+    // strides past the end of each block and drops the iterations in between.
+    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;

``````````

</details>


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