[Openmp-commits] [llvm] [openmp] [OpenMP][DeviceRTL] Fix the chunked static distribution dropping iterations (PR #216117)

[via Openmp-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Spencer Bryngelson (sbryngelson)

<details>
<summary>Changes</summary>

`StaticLoopChunker` drops loop iterations whenever a non-zero thread chunk is used, so a chunked static distribution silently computes the wrong answer.

Two things are wrong. `NormalizedLoopNestChunked` starts a thread at `BId * BlockChunk + TId`, which does not account for the chunk, so with `ThreadChunk = C` the threads start one index apart and then each run `C` consecutive iterations: they overlap at the front and leave the rest of the space unvisited. And `DistributeFor` defaults `BlockChunk` to `NumThreads`, which cannot hold one chunk of `C` iterations for each of `NumThreads` threads, so every thread past the block chunk gets nothing at all.

The path is only reachable with a non-zero thread chunk. clang lowers worksharing loops to `__kmpc_for_static_init_4` and never reaches these entries, and flang passes a hardcoded `0` for both chunks, so nothing in tree exercises it today and no in-tree behaviour changes. It becomes reachable as soon as a frontend forwards a `schedule(static, C)` chunk, which is what #<!-- -->214303 needs.

### Testing

New `offload/test/offloading/distribute_for_thread_chunk.c` records which thread ran each iteration for `NumIters=32`, `NumThreads=8`, `ThreadChunk=4`, and checks both that no iteration is skipped and that the mapping is the one `schedule(static,4)` prescribes. It calls the entry directly for the same reason as `distribute_for_no_loop_num_threads.c`: no C or C++ construct reaches it.

On gfx90a, thread recorded per iteration:

| | result |
|---|---|
| expected | `0 0 0 0 1 1 1 1 2 2 2 2 ... 7 7 7 7` |
| before | 17 of 32 iterations unwritten, 10 misplaced |
| after | 0 unwritten, 0 misplaced, mapping exactly as expected |

The default path is unchanged: with `ThreadChunk = 0` the mapping is identical before and after, and `distribute_for_no_loop_num_threads.c` still reports `unwritten: 0`.

`__kmpc_for_static_loop` reaches the same helper with `BlockChunk` hardcoded to `0`, which makes `KernelIteration` zero so the induction variable never advances. That entry is left alone here; it is a separate defect and no frontend passes it a chunk either.

The analysis, the reduced test case and this fix were produced with Claude; I reviewed and verified them.


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


2 Files Affected:

- (added) offload/test/offloading/distribute_for_thread_chunk.c (+75) 
- (modified) openmp/device/src/Workshare.cpp (+7-5) 


``````````diff
diff --git a/offload/test/offloading/distribute_for_thread_chunk.c b/offload/test/offloading/distribute_for_thread_chunk.c
new file mode 100644
index 0000000000000..bd3e8afcde31f
--- /dev/null
+++ b/offload/test/offloading/distribute_for_thread_chunk.c
@@ -0,0 +1,75 @@
+// Chunked static distribution: thread T owns iterations [T*chunk, T*chunk+chunk)
+// within each block chunk. The index used to start each thread has to account
+// for the chunk, and the block chunk has to cover one chunk per thread.
+//
+// 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 32
+#define NT 8
+#define CHUNK 4
+
+struct Args {
+  int *Tid;
+};
+
+#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)->Tid[I] = omp_get_thread_num();
+}
+#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 int Tid[N];
+  for (int I = 0; I < N; ++I)
+    Tid[I] = -1;
+
+#pragma omp target teams map(tofrom : Tid[0 : N]) num_teams(1) thread_limit(NT)
+  {
+#pragma omp parallel num_threads(NT)
+    {
+      struct Args A;
+      A.Tid = Tid;
+      __kmpc_distribute_for_static_loop_4u(0, body, &A, N, NT,
+                                           /*BlockChunk=*/0,
+                                           /*ThreadChunk=*/CHUNK,
+                                           /*OneIterationPerThread=*/0);
+    }
+  }
+
+  int Unwritten = 0, Misplaced = 0;
+  for (int I = 0; I < N; ++I) {
+    if (Tid[I] == -1)
+      ++Unwritten;
+    else if (Tid[I] != (I / CHUNK) % NT)
+      ++Misplaced;
+  }
+
+  // CHECK: unwritten: 0
+  // CHECK: misplaced: 0
+  printf("unwritten: %d\n", Unwritten);
+  printf("misplaced: %d\n", Misplaced);
+  return 0;
+}
diff --git a/openmp/device/src/Workshare.cpp b/openmp/device/src/Workshare.cpp
index a16ef93707989..fdb64c01eb41f 100644
--- a/openmp/device/src/Workshare.cpp
+++ b/openmp/device/src/Workshare.cpp
@@ -732,7 +732,7 @@ template <typename Ty> class StaticLoopChunker {
     Ty KernelIteration = NumBlocks * BlockChunk;
 
     // Start index in the chunked space.
-    Ty IV = BId * BlockChunk + TId;
+    Ty IV = BId * BlockChunk + TId * ThreadChunk;
     ASSERT(IV >= 0, "Bad index");
 
     // Cover the entire iteration space, assumptions in the caller might allow
@@ -900,14 +900,16 @@ template <typename Ty> class StaticLoopChunker {
     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;
-
     // If the thread chunk is not specified we pick a default now.
     if (ThreadChunk == 0)
       ThreadChunk = 1;
 
+    // If the block chunk is not specified we pick a default now. It has to
+    // cover one chunk for every thread, otherwise the threads past the block
+    // chunk get no iterations at all.
+    if (BlockChunk == 0)
+      BlockChunk = NumThreads * ThreadChunk;
+
     // If we know we have more threads (across all blocks) than iterations we
     // can indicate that to avoid an outer loop.
     if (OneIterationPerThread)

``````````

</details>


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