[Openmp-commits] [openmp] 29a956e - [offload][OpenMP] Add atomic cross-team reductions (#209298)
via Openmp-commits
openmp-commits at lists.llvm.org
Sun Aug 2 05:01:38 PDT 2026
Author: Robert Imschweiler
Date: 2026-08-02T14:01:32+02:00
New Revision: 29a956e40a1f617887ecfeb02ae844102c4a5a48
URL: https://github.com/llvm/llvm-project/commit/29a956e40a1f617887ecfeb02ae844102c4a5a48
DIFF: https://github.com/llvm/llvm-project/commit/29a956e40a1f617887ecfeb02ae844102c4a5a48.diff
LOG: [offload][OpenMP] Add atomic cross-team reductions (#209298)
Regular cross-team reductions have two phases: the intra-team reduction
and the inter-team reduction. Atomic cross-team reductions replace the
second phase with a atomic instruction which is used by the main thread
of each team to directly fold the result of the intra-team reduction
into the final result. Since this requires a combination of "data type"
and "combine operation" for which an atomic instruction is available,
only some (but very common) reductions can be transformed to atomic
reductions. In cases where multiple reductions are performed on the same
construct, the atomic path is only taken if all reductions can be
transformed. Otherwise, we fall back to the regular cross-team reduction
using a buffer with per-team slots. This is not strictly necessary, but
hybrid reductions would induce more complexity with questionable
benefit.
Selecting an atomic path might not be the best option for every
situation, which is why it is not enabled by default. Instead, it can be
enabled via `-fopenmp-target-atomic-reduction`. Note that enabling the
atomic path will not *force* atomic reductions. They will only be
applied if possible, as described above.
The performance (measured with https://github.com/ro-i/xteam-test @
c71339705091500f731e2a39f247d2660bacbdce, array size 177,777,777) is up
to +15% faster (aka, more throughput) for supported reductions on a
gfx942, with no noticeable regressions.
Example:
- sum reduction, type double: +10.22% faster
- sum reduction, type uint: +15.57% faster
- sum reduction, type ulong: +13.31% faster
On a gfx90a, there is little to negative benefit:
- sum reduction, type double: -4.32% faster (aka, slower)
- sum reduction, type uint: +3.08% faster
- sum reduction, type ulong: +1.68% faster
Claude assisted with this patch.
Added:
clang/test/Driver/openmp-target-atomic-reduction-flag.c
clang/test/OpenMP/target_teams_atomic_reduction_codegen.cpp
offload/test/offloading/xteam_atomic_reduction_usm.cpp
Modified:
clang/include/clang/Basic/LangOptions.def
clang/include/clang/Options/Options.td
clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
clang/lib/Driver/ToolChains/Clang.cpp
llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
offload/test/offloading/multiple_reductions.cpp
openmp/device/include/Interface.h
openmp/device/src/Kernel.cpp
Removed:
################################################################################
diff --git a/clang/include/clang/Basic/LangOptions.def b/clang/include/clang/Basic/LangOptions.def
index 3d63b9677e4df..0519514327355 100644
--- a/clang/include/clang/Basic/LangOptions.def
+++ b/clang/include/clang/Basic/LangOptions.def
@@ -241,6 +241,7 @@ LANGOPT(OpenMPNoThreadState , 1, 0, NotCompatible, "Assume that no thread in a
LANGOPT(OpenMPNoNestedParallelism , 1, 0, NotCompatible, "Assume that no thread in a parallel region will encounter a parallel region")
LANGOPT(OpenMPOffloadMandatory , 1, 0, NotCompatible, "Assert that offloading is mandatory and do not create a host fallback.")
LANGOPT(OpenMPForceUSM , 1, 0, NotCompatible, "Enable OpenMP unified shared memory mode via compiler.")
+LANGOPT(OpenMPTargetAtomicReduction , 1, 0, NotCompatible, "Use atomic operations for OpenMP GPU cross-team reductions where supported.")
LANGOPT(NoGPULib , 1, 0, NotCompatible, "Indicate a build without the standard GPU libraries.")
LANGOPT(HLSL, 1, 0, NotCompatible, "HLSL")
diff --git a/clang/include/clang/Options/Options.td b/clang/include/clang/Options/Options.td
index 2defcee88f741..2467ebd0abe19 100644
--- a/clang/include/clang/Options/Options.td
+++ b/clang/include/clang/Options/Options.td
@@ -4306,6 +4306,15 @@ def fopenmp_target_fast : Flag<["-"], "fopenmp-target-fast">,
def fno_openmp_target_fast : Flag<["-"], "fno-openmp-target-fast">,
Group<f_Group>, Flags<[NoArgumentUnused, HelpHidden]>,
Visibility<[ClangOption, FlangOption]>;
+def fopenmp_target_atomic_reduction : Flag<["-"], "fopenmp-target-atomic-reduction">,
+ Group<f_Group>, Flags<[NoArgumentUnused]>,
+ Visibility<[ClangOption, CC1Option]>,
+ HelpText<"Use atomic operations for OpenMP GPU cross-team reductions where "
+ "supported">,
+ MarshallingInfoFlag<LangOpts<"OpenMPTargetAtomicReduction">>;
+def fno_openmp_target_atomic_reduction : Flag<["-"], "fno-openmp-target-atomic-reduction">,
+ Group<f_Group>, Flags<[NoArgumentUnused, HelpHidden]>,
+ Visibility<[ClangOption, CC1Option]>;
defm openmp_optimistic_collapse : BoolFOption<"openmp-optimistic-collapse",
LangOpts<"OpenMPOptimisticCollapse">, DefaultFalse,
PosFlag<SetTrue, [], [ClangOption, CC1Option]>,
diff --git a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
index 79b2b275ba100..8f9000660d86b 100644
--- a/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
+++ b/clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp
@@ -14,6 +14,7 @@
#include "CGOpenMPRuntimeGPU.h"
#include "CGDebugInfo.h"
#include "CodeGenFunction.h"
+#include "TargetInfo.h"
#include "clang/AST/Attr.h"
#include "clang/AST/DeclOpenMP.h"
#include "clang/AST/OpenMPClause.h"
@@ -22,6 +23,8 @@
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
#include "llvm/Frontend/OpenMP/OMPGridValues.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
#include "llvm/TargetParser/NVPTXTargetParser.h"
using namespace clang;
@@ -1445,6 +1448,59 @@ static llvm::Value *castValueToType(CodeGenFunction &CGF, llvm::Value *Val,
TBAAAccessInfo());
}
+/// Extracts the built-in reduction operator from a combiner of the form `x = x
+/// <op> rhs` (or the min/max conditional), or nullopt if the shape is not
+/// recognized (e.g. user-defined reductions).
+static std::optional<BinaryOperatorKind>
+getReductionBinOpKind(const Expr *ReductionOp) {
+ const auto *Assign = dyn_cast<BinaryOperator>(ReductionOp);
+ if (!Assign || Assign->getOpcode() != BO_Assign)
+ return std::nullopt;
+ const Expr *RHS = Assign->getRHS();
+ // min/max are lowered as `x <cmp> rhs ? x : rhs`; the comparison identifies
+ // it.
+ if (const auto *ACO =
+ dyn_cast<AbstractConditionalOperator>(RHS->IgnoreParenImpCasts()))
+ RHS = ACO->getCond();
+ if (const auto *BO = dyn_cast<BinaryOperator>(RHS->IgnoreParenImpCasts()))
+ return BO->getOpcode();
+ return std::nullopt;
+}
+
+/// Maps a built-in reduction operator to an atomicrmw opcode for the atomic
+/// cross-team reduction fast path, or nullopt if there is no direct atomicrmw
+/// (e.g. user-defined, complex, fp min/max) so the buffer path is used instead.
+static std::optional<llvm::AtomicRMWInst::BinOp>
+getReductionAtomicRMWOp(BinaryOperatorKind BOK, QualType Ty) {
+ bool IsInt = Ty->isIntegerType();
+ bool IsSigned = Ty->hasSignedIntegerRepresentation();
+ switch (BOK) {
+ case BO_Add:
+ case BO_Sub: // A `-` reduction sums the partials, so it accumulates with add.
+ if (IsInt)
+ return llvm::AtomicRMWInst::Add;
+ if (Ty->isFloatingType())
+ return llvm::AtomicRMWInst::FAdd;
+ return std::nullopt;
+ case BO_And:
+ return IsInt ? std::optional(llvm::AtomicRMWInst::And) : std::nullopt;
+ case BO_Or:
+ return IsInt ? std::optional(llvm::AtomicRMWInst::Or) : std::nullopt;
+ case BO_Xor:
+ return IsInt ? std::optional(llvm::AtomicRMWInst::Xor) : std::nullopt;
+ case BO_LT: // min
+ if (IsInt)
+ return IsSigned ? llvm::AtomicRMWInst::Min : llvm::AtomicRMWInst::UMin;
+ return std::nullopt;
+ case BO_GT: // max
+ if (IsInt)
+ return IsSigned ? llvm::AtomicRMWInst::Max : llvm::AtomicRMWInst::UMax;
+ return std::nullopt;
+ default:
+ return std::nullopt;
+ }
+}
+
///
/// Design of OpenMP reductions on the GPU
///
@@ -1716,8 +1772,13 @@ void CGOpenMPRuntimeGPU::emitReduction(
const RecordDecl *ReductionRec = ::buildRecordForGlobalizedVars(
CGM.getContext(), PrivatesReductions, {}, VarFieldMap, 1);
- if (TeamsReduction)
- TeamsReductions.push_back(ReductionRec);
+ // The atomic cross-team reduction fast path is opt-in. Hand each eligible
+ // scalar reduction an atomic combiner; createReductionsGPU uses the atomic
+ // path only if every reduction in the set has one. Track whether that holds
+ // so we can skip the (then unused) per-team buffer registration.
+ bool UseAtomicReduction =
+ TeamsReduction && CGM.getLangOpts().OpenMPTargetAtomicReduction;
+ bool AllAtomicable = UseAtomicReduction;
// Source location for the ident struct
llvm::Value *RTLoc = emitUpdateLocation(CGF, Loc);
@@ -1780,6 +1841,45 @@ void CGOpenMPRuntimeGPU::emitReduction(
return InsertPointTy(CGF.Builder.GetInsertBlock(),
CGF.Builder.GetInsertPoint());
};
+
+ // For the atomic fast path, hand this reduction an atomic combiner if it is
+ // a scalar with a direct atomicrmw; otherwise the set is not fully
+ // atomicable and falls back to the buffer path.
+ if (UseAtomicReduction) {
+ std::optional<llvm::AtomicRMWInst::BinOp> AtomicOp;
+ if (EvalKind == llvm::OpenMPIRBuilder::EvalKind::Scalar) {
+ if (std::optional<BinaryOperatorKind> BOK =
+ getReductionBinOpKind(ReductionOps[Idx]))
+ AtomicOp = getReductionAtomicRMWOp(*BOK, Private->getType());
+ }
+ if (!AtomicOp) {
+ AllAtomicable = false;
+ } else {
+ llvm::AtomicRMWInst::BinOp Op = *AtomicOp;
+ llvm::Align Alignment =
+ CGM.getModule().getDataLayout().getPrefTypeAlign(ElementType);
+ // Device (agent) scope suffices: all teams accumulate on-device and the
+ // host reads the result only after the kernel (via map-back), so the
+ // far costlier system scope is unnecessary. The
+ // no.fine.grained/no.remote memory metadata is omitted so the atomic
+ // stays correct under USM.
+ llvm::SyncScope::ID SSID = CGF.getTargetHooks().getLLVMSyncScopeID(
+ CGF.getLangOpts(), SyncScope::DeviceScope,
+ llvm::AtomicOrdering::Monotonic, CGF.getLLVMContext());
+ AtomicReductionGen = [Op, Alignment,
+ SSID](InsertPointTy IP, llvm::Type *EltTy,
+ llvm::Value *LHS, llvm::Value *RHS)
+ -> llvm::OpenMPIRBuilder::InsertPointOrErrorTy {
+ llvm::IRBuilder<> Builder(IP.getBlock(), IP.getPoint());
+ llvm::Value *Val = Builder.CreateLoad(EltTy, RHS);
+ Builder.CreateAtomicRMW(Op, LHS, Val, Alignment,
+ llvm::AtomicOrdering::Monotonic, SSID);
+ return InsertPointTy(Builder.GetInsertBlock(),
+ Builder.GetInsertPoint());
+ };
+ }
+ }
+
ReductionInfos.emplace_back(llvm::OpenMPIRBuilder::ReductionInfo(
ElementType, Variable, PrivateVariable, EvalKind,
/*ReductionGen=*/nullptr, ReductionGen, AtomicReductionGen,
@@ -1787,6 +1887,11 @@ void CGOpenMPRuntimeGPU::emitReduction(
Idx++;
}
+ // The atomic path folds directly into the mapped variable and needs no
+ // per-team buffer; register the record for buffer allocation otherwise.
+ if (TeamsReduction && !AllAtomicable)
+ TeamsReductions.push_back(ReductionRec);
+
bool IsSPMD = getExecutionMode() == CGOpenMPRuntimeGPU::EM_SPMD;
llvm::OpenMPIRBuilder::InsertPointTy AfterIP =
cantFail(OMPBuilder.createReductionsGPU(
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
index 7f3b503153ab1..9e90994c178dd 100644
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -7128,6 +7128,12 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
/*Default=*/TargetFastUsed))
CmdArgs.push_back("-fopenmp-assume-no-nested-parallelism");
+ // Handle -fopenmp-target-atomic-reduction.
+ if (Args.hasFlag(options::OPT_fopenmp_target_atomic_reduction,
+ options::OPT_fno_openmp_target_atomic_reduction,
+ /*Default=*/false))
+ CmdArgs.push_back("-fopenmp-target-atomic-reduction");
+
if (Args.hasArg(options::OPT_fopenmp_offload_mandatory))
CmdArgs.push_back("-fopenmp-offload-mandatory");
if (Args.hasArg(options::OPT_fopenmp_force_usm))
diff --git a/clang/test/Driver/openmp-target-atomic-reduction-flag.c b/clang/test/Driver/openmp-target-atomic-reduction-flag.c
new file mode 100644
index 0000000000000..59362ddd78e94
--- /dev/null
+++ b/clang/test/Driver/openmp-target-atomic-reduction-flag.c
@@ -0,0 +1,18 @@
+// REQUIRES: x86-registered-target, amdgpu-registered-target
+
+// Not passed by default.
+// RUN: %clang -### -fopenmp=libomp -fopenmp-targets=amdgcn-amd-amdhsa -Xopenmp-target=amdgcn-amd-amdhsa -march=gfx90a -nogpulib %s 2>&1 \
+// RUN: | FileCheck -check-prefix=DEFAULT %s
+
+// Passed through to -cc1 when requested.
+// RUN: %clang -### -fopenmp=libomp -fopenmp-targets=amdgcn-amd-amdhsa -Xopenmp-target=amdgcn-amd-amdhsa -march=gfx90a -nogpulib -fopenmp-target-atomic-reduction %s 2>&1 \
+// RUN: | FileCheck -check-prefix=ENABLE %s
+
+// Explicit disable wins over a preceding enable and is not passed through.
+// RUN: %clang -### -fopenmp=libomp -fopenmp-targets=amdgcn-amd-amdhsa -Xopenmp-target=amdgcn-amd-amdhsa -march=gfx90a -nogpulib -fopenmp-target-atomic-reduction -fno-openmp-target-atomic-reduction %s 2>&1 \
+// RUN: | FileCheck -check-prefix=DEFAULT %s
+
+// DEFAULT-NOT: {{"-f(no-)?openmp-target-atomic-reduction"}}
+
+// ENABLE: "-fopenmp-target-atomic-reduction"
+// ENABLE-NOT: "-fno-openmp-target-atomic-reduction"
diff --git a/clang/test/OpenMP/target_teams_atomic_reduction_codegen.cpp b/clang/test/OpenMP/target_teams_atomic_reduction_codegen.cpp
new file mode 100644
index 0000000000000..dea0ecff52be4
--- /dev/null
+++ b/clang/test/OpenMP/target_teams_atomic_reduction_codegen.cpp
@@ -0,0 +1,90 @@
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple x86_64-unknown-linux-gnu \
+// RUN: -fopenmp-targets=amdgcn-amd-amdhsa -emit-llvm-bc %s -o %t-host.bc
+// RUN: %clang_cc1 -verify -fopenmp -x c++ -triple amdgcn-amd-amdhsa \
+// RUN: -fopenmp-targets=amdgcn-amd-amdhsa -fopenmp-is-target-device \
+// RUN: -fopenmp-target-atomic-reduction -fopenmp-host-ir-file-path %t-host.bc \
+// RUN: -emit-llvm %s -o %t.ll
+// RUN: FileCheck --check-prefix=ATOMIC %s < %t.ll
+// RUN: FileCheck --check-prefix=BUFFER %s < %t.ll
+// RUN: FileCheck %s < %t.ll
+
+// expected-no-diagnostics
+
+// ATOMIC-DAG: call i32 @__kmpc_is_team_main_thread(
+// ATOMIC-DAG: atomicrmw add ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw add ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw and ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw or ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw xor ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw min ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw max ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw umin ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw umax ptr {{.*}} syncscope("agent") monotonic
+// ATOMIC-DAG: atomicrmw fadd ptr {{.*}} syncscope("agent") monotonic
+void atomicable_matrix(int n) {
+ int si = 0, su = 0, an = ~0, orv = 0, xr = 0, smn = 0x7fffffff,
+ smx = -0x7fffffff - 1;
+ unsigned umn = ~0u, umx = 0u;
+ float f = 0.0f;
+#pragma omp target teams distribute parallel for reduction(+ : si) \
+ reduction(- : su) reduction(& : an) reduction(| : orv) reduction(^ : xr) \
+ reduction(min : smn) reduction(max : smx) reduction(min : umn) \
+ reduction(max : umx) reduction(+ : f)
+ for (int i = 0; i < n; ++i) {
+ si += i;
+ su -= i;
+ an &= i;
+ orv |= i;
+ xr ^= i;
+ smn = i < smn ? i : smn;
+ smx = i > smx ? i : smx;
+ umn = (unsigned)i < umn ? (unsigned)i : umn;
+ umx = (unsigned)i > umx ? (unsigned)i : umx;
+ f += i;
+ }
+}
+
+// A fp multiply and a fp max reduction have no direct atomicrmw, so each falls
+// back to the buffer path even with -fopenmp-target-atomic-reduction. Together with
+// the mixed construct below there are exactly three buffered writebacks.
+//
+// BUFFER-COUNT-3: call i32 @__kmpc_gpu_xteam_reduce_nowait(
+// BUFFER-NOT: call i32 @__kmpc_gpu_xteam_reduce_nowait(
+double nonatomicable_mul(int n) {
+ double p = 1.0;
+#pragma omp target teams distribute parallel for reduction(* : p)
+ for (int i = 0; i < n; ++i)
+ p *= 1.0;
+ return p;
+}
+
+double nonatomicable_fpmax(int n) {
+ double mx = 0.0;
+#pragma omp target teams distribute parallel for reduction(max : mx)
+ for (int i = 0; i < n; ++i)
+ mx = (double)i > mx ? (double)i : mx;
+ return mx;
+}
+
+// A single non-atomicable reduction (fp multiply) poisons the whole set: the
+// atomic path requires *all* reductions to have an atomicrmw, so even the int
+// sum here is combined through the buffer, not atomically. Hence no team-main
+// guard and no `atomicrmw add` for `s` from this construct (asserted by the
+// module-wide counts above).
+//
+// CHECK-COUNT-1: call i32 @__kmpc_is_team_main_thread(
+// CHECK-NOT: call i32 @__kmpc_is_team_main_thread(
+// CHECK-NOT: atomicrmw fmul
+// CHECK-NOT: atomicrmw fmax
+
+double mixed(int n) {
+ int s = 0;
+ double p = 1.0;
+#pragma omp target teams distribute parallel for reduction(+ : s) \
+ reduction(* : p)
+ for (int i = 0; i < n; ++i) {
+ s += i;
+ p *= 1.0;
+ }
+ return s + p;
+}
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
index e253c838e5e28..ea1cef97900f5 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -505,6 +505,7 @@ __OMP_RTL(__kmpc_end_sharing_variables, false, Void, )
__OMP_RTL(__kmpc_get_shared_variables, false, Void, VoidPtrPtrPtr)
__OMP_RTL(__kmpc_parallel_level, false, Int16, IdentPtr, Int32)
__OMP_RTL(__kmpc_is_spmd_exec_mode, false, Int8, )
+__OMP_RTL(__kmpc_is_team_main_thread, false, Int32, )
__OMP_RTL(__kmpc_barrier_simple_spmd, false, Void, IdentPtr, Int32)
__OMP_RTL(__kmpc_barrier_simple_generic, false, Void, IdentPtr, Int32)
@@ -1104,6 +1105,7 @@ __OMP_RTL_ATTRS(__kmpc_shuffle_int64, AttributeSet(), AttributeSet(),
ParamAttrs(AttributeSet(), SExt, SExt))
__OMP_RTL_ATTRS(__kmpc_is_spmd_exec_mode, AttributeSet(), SExt, ParamAttrs())
+__OMP_RTL_ATTRS(__kmpc_is_team_main_thread, AttributeSet(), SExt, ParamAttrs())
#undef __OMP_RTL_ATTRS
#undef OMP_RTL_ATTRS
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index fc3812502fb2e..a227397b06f4e 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4650,6 +4650,16 @@ checkReductionInfos(ArrayRef<OpenMPIRBuilder::ReductionInfo> ReductionInfos,
}
}
+// The atomic cross-team reduction fast path applies when every reduction in the
+// set can be represented by an atomicrmw. Clang only populates it for scalar
+// reductions with a supported atomic operator.
+static bool isAtomicableReductionSet(
+ ArrayRef<OpenMPIRBuilder::ReductionInfo> ReductionInfos) {
+ return all_of(ReductionInfos, [](const OpenMPIRBuilder::ReductionInfo &RI) {
+ return static_cast<bool>(RI.AtomicReductionGen);
+ });
+}
+
OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createReductionsGPU(
const LocationDescription &Loc, InsertPointTy AllocaIP,
InsertPointTy CodeGenIP, ArrayRef<ReductionInfo> ReductionInfos,
@@ -4791,6 +4801,9 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createReductionsGPU(
// copied back. (Basically RL, appropriately casted if necessary.)
Value *RLForCopyBack = RL;
+ bool IsAtomicReduction =
+ IsTeamsReduction && isAtomicableReductionSet(ReductionInfos);
+
if (!IsTeamsReduction) {
Value *SarFuncCast =
Builder.CreatePointerBitCastOrAddrSpaceCast(*SarFunc, FuncPtrTy);
@@ -4801,6 +4814,12 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createReductionsGPU(
Function *Pv2Ptr = getOrCreateRuntimeFunctionPtr(
RuntimeFunction::OMPRTL___kmpc_nvptx_parallel_reduce_nowait_v2);
Res = createRuntimeFunctionCall(Pv2Ptr, Args);
+ } else if (IsAtomicReduction) {
+ // Atomic cross-team reduction fast path: determine the team's main thread
+ // that is later to fold its value atomically into the mapped variable.
+ Function *IsMainThreadFn = getOrCreateRuntimeFunctionPtr(
+ RuntimeFunction::OMPRTL___kmpc_is_team_main_thread);
+ Res = createRuntimeFunctionCall(IsMainThreadFn, {});
} else {
CodeGenIP = Builder.saveIP();
StructType *ReductionsBufferTy = StructType::create(
@@ -4939,6 +4958,19 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createReductionsGPU(
// Add emission of __kmpc_end_reduce{_nowait}(<gtid>);
for (auto En : enumerate(ReductionInfos)) {
const ReductionInfo &RI = En.value();
+
+ // Atomic cross-team fast path: each team's main thread folds its
+ // team-reduced value directly into the mapped reduction variable with a
+ // single atomicrmw.
+ if (IsAtomicReduction) {
+ InsertPointOrErrorTy AfterIP = RI.AtomicReductionGen(
+ Builder.saveIP(), RI.ElementType, RI.Variable, RI.PrivateVariable);
+ if (!AfterIP)
+ return AfterIP.takeError();
+ Builder.restoreIP(*AfterIP);
+ continue;
+ }
+
Type *ValueType = RI.ElementType;
Value *RedValue = RI.Variable;
diff --git a/offload/test/offloading/multiple_reductions.cpp b/offload/test/offloading/multiple_reductions.cpp
index 6c759416feaf5..9d993b5f48e00 100644
--- a/offload/test/offloading/multiple_reductions.cpp
+++ b/offload/test/offloading/multiple_reductions.cpp
@@ -1,5 +1,9 @@
// RUN: %libomptarget-compilexx-run-and-check-generic
// RUN: %libomptarget-compileoptxx-run-and-check-generic
+// RUN: %libomptarget-compilexx-generic -fopenmp-target-atomic-reduction && \
+// RUN: %libomptarget-run-generic | %fcheck-generic
+// RUN: %libomptarget-compileoptxx-generic -fopenmp-target-atomic-reduction && \
+// RUN: %libomptarget-run-generic | %fcheck-generic
// UNSUPPORTED: intelgpu
// See array_reductions.cpp for a clone of this file for array reductions.
diff --git a/offload/test/offloading/xteam_atomic_reduction_usm.cpp b/offload/test/offloading/xteam_atomic_reduction_usm.cpp
new file mode 100644
index 0000000000000..ef00e0a2ee6df
--- /dev/null
+++ b/offload/test/offloading/xteam_atomic_reduction_usm.cpp
@@ -0,0 +1,64 @@
+// Validate the atomic cross-team reduction fast path under unified shared
+// memory.
+//
+// RUN: %libomptarget-compilexx-generic -fopenmp-force-usm \
+// RUN: -fopenmp-target-atomic-reduction && env HSA_XNACK=1 \
+// RUN: %libomptarget-run-generic | %fcheck-generic
+// RUN: %libomptarget-compileoptxx-generic -fopenmp-force-usm \
+// RUN: -fopenmp-target-atomic-reduction && env HSA_XNACK=1 \
+// RUN: %libomptarget-run-generic | %fcheck-generic
+//
+// REQUIRES: amdgpu, unified_shared_memory
+
+#include <cassert>
+#include <climits>
+#include <cstdio>
+#include <omp.h>
+
+int main() {
+ const int N = 1 << 16;
+
+ long *sum = (long *)omp_alloc(sizeof(long), llvm_omp_target_shared_mem_alloc);
+ *sum = 7;
+#pragma omp target teams distribute parallel for reduction(+ : sum[0])
+ for (int i = 0; i < N; ++i)
+ sum[0] += i;
+ assert(sum[0] == 7 + (long)(N - 1) * N / 2 && "atomic + reduction incorrect");
+ omp_free(sum, llvm_omp_target_shared_mem_alloc);
+
+ double *fsum =
+ (double *)omp_alloc(sizeof(double), llvm_omp_target_shared_mem_alloc);
+ *fsum = 0.5;
+#pragma omp target teams distribute parallel for reduction(+ : fsum[0])
+ for (int i = 0; i < N; ++i)
+ fsum[0] += 1.0;
+ assert(fsum[0] == 0.5 + (double)N && "atomic fp + reduction incorrect");
+ omp_free(fsum, llvm_omp_target_shared_mem_alloc);
+
+ int *mx = (int *)omp_alloc(sizeof(int), llvm_omp_target_shared_mem_alloc);
+ int *mn = (int *)omp_alloc(sizeof(int), llvm_omp_target_shared_mem_alloc);
+ *mx = INT_MIN;
+ *mn = INT_MAX;
+#pragma omp target teams distribute parallel for reduction(max : mx[0]) \
+ reduction(min : mn[0])
+ for (int i = 0; i < N; ++i) {
+ mx[0] = i > mx[0] ? i : mx[0];
+ mn[0] = i < mn[0] ? i : mn[0];
+ }
+ assert(mx[0] == N - 1 && mn[0] == 0 && "atomic min/max reduction incorrect");
+ omp_free(mx, llvm_omp_target_shared_mem_alloc);
+ omp_free(mn, llvm_omp_target_shared_mem_alloc);
+
+ unsigned *bits =
+ (unsigned *)omp_alloc(sizeof(unsigned), llvm_omp_target_shared_mem_alloc);
+ *bits = 0;
+#pragma omp target teams distribute parallel for reduction(| : bits[0])
+ for (int i = 0; i < N; ++i)
+ bits[0] |= (unsigned)i;
+ assert(bits[0] == (unsigned)(N - 1) && "atomic | reduction incorrect");
+ omp_free(bits, llvm_omp_target_shared_mem_alloc);
+
+ printf("SUCCESS\n");
+ // CHECK: SUCCESS
+ return 0;
+}
diff --git a/openmp/device/include/Interface.h b/openmp/device/include/Interface.h
index cf455bf030270..601694871597b 100644
--- a/openmp/device/include/Interface.h
+++ b/openmp/device/include/Interface.h
@@ -225,6 +225,8 @@ struct KernelEnvironmentTy;
int8_t __kmpc_is_spmd_exec_mode();
+int32_t __kmpc_is_team_main_thread();
+
int32_t __kmpc_target_init(KernelEnvironmentTy &KernelEnvironment,
KernelLaunchEnvironmentTy *KernelLaunchEnvironment);
diff --git a/openmp/device/src/Kernel.cpp b/openmp/device/src/Kernel.cpp
index 600aa0c8528b7..2f40e44496c5e 100644
--- a/openmp/device/src/Kernel.cpp
+++ b/openmp/device/src/Kernel.cpp
@@ -173,4 +173,9 @@ void __kmpc_target_deinit() {
}
int8_t __kmpc_is_spmd_exec_mode() { return mapping::isSPMDMode(); }
+
+/// Whether the calling thread is its team's main thread.
+int32_t __kmpc_is_team_main_thread() {
+ return mapping::isInitialThreadInLevel0(mapping::isSPMDMode());
+}
}
More information about the Openmp-commits
mailing list