[clang] [llvm] [offload][OpenMP] Improve cross-team reduction grid selection (PR #208253)
Robert Imschweiler via cfe-commits
cfe-commits at lists.llvm.org
Sun Jul 12 07:55:47 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/208253
>From df6b80d3e4fd1a24c35e73bc4f4f58ca71fcec4a Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 8 Jul 2026 11:08:21 -0500
Subject: [PATCH 1/2] [offload][OpenMP] Improve cross-team reduction grid
selection
The default cross-team reduction algorithm benefits from larger and
fewer teams. Implement that by using 2 x the default number of threads
and 1/2 x the default number of teams for reduction kernels. This
doesn't change the default total number of threads, it just
redistributes them.
This is a first, rather simple heuristic, derived from (a subset of)
what AOMP does.
The performance benefits I observed for the reduction tests in
https://github.com/ro-i/xteam-test on a gfx942
(c71339705091500f731e2a39f247d2660bacbdce) are up to a few percent, with
no regressions.
Claude assisted with this patch.
---
clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp | 22 ++++++++++++++++
.../common/include/PluginInterface.h | 5 ++++
.../common/src/PluginInterface.cpp | 26 ++++++++++++++++---
3 files changed, 49 insertions(+), 4 deletions(-)
diff --git a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
index 3fccd3a291d37..bd36758c7eea0 100644
--- a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
+++ b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
@@ -28,6 +28,19 @@ using namespace clang;
using namespace CodeGen;
using namespace llvm::omp;
+/// Returns true if \p D is a teams reduction directive. This corresponds to the
+/// condition under which CGOpenMPRuntimeGPU::emitReduction builds a teams
+/// reduction (its ReductionKind is OMPD_teams and there is at least one
+/// reduction clause).
+static bool hasTeamsReduction(const OMPExecutableDirective &D) {
+ if (!isOpenMPTeamsDirective(D.getDirectiveKind()))
+ return false;
+ for (const auto *C : D.getClausesOfKind<OMPReductionClause>())
+ if (C->getModifier() != OMPC_REDUCTION_inscan)
+ return true;
+ return false;
+}
+
namespace {
/// Pre(post)-action for different OpenMP constructs specialized for NVPTX.
class NVPTXActionTy final : public PrePostActionTy {
@@ -751,6 +764,15 @@ void CGOpenMPRuntimeGPU::emitKernelInit(const OMPExecutableDirective &D,
: llvm::omp::OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC;
computeMinAndMaxThreadsAndTeams(D, CGF, Attrs);
+ // Set MaxThreads to twice the default workgroup size for SPMD cross-team
+ // reductions. For these reductions, we want to have fewer, larger teams to
+ // make better use of intra-team reduction and not over-emphasize the
+ // inter-team part. Note that this only buys us more freedom. The actual grid
+ // size for the launch is going to be chosen later by the offload plugin.
+ if (IsSPMD && Attrs.MaxThreads.front() < 0 && hasTeamsReduction(D))
+ Attrs.MaxThreads.front() =
+ 2 * CGM.getTarget().getGridValue().GV_Default_WG_Size;
+
CGBuilderTy &Bld = CGF.Builder;
Bld.restoreIP(OMPBuilder.createTargetInit(Bld, Attrs));
if (!IsSPMD)
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index dc21abf1a334a..308336e105524 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -574,6 +574,11 @@ struct GenericKernelTy {
DeviceImageTy *ImagePtr = nullptr;
protected:
+ /// Whether the kernel performs a cross-team reduction.
+ bool isCrossTeamReduction() const {
+ return KernelEnvironment.Configuration.ReductionDataSize != 0;
+ }
+
/// The preferred number of threads to run the kernel.
uint32_t PreferredNumThreads;
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 70e10ca6be24a..6a08f9fe1853c 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -119,8 +119,7 @@ GenericKernelTy::getKernelLaunchEnvironment(
KernelArgs.Version < OMP_KERNEL_ARG_MIN_VERSION_WITH_DYN_PTR)
return nullptr;
- const auto &RedCfg = KernelEnvironment.Configuration;
- const bool NeedsReductionBuffer = RedCfg.ReductionDataSize != 0;
+ const bool NeedsReductionBuffer = isCrossTeamReduction();
if (NeedsReductionBuffer && KernelArgs.Version < OMP_KERNEL_ARG_VERSION)
return Plugin::error(ErrorCode::INVALID_BINARY,
"kernel was built against an older OpenMP "
@@ -150,6 +149,7 @@ GenericKernelTy::getKernelLaunchEnvironment(
LocalKLE.ReductionBuffer = nullptr;
if (NeedsReductionBuffer) {
+ const auto &RedCfg = KernelEnvironment.Configuration;
// Use number of teams many buffer elements.
auto AllocOrErr = GenericDevice.dataAlloc(
uint64_t(RedCfg.ReductionDataSize) * NumBlocks0,
@@ -390,8 +390,13 @@ GenericKernelTy::getEffectiveNumThreads(GenericDeviceTy &GenericDevice,
if (UserThreadLimit > 0 && isGenericMode())
UserThreadLimit += GenericDevice.getWarpSize();
- return std::min(MaxNumThreads, (UserThreadLimit > 0) ? UserThreadLimit
- : PreferredNumThreads);
+ // Favor fewer, larger teams for cross-team reductions. We deliberately
+ // increased MaxNumThreads in Codegen.
+ return std::min(
+ MaxNumThreads,
+ (UserThreadLimit > 0)
+ ? UserThreadLimit
+ : (isCrossTeamReduction() ? MaxNumThreads : PreferredNumThreads));
}
uint32_t GenericKernelTy::getEffectiveNumBlocks(
@@ -422,6 +427,19 @@ uint32_t GenericKernelTy::getEffectiveNumBlocks(
// will execute one iteration of the loop; rounded up to the nearest
// integer. However, if that results in too few blocks, we artificially
// reduce the thread count per block to increase the outer parallelism.
+
+ // Favor fewer, larger teams for cross-team reductions. As a heuristic, we
+ // aim for double the default number of threads and half the default
+ // number of blocks.
+ if (isCrossTeamReduction()) {
+ // Never launch more teams than the trip count needs.
+ uint64_t NumBlocks =
+ std::min<uint64_t>(std::max<uint64_t>(DefaultNumBlocks / 2, 1),
+ ((LoopTripCount - 1) / EffectiveNumThreads) + 1);
+ return std::min(NumBlocks, uint64_t(GenericDevice.getBlockLimit(
+ EffectiveNumThreads)));
+ }
+
auto MinThreads = GenericDevice.getMinThreadsForLowTripCountLoop();
MinThreads = std::min(MinThreads, EffectiveNumThreads);
>From 21bc0e987c91bfa9797bd51526294eaa7f5bf26b Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Sun, 12 Jul 2026 09:53:24 -0500
Subject: [PATCH 2/2] add test
---
.../offloading/xteam_reduction_default_grid.c | 61 +++++++++++++++++++
1 file changed, 61 insertions(+)
create mode 100644 offload/test/offloading/xteam_reduction_default_grid.c
diff --git a/offload/test/offloading/xteam_reduction_default_grid.c b/offload/test/offloading/xteam_reduction_default_grid.c
new file mode 100644
index 0000000000000..90cdfed869252
--- /dev/null
+++ b/offload/test/offloading/xteam_reduction_default_grid.c
@@ -0,0 +1,61 @@
+// RUN: %libomptarget-compile-generic
+// RUN: env LIBOMPTARGET_INFO=16 \
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+// RUN: %libomptarget-compileopt-generic
+// RUN: env LIBOMPTARGET_INFO=16 \
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+
+// REQUIRES: amdgpu
+
+// A cross-team reduction launches with twice the threads per team and half the
+// teams of an equivalent non-reduction loop, keeping the total thread count of
+// the default grid.
+
+#include <stdio.h>
+#include <stdlib.h>
+
+// ceil(SMALL_N / (2 * default workgroup size)) = 3 teams, far below the
+// half-default team cap on any GPU.
+#define SMALL_N (3 * 512)
+
+int main(void) {
+ // High trip count so both grids hit the default team-count cap.
+ const int N = 4 * 1024 * 1024;
+ int *a = (int *)malloc(sizeof(int) * N);
+ for (int i = 0; i < N; ++i)
+ a[i] = 1;
+
+ // Reference: capture the default grid
+ // CHECK: Launching kernel {{.*}} with {{\[}}[[#REFB:]],1,1] blocks and
+ // CHECK-SAME: {{\[}}[[#REFT:]],1,1] threads in SPMD mode
+#pragma omp target teams distribute parallel for map(tofrom : a[0 : N])
+ for (int i = 0; i < N; ++i)
+ a[i] += 1;
+
+ // High trip count: half teams, double threads
+ // CHECK: Launching kernel {{.*}} with {{\[}}[[#div(REFB,2)]],1,1] blocks and
+ // CHECK-SAME: {{\[}}[[#mul(REFT,2)]],1,1] threads in SPMD mode
+ long long sum = 0;
+#pragma omp target teams distribute parallel for reduction(+ : sum) \
+ map(tofrom : a[0 : N])
+ for (int i = 0; i < N; ++i)
+ sum += a[i];
+
+ // Small trip count: the trip-count bound wins, so teams follow the trip count
+ // while threads stay widened.
+ // CHECK: Launching kernel {{.*}} with [3,1,1] blocks and
+ // CHECK-SAME: {{\[}}[[#mul(REFT,2)]],1,1] threads in SPMD mode
+ long long small_sum = 0;
+#pragma omp target teams distribute parallel for reduction(+ : small_sum) \
+ map(tofrom : a[0 : SMALL_N])
+ for (int i = 0; i < SMALL_N; ++i)
+ small_sum += a[i];
+
+ // CHECK: sum = 8388608
+ // CHECK: small_sum = 3072
+ printf("sum = %lld\n", sum);
+ printf("small_sum = %lld\n", small_sum);
+
+ free(a);
+ return 0;
+}
More information about the cfe-commits
mailing list