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

[Spencer Bryngelson via Openmp-commits]

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

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


>From df13b38b517c76a6045c7831ac4cd3d4ce36e4ae Mon Sep 17 00:00:00 2001
From: Spencer Bryngelson <sbryngelson at gmail.com>
Date: Thu, 13 Aug 2026 11:52:57 -0500
Subject: [PATCH] [OpenMP][DeviceRTL] Fix the chunked static distribution
 dropping iterations

NormalizedLoopNestChunked starts thread T at BId * BlockChunk + TId, so with a
thread chunk the threads start one index apart and then run consecutive
chunks: they overlap at the front and leave the rest of the space unvisited.
The start index has to account for the chunk.

DistributeFor also defaults BlockChunk to NumThreads, which cannot hold one
chunk of ThreadChunk iterations for each of NumThreads threads, so the threads
past the block chunk get nothing.

The path is only reachable with a non-zero thread chunk, which no frontend
passes today, so nothing in tree changes behaviour.

Assisted-by: Claude <noreply at anthropic.com>
---
 .../offloading/distribute_for_thread_chunk.c  | 75 +++++++++++++++++++
 openmp/device/src/Workshare.cpp               | 12 +--
 2 files changed, 82 insertions(+), 5 deletions(-)
 create mode 100644 offload/test/offloading/distribute_for_thread_chunk.c

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)