[Openmp-commits] [llvm] [openmp] [OpenMP][offload] Trade threads for blocks in Generic-SPMD kernels (PR #234191)
Robert Imschweiler via Openmp-commits
openmp-commits at lists.llvm.org
Fri Oct 23 07:58:04 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/234191
>From 0535caf68deaf6011a3bb4d8ed79132317a14f1d Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 23 Oct 2026 09:50:40 -0500
Subject: [PATCH 1/3] [OpenMP][offload] Add NumComputeUnits and Backend to
GeometryProps
---
offload/include/device.h | 10 ++++++++++
offload/libompaccsupport/device.cpp | 13 +++++++++++++
2 files changed, 23 insertions(+)
diff --git a/offload/include/device.h b/offload/include/device.h
index 11dc6931a97cc..a849c383d595f 100644
--- a/offload/include/device.h
+++ b/offload/include/device.h
@@ -101,6 +101,10 @@ struct DeviceGeometryPropsTy {
uint32_t DefaultNumBlocks = 0;
/// Maximum amount of shared memory, in bytes, a block can use in a kernel.
uint64_t MaxBlockSharedMemSize = 0;
+ /// Number of compute units, or zero if the device does not report it.
+ uint32_t NumComputeUnits = 0;
+ /// Backend of the platform of the device.
+ ol_platform_backend_t Backend = OL_PLATFORM_BACKEND_UNKNOWN;
};
struct PluginManager;
@@ -133,6 +137,9 @@ struct DeviceTy {
/// Whether this device executes kernels on the host CPU.
bool isHostDevice() const { return IsHostDevice; }
+ /// Get this device's backend.
+ ol_platform_backend_t getBackend() const { return Backend; }
+
/// Load \p Img onto the device and return the resulting program.
llvm::Expected<ProgramTy> loadBinary(__tgt_device_image *Img);
@@ -373,6 +380,9 @@ struct DeviceTy {
/// Whether this is a host device.
bool IsHostDevice = false;
+ /// The backend of this device.
+ ol_platform_backend_t Backend = OL_PLATFORM_BACKEND_UNKNOWN;
+
/// Handler to collect and organize host-2-device mapping information.
MappingInfoTy MappingInfo;
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 0bbfc9ce8698f..73fe1b0fced6b 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -111,6 +111,13 @@ llvm::Error DeviceTy::initGeometryProps() {
if (auto Err = Query(DeviceInfo::WORK_GROUP_LOCAL_MEM_SIZE,
GeometryProps.MaxBlockSharedMemSize))
return Err;
+ // Not every device reports its number of compute units; zero turns off the
+ // heuristics that need it.
+ GeometryProps.NumComputeUnits =
+ getInfo<uint32_t>(DeviceInfo::NUM_COMPUTE_UNITS);
+ // Some defaults depend on the backend. If it is unknown, the generic ones
+ // apply.
+ GeometryProps.Backend = Backend;
return llvm::Error::success();
}
@@ -131,6 +138,12 @@ llvm::Error DeviceTy::init() {
IsHostDevice =
getInfo<ol_device_type_t>(DeviceInfo::TYPE) == OL_DEVICE_TYPE_HOST;
+ // Needs to be initialized before the GeometryProps.
+ if (auto *Platform = getInfo<ol_platform_handle_t>(DeviceInfo::PLATFORM);
+ olGetPlatformInfo(Platform, OL_PLATFORM_INFO_BACKEND,
+ sizeof(GeometryProps.Backend), &GeometryProps.Backend))
+ Backend = OL_PLATFORM_BACKEND_UNKNOWN;
+
getInfoString(DeviceInfo::UID, Uid);
if (Uid.empty())
ODBG(ODT_Init) << "Failed to get UID for device " << DeviceID;
>From bab27b6e6c310ea2437ed1f7b7b82637d10d8a9f Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Thu, 15 Oct 2026 16:53:53 -0500
Subject: [PATCH 2/3] [OpenMP][offload] Trade threads for blocks in
Generic-SPMD kernels
The threads of a Generic-SPMD block run the inner parallel loops, whose
trip counts are unknown. Halve the threads down to one warp while the
distribute loop has an iteration for every block, and default to 6
blocks per CU on AMDGPU (LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU).
libomptarget reads the number of compute units and the platform backend
of the device through liboffload, as part of its geometry properties.
This is an upstream port of the corresponding downstream.
Claude assisted with this patch.
---
offload/libompaccsupport/device.cpp | 46 +++++++++++++++--
offload/test/offloading/generic_spmd_grid.c | 56 +++++++++++++++++++++
openmp/docs/design/Runtimes.md | 11 ++++
3 files changed, 108 insertions(+), 5 deletions(-)
create mode 100644 offload/test/offloading/generic_spmd_grid.c
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 73fe1b0fced6b..6ffcad5ed0e48 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -800,19 +800,38 @@ static uint32_t getEffectiveNumThreads(const DeviceGeometryPropsTy &Props,
: KernelEnv.PreferredNumThreads);
}
+/// Get the number of teams the user set with OMP_NUM_TEAMS, or zero.
+static uint32_t getOMPNumTeams() {
+ static Int32Envar OMP_NumTeams("OMP_NUM_TEAMS");
+ return OMP_NumTeams > 0 ? uint32_t(OMP_NumTeams) : 0;
+}
+
/// Get the maximum number of blocks the device can launch for a kernel using
/// \p NumThreads threads per block, further limited by OMP_NUM_TEAMS if the
/// user set it.
static uint32_t getDeviceBlockLimit(GenericDeviceTy &GenericDevice,
uint32_t NumThreads) {
- static Int32Envar OMP_NumTeams("OMP_NUM_TEAMS");
-
uint32_t BlockLimit = GenericDevice.getBlockLimit(NumThreads);
- if (OMP_NumTeams > 0)
- return std::min(BlockLimit, uint32_t(OMP_NumTeams));
+ if (uint32_t NumTeams = getOMPNumTeams())
+ return std::min(BlockLimit, NumTeams);
return BlockLimit;
}
+/// Get the default number of blocks for a kernel in Generic-SPMD mode.
+static uint32_t
+getDefaultNumBlocksGenericSPMD(const DeviceGeometryPropsTy &Props) {
+ // Envar to set the default number of blocks per compute unit for a kernel in
+ // Generic-SPMD mode on AMDGPU devices. Zero uses the default number of blocks
+ // of other kernels.
+ static UInt32Envar OMPX_GenericSPMDTeamsPerCU(
+ "LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU", 6);
+
+ if (Props.Backend != OL_PLATFORM_BACKEND_AMDGPU ||
+ !OMPX_GenericSPMDTeamsPerCU.get() || !Props.NumComputeUnits)
+ return Props.DefaultNumBlocks;
+ return Props.NumComputeUnits * OMPX_GenericSPMDTeamsPerCU.get();
+}
+
/// Get the effective number of blocks for the kernel based on the
/// user-defined number of blocks and the loop trip count.
/// The number of threads \p EffectiveNumThreads can be adjusted by this
@@ -839,7 +858,9 @@ static uint32_t getEffectiveNumBlocks(
return LoopTripCount > 0 ? (((LoopTripCount - 1) / EffectiveNumThreads) + 1)
: 1;
- uint64_t DefaultNumBlocks = Props.DefaultNumBlocks;
+ uint64_t DefaultNumBlocks = KernelEnv.isGenericSPMDMode()
+ ? getDefaultNumBlocksGenericSPMD(Props)
+ : Props.DefaultNumBlocks;
uint64_t TripCountNumBlocks = std::numeric_limits<uint64_t>::max();
if (LoopTripCount > 0) {
if (KernelEnv.isSPMDMode()) {
@@ -915,6 +936,21 @@ static uint32_t getEffectiveNumBlocks(
"LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT", true);
if (OMPX_ReuseBlocksForHighTripCount.get())
PreferredNumBlocks = std::min(TripCountNumBlocks, DefaultNumBlocks);
+
+ // In Generic-SPMD mode, the threads of a block run the iterations of the
+ // inner parallel loops, whose trip counts are unknown. If they are lower than
+ // the number of threads, most threads idle. So trade threads for blocks, down
+ // to one warp per block, as long as the distribute loop has an iteration for
+ // every block. This keeps the device as busy if the inner trip counts are
+ // high. Neither is done if the user set the number of threads or teams.
+ if (KernelEnv.isGenericSPMDMode() && !IsNumThreadsFromUser &&
+ !IsNumThreadsStrict && !getOMPNumTeams()) {
+ while (uint64_t(PreferredNumBlocks) * 2 <= LoopTripCount &&
+ EffectiveNumThreads > Props.WarpSize) {
+ EffectiveNumThreads /= 2;
+ PreferredNumBlocks *= 2;
+ }
+ }
return std::min(PreferredNumBlocks,
getDeviceBlockLimit(GenericDevice, EffectiveNumThreads));
}
diff --git a/offload/test/offloading/generic_spmd_grid.c b/offload/test/offloading/generic_spmd_grid.c
new file mode 100644
index 0000000000000..8d2c0037b19d0
--- /dev/null
+++ b/offload/test/offloading/generic_spmd_grid.c
@@ -0,0 +1,56 @@
+// clang-format off
+// RUN: %libomptarget-compileopt-generic
+// RUN: env LIBOMPTARGET_INFO=16 \
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+
+// REQUIRES: amdgpu
+// clang-format on
+
+// The number of blocks and threads of kernels in Generic-SPMD mode: by default
+// six blocks per compute unit, then threads traded for blocks down to one
+// wavefront while the distribute loop has an iteration for every block. A
+// thread_limit clause turns the trade off. The number of compute units is
+// taken from the default grid of an SPMD kernel, four blocks per compute unit.
+
+#include <stdio.h>
+
+#define N (1 << 22)
+#define M 64
+
+int A[N];
+
+int main() {
+#pragma omp target teams distribute parallel for thread_limit(256) \
+ map(tofrom : A)
+ for (int I = 0; I < N; ++I)
+ A[I] = 1;
+ // CHECK: Launching kernel {{.*}} with [[[#BLK4:]],1,1] blocks and [256,1,1] threads in SPMD mode
+
+#pragma omp target teams distribute thread_limit(256) map(tofrom : A)
+ for (int I = 0; I < N / M; ++I) {
+#pragma omp parallel for
+ for (int J = 0; J < M; ++J)
+ A[I * M + J] += 1;
+ }
+ // CHECK: Launching kernel {{.*}} with [[[#div(mul(BLK4,3),2)]],1,1] blocks and [256,1,1] threads in Generic-SPMD mode
+
+#pragma omp target teams distribute map(tofrom : A)
+ for (int I = 0; I < N / M; ++I) {
+#pragma omp parallel for
+ for (int J = 0; J < M; ++J)
+ A[I * M + J] += 1;
+ }
+ // CHECK: Launching kernel {{.*}} with [{{[0-9]+}},1,1] blocks and [{{32|64}},1,1] threads in Generic-SPMD mode
+
+#pragma omp target teams distribute map(tofrom : A)
+ for (int I = 0; I < 10; ++I) {
+#pragma omp parallel for
+ for (int J = 0; J < M; ++J)
+ A[I * M + J] += 1;
+ }
+ // CHECK: Launching kernel {{.*}} with [10,1,1] blocks and [{{[0-9]+}},1,1] threads in Generic-SPMD mode
+
+ // CHECK: A[0]=4 A[N-1]=3
+ printf("A[0]=%d A[N-1]=%d\n", A[0], A[N - 1]);
+ return 0;
+}
diff --git a/openmp/docs/design/Runtimes.md b/openmp/docs/design/Runtimes.md
index 6fae6a84a549a..eb25754cefb9d 100644
--- a/openmp/docs/design/Runtimes.md
+++ b/openmp/docs/design/Runtimes.md
@@ -1401,6 +1401,17 @@ number of compute units (CUs) of the AMDGPU device. The default number of teams
is `#default_teams = #teams_per_CU * #CUs`. The default value of teams per CU
is `4`.
+#### LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU
+
+This environment variable controls the default number of teams relative to the
+number of compute units (CUs) of the AMDGPU device for kernels in Generic-SPMD
+mode, e.g., `target teams distribute` with a nested `parallel for`. The default
+value is `6`; `0` uses {ref}`LIBOMPTARGET_AMDGPU_TEAMS_PER_CU` instead. Unless
+the user sets the number of threads or teams, the runtime then halves the
+number of threads of such a kernel, down to one wavefront, and doubles the
+number of teams, as long as the distribute loop has an iteration for every
+team.
+
#### LIBOMPTARGET_AMDGPU_MAX_ASYNC_COPY_BYTES
This environment variable specifies the maximum size in bytes where the memory
>From 16351ecb8272ddcf1be08adaabee89827ee4ae88 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 23 Oct 2026 09:03:44 -0500
Subject: [PATCH 3/3] rename env var
---
offload/libompaccsupport/device.cpp | 6 +++---
openmp/docs/design/Runtimes.md | 16 ++++++++--------
2 files changed, 11 insertions(+), 11 deletions(-)
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 6ffcad5ed0e48..55781b3a1985f 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -821,10 +821,10 @@ static uint32_t getDeviceBlockLimit(GenericDeviceTy &GenericDevice,
static uint32_t
getDefaultNumBlocksGenericSPMD(const DeviceGeometryPropsTy &Props) {
// Envar to set the default number of blocks per compute unit for a kernel in
- // Generic-SPMD mode on AMDGPU devices. Zero uses the default number of blocks
- // of other kernels.
+ // Generic-SPMD mode. Currently only used on AMDGPU devices. Zero uses the
+ // default number of blocks of other kernels.
static UInt32Envar OMPX_GenericSPMDTeamsPerCU(
- "LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU", 6);
+ "LIBOMPTARGET_GENERIC_SPMD_TEAMS_PER_CU", 6);
if (Props.Backend != OL_PLATFORM_BACKEND_AMDGPU ||
!OMPX_GenericSPMDTeamsPerCU.get() || !Props.NumComputeUnits)
diff --git a/openmp/docs/design/Runtimes.md b/openmp/docs/design/Runtimes.md
index eb25754cefb9d..8a0630d1159bf 100644
--- a/openmp/docs/design/Runtimes.md
+++ b/openmp/docs/design/Runtimes.md
@@ -1401,16 +1401,16 @@ number of compute units (CUs) of the AMDGPU device. The default number of teams
is `#default_teams = #teams_per_CU * #CUs`. The default value of teams per CU
is `4`.
-#### LIBOMPTARGET_AMDGPU_GENERIC_SPMD_TEAMS_PER_CU
+#### LIBOMPTARGET_GENERIC_SPMD_TEAMS_PER_CU
This environment variable controls the default number of teams relative to the
-number of compute units (CUs) of the AMDGPU device for kernels in Generic-SPMD
-mode, e.g., `target teams distribute` with a nested `parallel for`. The default
-value is `6`; `0` uses {ref}`LIBOMPTARGET_AMDGPU_TEAMS_PER_CU` instead. Unless
-the user sets the number of threads or teams, the runtime then halves the
-number of threads of such a kernel, down to one wavefront, and doubles the
-number of teams, as long as the distribute loop has an iteration for every
-team.
+number of compute units (CUs) for kernels in Generic-SPMD mode, e.g., `target
+teams distribute` with a nested `parallel for`. The default value is `6`
+(currently only active for AMDGPU); `0` uses
+{ref}`LIBOMPTARGET_AMDGPU_TEAMS_PER_CU` instead. Unless the user sets the
+number of threads or teams, the runtime then halves the number of threads of
+such a kernel, down to one wavefront, and doubles the number of teams, as long
+as the distribute loop has an iteration for every team.
#### LIBOMPTARGET_AMDGPU_MAX_ASYNC_COPY_BYTES
More information about the Openmp-commits
mailing list