[Openmp-commits] [llvm] [openmp] [OpenMP][offload] Size the grid of cross-team reductions by their size (PR #234190)
Robert Imschweiler via Openmp-commits
openmp-commits at lists.llvm.org
Fri Oct 23 07:58:02 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/234190
>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 ff131ad7a14b8054391df530d3ae68fcc0f929f5 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Thu, 15 Oct 2026 16:51:39 -0500
Subject: [PATCH 2/3] [OpenMP][offload] Size the grid of cross-team reductions
by their size
Every block adds to the cost of combining the partial results of a
cross-team reduction. Launch one block per compute unit for reductions
that work with a reduction data size of more than 32 bytes, and enough
blocks for a per-device number of waves per compute unit (16 on AMDGPU)
for the smaller reduction data sizes, capped by the trip count. Such
kernels keep their number of threads: the low-trip-count reduction of
threads does not apply to them. 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 | 60 ++++++++++++++++
.../test/offloading/teams_reduction_grid.c | 68 +++++++++++++++++++
openmp/docs/design/Runtimes.md | 15 ++++
3 files changed, 143 insertions(+)
create mode 100644 offload/test/offloading/teams_reduction_grid.c
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 73fe1b0fced6b..ee4c98ed91a6a 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -813,6 +813,61 @@ static uint32_t getDeviceBlockLimit(GenericDeviceTy &GenericDevice,
return BlockLimit;
}
+/// Threshold for the reduction data size (in bytes) above which a cross-team
+/// reduction classifies as "large".
+static constexpr uint32_t LargeTeamsReductionBytes = 32;
+
+/// Get the number of blocks (each having \p NumThreads threads) to launch a
+/// kernel with a cross-team reduction, or zero to handle it like a
+/// non-reduction kernel.
+///
+/// Every block adds to the cost of combining the partial results of the blocks.
+/// For a "large" reduction (see \p LargeTeamsReductionBytes), this cost
+/// dominates -> one block / compute unit. A small/scalar reduction prefers more
+/// blocks -> as many as needed given the device's number of waves per compute
+/// unit, but at least one / compute unit.
+static uint64_t
+getXTeamReductionNumBlocks(const DeviceGeometryPropsTy &Props,
+ uint32_t NumThreads, uint64_t LoopTripCount,
+ const KernelLaunchInfoTy &KernelEnv) {
+ if (!KernelEnv.ReductionDataSize || !NumThreads)
+ return 0;
+
+ // In Generic and Generic-SPMD mode, each block runs one iteration of the
+ // distribute loop. For these modes, there currently is only performance data
+ // for large reductions.
+ bool IsLarge = KernelEnv.ReductionDataSize > LargeTeamsReductionBytes;
+ bool IsSPMD = KernelEnv.isSPMDMode();
+ if (!IsSPMD && !(IsLarge && (KernelEnv.isGenericMode() ||
+ KernelEnv.isGenericSPMDMode())))
+ return 0;
+
+ // Envar to override the device's number of waves per compute unit for a
+ // kernel with a scalar cross-team reduction. Zero keeps the device's value.
+ static UInt32Envar OMPX_TeamsReductionWavesPerCU(
+ "LIBOMPTARGET_TEAMS_REDUCTION_WAVES_PER_CU", 0);
+ uint32_t ScalarWavesPerCU = OMPX_TeamsReductionWavesPerCU.get();
+ // Measured on MI210, MI300A, MI300X, MI350X, and MI450: scalar reductions
+ // with many iterations prefer more waves, ones with fewer iterations or more
+ // memory traffic prefer less waves. 16 waves were never far from the best.
+ if (!ScalarWavesPerCU && Props.Backend == OL_PLATFORM_BACKEND_AMDGPU)
+ ScalarWavesPerCU = 16;
+ uint64_t NumCUs = Props.NumComputeUnits;
+ if (!ScalarWavesPerCU || !NumCUs)
+ return 0;
+
+ uint32_t WavesPerBlock = llvm::divideCeil(NumThreads, Props.WarpSize);
+ uint64_t BlocksPerCU =
+ IsLarge ? 1 : std::max(1u, ScalarWavesPerCU / WavesPerBlock);
+ uint64_t NumBlocks = NumCUs * BlocksPerCU;
+ // No more blocks than there is work for.
+ if (LoopTripCount > 0)
+ NumBlocks =
+ std::min(NumBlocks, IsSPMD ? llvm::divideCeil(LoopTripCount, NumThreads)
+ : LoopTripCount);
+ return NumBlocks;
+}
+
/// 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,6 +894,11 @@ static uint32_t getEffectiveNumBlocks(
return LoopTripCount > 0 ? (((LoopTripCount - 1) / EffectiveNumThreads) + 1)
: 1;
+ if (uint64_t NumBlocks = getXTeamReductionNumBlocks(
+ Props, EffectiveNumThreads, LoopTripCount, KernelEnv))
+ return std::min(NumBlocks, uint64_t(getDeviceBlockLimit(
+ GenericDevice, EffectiveNumThreads)));
+
uint64_t DefaultNumBlocks = Props.DefaultNumBlocks;
uint64_t TripCountNumBlocks = std::numeric_limits<uint64_t>::max();
if (LoopTripCount > 0) {
diff --git a/offload/test/offloading/teams_reduction_grid.c b/offload/test/offloading/teams_reduction_grid.c
new file mode 100644
index 0000000000000..281d04a335374
--- /dev/null
+++ b/offload/test/offloading/teams_reduction_grid.c
@@ -0,0 +1,68 @@
+// clang-format off
+// RUN: %libomptarget-compile-generic
+// RUN: env LIBOMPTARGET_INFO=16 \
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+// RUN: env LIBOMPTARGET_INFO=16 OMP_NUM_TEAMS=2 \
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic --check-prefix=CAP
+
+// REQUIRES: amdgpu
+// clang-format on
+
+// The number of blocks of kernels with a cross-team reduction: one block per
+// compute unit for a large reduction, and for a scalar reduction as many as
+// fill 16 waves per compute unit, but at least one (one block of 1024 threads
+// has at least 16 waves). The number of compute units is taken from the
+// default grid of a kernel without a reduction, four blocks per compute unit.
+// No more blocks than there are iterations, the num_teams clause is honored,
+// and OMP_NUM_TEAMS caps the number of blocks.
+
+#include <stdio.h>
+
+#define N (1 << 20)
+#define M 8
+
+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
+ // CAP: Launching kernel {{.*}} with [2,1,1] blocks and [256,1,1] threads
+
+ double Sums[M] = {0};
+#pragma omp target teams distribute parallel for thread_limit(256) \
+ reduction(+ : Sums[0 : M])
+ for (int I = 0; I < N; ++I)
+ Sums[I % M] += 1;
+ // CHECK: Launching kernel {{.*}} with [[[#div(BLK4,4)]],1,1] blocks and [256,1,1] threads
+ // CAP: Launching kernel {{.*}} with [2,1,1] blocks and [256,1,1] threads
+
+ double Sum1 = 0;
+#pragma omp target teams distribute parallel for thread_limit(1024) \
+ reduction(+ : Sum1)
+ for (int I = 0; I < N; ++I)
+ Sum1 += 1;
+ // CHECK: Launching kernel {{.*}} with [[[#div(BLK4,4)]],1,1] blocks and [1024,1,1] threads
+ // CAP: Launching kernel {{.*}} with [2,1,1] blocks and [1024,1,1] threads
+
+ double Sum2 = 0;
+#pragma omp target teams distribute parallel for reduction(+ : Sum2)
+ for (int I = 0; I < 16; ++I)
+ Sum2 += 1;
+ // CHECK: Launching kernel {{.*}} with [1,1,1] blocks
+
+ double Sum3 = 0;
+#pragma omp target teams distribute parallel for num_teams(3) \
+ thread_limit(256) reduction(+ : Sum3)
+ for (int I = 0; I < N; ++I)
+ Sum3 += 1;
+ // CHECK: Launching kernel {{.*}} with [3,1,1] blocks and [256,1,1] threads
+
+ // CHECK: Sums[0]=131072 Sum1=1048576 Sum2=16 Sum3=1048576
+ // CAP: Sums[0]=131072 Sum1=1048576 Sum2=16 Sum3=1048576
+ printf("Sums[0]=%.0f Sum1=%.0f Sum2=%.0f Sum3=%.0f\n", Sums[0], Sum1, Sum2,
+ Sum3);
+ return 0;
+}
diff --git a/openmp/docs/design/Runtimes.md b/openmp/docs/design/Runtimes.md
index 6fae6a84a549a..37136938cf632 100644
--- a/openmp/docs/design/Runtimes.md
+++ b/openmp/docs/design/Runtimes.md
@@ -717,6 +717,7 @@ variables is defined below.
> - `LIBOMPTARGET_JIT_SAVE_IMAGE_FILENAME=<out:Filename> (device image file)`
> - `LIBOMPTARGET_MIN_THREADS_FOR_LOW_TRIP_COUNT=<Num> (default: 32)`
> - `LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT=[TRUE/FALSE] (default TRUE)`
+> - `LIBOMPTARGET_TEAMS_REDUCTION_WAVES_PER_CU=<Num> (default: per device)`
> - `OFFLOAD_TRACK_ALLOCATION_TRACES=[TRUE/FALSE] (default FALSE)`
> - `OFFLOAD_TRACK_NUM_KERNEL_LAUNCH_TRACES=<Num> (default 0)`
@@ -1139,6 +1140,20 @@ This environment variable can be used to control how the OpenMP runtime assigns
blocks to loops with high trip counts. By default we reuse existing blocks
rather than spawning new blocks.
+#### LIBOMPTARGET_TEAMS_REDUCTION_WAVES_PER_CU
+
+This environment variable sets the number of waves per compute unit for a
+kernel with a scalar cross-team reduction, e.g., `target teams distribute
+parallel for reduction(...)`. The number of blocks of a kernel with a
+cross-team reduction follows from the size of the reduction: if the data
+reduced per team is larger than 32 bytes, one block per compute unit;
+otherwise, as many blocks as fill this number of waves per compute unit, but
+at least one per compute unit. No more blocks are launched than the loop has
+iterations for. In Generic mode, only large reductions are sized this way. The
+`num_teams` clause sets the number of blocks, and `OMP_NUM_TEAMS` caps it. The
+default on AMDGPU devices is `16`. On other devices, the default is `0`: kernels
+with a cross-team reduction get the same number of blocks as other kernels.
+
#### OFFLOAD_TRACK_ALLOCATION_TRACES
This environment variable determines if the stack traces of allocations and
>From 98e91d5a016633a7ad8b1f5c9a5cdb284ca650dc Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Thu, 22 Oct 2026 08:24:25 -0500
Subject: [PATCH 3/3] [OpenMP][offload] Size only reductions that fill the
device by their size
A cross-team reduction whose loop is too short for one block per compute
unit now gets the grid of a kernel without a reduction. In SPMD mode,
that is the low-trip-count rule, which lowers the number of threads per
block to spread the iterations over more compute units. Keeping all
threads made such reductions, e.g. one of 1792 iterations, up to 2.5x
slower on MI210.
Claude assisted with this patch.
---
offload/libompaccsupport/device.cpp | 16 +++++---
.../test/offloading/teams_reduction_grid.c | 41 ++++++++++++++-----
openmp/docs/design/Runtimes.md | 10 +++--
3 files changed, 46 insertions(+), 21 deletions(-)
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index ee4c98ed91a6a..47036c687effd 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -825,7 +825,8 @@ static constexpr uint32_t LargeTeamsReductionBytes = 32;
/// For a "large" reduction (see \p LargeTeamsReductionBytes), this cost
/// dominates -> one block / compute unit. A small/scalar reduction prefers more
/// blocks -> as many as needed given the device's number of waves per compute
-/// unit, but at least one / compute unit.
+/// unit, but at least one / compute unit. A loop too short to give every
+/// compute unit a block is handled like a non-reduction kernel.
static uint64_t
getXTeamReductionNumBlocks(const DeviceGeometryPropsTy &Props,
uint32_t NumThreads, uint64_t LoopTripCount,
@@ -860,11 +861,14 @@ getXTeamReductionNumBlocks(const DeviceGeometryPropsTy &Props,
uint64_t BlocksPerCU =
IsLarge ? 1 : std::max(1u, ScalarWavesPerCU / WavesPerBlock);
uint64_t NumBlocks = NumCUs * BlocksPerCU;
- // No more blocks than there is work for.
- if (LoopTripCount > 0)
- NumBlocks =
- std::min(NumBlocks, IsSPMD ? llvm::divideCeil(LoopTripCount, NumThreads)
- : LoopTripCount);
+ if (LoopTripCount > 0) {
+ // No more blocks than there is work for.
+ uint64_t TripCountNumBlocks =
+ IsSPMD ? llvm::divideCeil(LoopTripCount, NumThreads) : LoopTripCount;
+ if (TripCountNumBlocks < NumCUs)
+ return 0;
+ NumBlocks = std::min(NumBlocks, TripCountNumBlocks);
+ }
return NumBlocks;
}
diff --git a/offload/test/offloading/teams_reduction_grid.c b/offload/test/offloading/teams_reduction_grid.c
index 281d04a335374..25510828449bf 100644
--- a/offload/test/offloading/teams_reduction_grid.c
+++ b/offload/test/offloading/teams_reduction_grid.c
@@ -13,9 +13,11 @@
// fill 16 waves per compute unit, but at least one (one block of 1024 threads
// has at least 16 waves). The number of compute units is taken from the
// default grid of a kernel without a reduction, four blocks per compute unit.
-// No more blocks than there are iterations, the num_teams clause is honored,
-// and OMP_NUM_TEAMS caps the number of blocks.
+// No more blocks than there are iterations, a loop too short for one block per
+// compute unit gets the grid of a kernel without a reduction, the num_teams
+// clause is honored, and OMP_NUM_TEAMS caps the number of blocks.
+#include <omp.h>
#include <stdio.h>
#define N (1 << 20)
@@ -32,10 +34,14 @@ int main() {
// CAP: Launching kernel {{.*}} with [2,1,1] blocks and [256,1,1] threads
double Sums[M] = {0};
+ int NumCUs = 0;
#pragma omp target teams distribute parallel for thread_limit(256) \
- reduction(+ : Sums[0 : M])
- for (int I = 0; I < N; ++I)
+ reduction(+ : Sums[0 : M]) map(from : NumCUs)
+ for (int I = 0; I < N; ++I) {
Sums[I % M] += 1;
+ if (I == 0)
+ NumCUs = omp_get_num_teams();
+ }
// CHECK: Launching kernel {{.*}} with [[[#div(BLK4,4)]],1,1] blocks and [256,1,1] threads
// CAP: Launching kernel {{.*}} with [2,1,1] blocks and [256,1,1] threads
@@ -47,11 +53,23 @@ int main() {
// CHECK: Launching kernel {{.*}} with [[[#div(BLK4,4)]],1,1] blocks and [1024,1,1] threads
// CAP: Launching kernel {{.*}} with [2,1,1] blocks and [1024,1,1] threads
+ // Two blocks of 256 threads per compute unit.
double Sum2 = 0;
-#pragma omp target teams distribute parallel for reduction(+ : Sum2)
- for (int I = 0; I < 16; ++I)
+#pragma omp target teams distribute parallel for thread_limit(256) \
+ reduction(+ : Sum2)
+ for (int I = 0; I < NumCUs * 512; ++I)
Sum2 += 1;
- // CHECK: Launching kernel {{.*}} with [1,1,1] blocks
+ // CHECK: Launching kernel {{.*}} with [[[#div(BLK4,2)]],1,1] blocks and [256,1,1] threads
+ // CAP: Launching kernel {{.*}} with [2,1,1] blocks and [256,1,1] threads
+
+ // Too short for one block of the default size per compute unit: as for a
+ // kernel without a reduction, one block of 32 threads per compute unit.
+ double Sum4 = 0;
+#pragma omp target teams distribute parallel for reduction(+ : Sum4)
+ for (int I = 0; I < NumCUs * 32; ++I)
+ Sum4 += 1;
+ // CHECK: Launching kernel {{.*}} with [[[#div(BLK4,4)]],1,1] blocks and [32,1,1] threads
+ // CAP: Launching kernel {{.*}} with [2,1,1] blocks and [32,1,1] threads
double Sum3 = 0;
#pragma omp target teams distribute parallel for num_teams(3) \
@@ -60,9 +78,10 @@ int main() {
Sum3 += 1;
// CHECK: Launching kernel {{.*}} with [3,1,1] blocks and [256,1,1] threads
- // CHECK: Sums[0]=131072 Sum1=1048576 Sum2=16 Sum3=1048576
- // CAP: Sums[0]=131072 Sum1=1048576 Sum2=16 Sum3=1048576
- printf("Sums[0]=%.0f Sum1=%.0f Sum2=%.0f Sum3=%.0f\n", Sums[0], Sum1, Sum2,
- Sum3);
+ // CHECK: Sums[0]=131072 Sum1=1048576 Sum2=[[#mul(div(BLK4,4),512)]]
+ // CHECK-SAME: Sum3=1048576 Sum4=[[#mul(div(BLK4,4),32)]]
+ // CAP: Sums[0]=131072 Sum1=1048576 Sum2=1024 Sum3=1048576 Sum4=64
+ printf("Sums[0]=%.0f Sum1=%.0f Sum2=%.0f Sum3=%.0f Sum4=%.0f\n", Sums[0],
+ Sum1, Sum2, Sum3, Sum4);
return 0;
}
diff --git a/openmp/docs/design/Runtimes.md b/openmp/docs/design/Runtimes.md
index 37136938cf632..66b0f843be9b2 100644
--- a/openmp/docs/design/Runtimes.md
+++ b/openmp/docs/design/Runtimes.md
@@ -1149,10 +1149,12 @@ cross-team reduction follows from the size of the reduction: if the data
reduced per team is larger than 32 bytes, one block per compute unit;
otherwise, as many blocks as fill this number of waves per compute unit, but
at least one per compute unit. No more blocks are launched than the loop has
-iterations for. In Generic mode, only large reductions are sized this way. The
-`num_teams` clause sets the number of blocks, and `OMP_NUM_TEAMS` caps it. The
-default on AMDGPU devices is `16`. On other devices, the default is `0`: kernels
-with a cross-team reduction get the same number of blocks as other kernels.
+iterations for. (A loop too short for one block per compute unit gets the same
+grid as a kernel without a reduction.) In Generic mode, only large reductions
+are sized this way. The `num_teams` clause sets the number of blocks, and
+`OMP_NUM_TEAMS` caps it. The default on AMDGPU devices is `16`. On other
+devices, the default is `0`: kernels with a cross-team reduction get the same
+number of blocks as other kernels.
#### OFFLOAD_TRACK_ALLOCATION_TRACES
More information about the Openmp-commits
mailing list