[Mlir-commits] [mlir] [mlir][OpenACC] Atomicize contended shared array reduction updates (PR #212971)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 30 04:15:27 PDT 2026
https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/212971
>From bec6d16597b6fce306193fb4b3312d70f1ccc03d Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 29 Jul 2026 22:41:15 -0700
Subject: [PATCH 1/4] [mlir][OpenACC] Honor array reduction storage scope in
ACCCGToGPU
Per-element gpu.all_reduce is only valid for thread_x-private accumulators.
Gang-/worker-scoped array temps (including shared memory) must no-op the
accumulate when block context already holds the partial.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 116 ++++++++++++++----
1 file changed, 92 insertions(+), 24 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 0d30f54f12447..ff5faf3549200 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -777,9 +777,23 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
std::optional<int64_t>
ACCCGToGPULowering::isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
MemRefType baseTy) {
- // Cross-thread array reduction accumulators must stay per-thread.
- if (perThreadArrayReductionAccum(privateLocal.getResult()))
- return std::nullopt;
+ // Cross-thread array reduction accumulators must stay per-thread when their
+ // storage scope includes thread_x. Gang-/worker-scoped array temps remain
+ // eligible for shared memory.
+ if (perThreadArrayReductionAccum(privateLocal.getResult())) {
+ bool storageHasThreadX = true;
+ if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal)) {
+ storageHasThreadX =
+ llvm::any_of(dims.getArray(), [](auto d) { return d.isThreadX(); });
+ } else if (GPUParallelDimsAttr dims =
+ getPrivatizeOp(privateLocal, computeRegion)
+ .getParDimsAttr()) {
+ storageHasThreadX =
+ llvm::any_of(dims.getArray(), [](auto d) { return d.isThreadX(); });
+ }
+ if (storageHasThreadX)
+ return std::nullopt;
+ }
ModuleOp module = computeRegion->getParentOfType<ModuleOp>();
FailureOr<bool> isCandidate = isPrivateLocalSharedMemoryCandidate(
privateLocal, computeRegion, module, defaultPolicy, &accSupport);
@@ -2707,10 +2721,21 @@ void ACCCGToGPULowering::processPrivateLocal(
} else {
// Hoisted acc.privatize: allocate per-thread stack storage in the launch
// body. Cross-thread array reduction accumulators are per-thread too, so
- // the accumulate can reduce each element across threads.
+ // the accumulate can reduce each element across threads. Skip when storage
+ // par_dims lack thread_x (gang-/worker-scoped array temp).
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
- if ((isThreadXPrivatize(privatizeOp) || arrayAccum) &&
+ auto storageHasThreadX = [&]() {
+ if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal))
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ if (GPUParallelDimsAttr dims = privatizeOp.getParDimsAttr())
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ return true;
+ };
+ if ((isThreadXPrivatize(privatizeOp) ||
+ (arrayAccum && storageHasThreadX())) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -2798,8 +2823,18 @@ void ACCCGToGPULowering::processPrivateLocal(
parDimsPair = computeActiveAndInactiveParDims(privateLocal, nullptr);
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
+ auto storageHasThreadX = [&]() {
+ if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal))
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ if (GPUParallelDimsAttr dims =
+ getPrivatizeOp(privateLocal, computeRegion).getParDimsAttr())
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ return true;
+ };
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
- if ((parDim.isThreadX() || arrayAccum) &&
+ if ((parDim.isThreadX() || (arrayAccum && storageHasThreadX())) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -3526,30 +3561,63 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
// Per-element gpu.all_reduce is only correct when each thread owns its own
// accumulator copy. For a statically-shaped accumulator, classify from the
- // operand: an explicit shared allocation is block-shared regardless of size,
- // and anything else (a per-thread stack alloca, or a view over one) is
- // per-thread when it fits the per-thread stack budget and block-shared when
- // it is too large. For a dynamically-shaped accumulator the type conveys no
- // size, so classify from par_dims (which the producer sets to the reduction's
- // actual parallel scope): a thread dimension means per-thread storage.
+ // operand: an explicit shared/heap allocation is block-shared regardless of
+ // size, and a stack alloca (or a view over one) is per-thread when it fits
+ // the per-thread stack budget. Storage `acc.par_dims` without thread_x means
+ // gang-/worker-scoped privacy (shared among vector lanes) even if the type
+ // would fit on the stack. For a dynamically-shaped accumulator the type
+ // conveys no size, so classify from storage/accumulate thread_x dims.
+ auto storageIsThreadXPrivate = [&](Value v) -> bool {
+ if (acc::PrivateLocalOp privateLocal = getPrivateLocalForMemref(v)) {
+ if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal)) {
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ }
+ if (acc::PrivatizeOp privatize =
+ getPrivatizeOp(privateLocal, computeRegion)) {
+ if (GPUParallelDimsAttr dims = privatize.getParDimsAttr()) {
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ }
+ }
+ }
+ if (Operation *root = unwrapMemRefConversion(v).getDefiningOp()) {
+ if (GPUParallelDimsAttr dims = getParDimsAttr(root)) {
+ return llvm::any_of(dims.getArray(),
+ [](auto d) { return d.isThreadX(); });
+ }
+ }
+ return true;
+ };
bool isPerThreadPrivate;
+ Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
+ bool isSharedStorage = isa_and_nonnull<memref::AllocOp>(rootOp) ||
+ isa_and_nonnull<acc::GPUSharedMemoryOp>(rootOp);
+ if (auto addrSpace = dyn_cast_if_present<gpu::AddressSpaceAttr>(
+ memrefTy.getMemorySpace())) {
+ isSharedStorage |=
+ addrSpace.getValue() == gpu::GPUDialect::getWorkgroupAddressSpace();
+ }
if (memrefTy.hasStaticShape()) {
- Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
isPerThreadPrivate =
- !isa_and_nonnull<memref::AllocOp>(rootOp) &&
- canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack);
+ !isSharedStorage &&
+ canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack) &&
+ storageIsThreadXPrivate(op.getMemref());
} else {
- isPerThreadPrivate = llvm::any_of(
- op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) { return d.isThreadX(); });
+ isPerThreadPrivate = !isSharedStorage &&
+ storageIsThreadXPrivate(op.getMemref()) &&
+ llvm::any_of(op.getParDims().getArray(),
+ [](mlir::acc::GPUParallelDimAttr d) {
+ return d.isThreadX();
+ });
}
if (!isPerThreadPrivate) {
- // Block-shared accumulator: no-op only when the accumulate spans a block
- // dim (threads distribute distinct elements, so the block partial is in
- // place and the atomic combine finishes it). A thread-only shared
- // reduction, where several threads reduce into the same element, is not yet
- // supported.
- if (hasBlockDim) {
+ // Block-shared accumulator: no-op when the accumulate spans a block dim or
+ // sits in block context (threads already hold the block partial; the
+ // atomic combine finishes across blocks). A thread-only shared reduction
+ // with no block context, where several threads reduce into the same
+ // element, is not yet supported.
+ if (hasBlockDim || reductionHasBlockContext(op)) {
eraseDeadBounds();
} else {
(void)accSupport.emitNYI(
>From 826ff99064b12ff67a321b047e8c2181020ece2e Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 02:05:09 -0700
Subject: [PATCH 2/4] [mlir][OpenACC] Compact array reduction storage-scope
helpers
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 125 +++++++-----------
1 file changed, 48 insertions(+), 77 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index ff5faf3549200..a1bf283ca7e58 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -382,6 +382,21 @@ static Value castPointerLikeTypeIfNeeded(OpBuilder &builder, Location loc,
/// looking through view/cast ops.
static acc::PrivateLocalOp getPrivateLocalForMemref(Value memref);
+/// Returns the dimensions that own \p privateLocal.
+static GPUParallelDimsAttr
+getPrivateParDims(acc::PrivateLocalOp privateLocal,
+ acc::ComputeRegionOp computeRegion);
+
+/// True when the storage backing \p privateLocal is thread_x-private. An
+/// unknown scope conservatively counts as per-thread.
+static bool storageHasThreadX(acc::PrivateLocalOp privateLocal,
+ acc::ComputeRegionOp computeRegion) {
+ GPUParallelDimsAttr dims = getPrivateParDims(privateLocal, computeRegion);
+ return !dims || llvm::any_of(dims.getArray(), [](GPUParallelDimAttr d) {
+ return d.isThreadX();
+ });
+}
+
/// Returns the sole user of \p v, or null if it has zero or multiple uses.
static Operation *getOnlyUser(Value v) {
if (!v.hasOneUse())
@@ -780,20 +795,9 @@ ACCCGToGPULowering::isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
// Cross-thread array reduction accumulators must stay per-thread when their
// storage scope includes thread_x. Gang-/worker-scoped array temps remain
// eligible for shared memory.
- if (perThreadArrayReductionAccum(privateLocal.getResult())) {
- bool storageHasThreadX = true;
- if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal)) {
- storageHasThreadX =
- llvm::any_of(dims.getArray(), [](auto d) { return d.isThreadX(); });
- } else if (GPUParallelDimsAttr dims =
- getPrivatizeOp(privateLocal, computeRegion)
- .getParDimsAttr()) {
- storageHasThreadX =
- llvm::any_of(dims.getArray(), [](auto d) { return d.isThreadX(); });
- }
- if (storageHasThreadX)
- return std::nullopt;
- }
+ if (perThreadArrayReductionAccum(privateLocal.getResult()) &&
+ storageHasThreadX(privateLocal, computeRegion))
+ return std::nullopt;
ModuleOp module = computeRegion->getParentOfType<ModuleOp>();
FailureOr<bool> isCandidate = isPrivateLocalSharedMemoryCandidate(
privateLocal, computeRegion, module, defaultPolicy, &accSupport);
@@ -1239,13 +1243,14 @@ static bool isRedundantChainAccumulate(acc::ReductionAccumulateOp op) {
return false;
}
-/// Returns the dimensions that own \p privateLocal.
static GPUParallelDimsAttr
getPrivateParDims(acc::PrivateLocalOp privateLocal,
acc::ComputeRegionOp computeRegion) {
if (GPUParallelDimsAttr parDims = acc::getParDimsAttr(privateLocal))
return parDims;
- return getPrivatizeOp(privateLocal, computeRegion).getParDimsAttr();
+ if (acc::PrivatizeOp privatize = getPrivatizeOp(privateLocal, computeRegion))
+ return privatize.getParDimsAttr();
+ return {};
}
/// True when \p privateLocal has one private slot per ThreadY row.
@@ -2725,17 +2730,8 @@ void ACCCGToGPULowering::processPrivateLocal(
// par_dims lack thread_x (gang-/worker-scoped array temp).
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
- auto storageHasThreadX = [&]() {
- if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal))
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- if (GPUParallelDimsAttr dims = privatizeOp.getParDimsAttr())
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- return true;
- };
if ((isThreadXPrivatize(privatizeOp) ||
- (arrayAccum && storageHasThreadX())) &&
+ (arrayAccum && storageHasThreadX(privateLocal, computeRegion))) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -2823,18 +2819,9 @@ void ACCCGToGPULowering::processPrivateLocal(
parDimsPair = computeActiveAndInactiveParDims(privateLocal, nullptr);
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
- auto storageHasThreadX = [&]() {
- if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal))
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- if (GPUParallelDimsAttr dims =
- getPrivatizeOp(privateLocal, computeRegion).getParDimsAttr())
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- return true;
- };
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
- if ((parDim.isThreadX() || (arrayAccum && storageHasThreadX())) &&
+ if ((parDim.isThreadX() ||
+ (arrayAccum && storageHasThreadX(privateLocal, computeRegion))) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -3568,28 +3555,17 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
// would fit on the stack. For a dynamically-shaped accumulator the type
// conveys no size, so classify from storage/accumulate thread_x dims.
auto storageIsThreadXPrivate = [&](Value v) -> bool {
- if (acc::PrivateLocalOp privateLocal = getPrivateLocalForMemref(v)) {
- if (GPUParallelDimsAttr dims = getParDimsAttr(privateLocal)) {
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- }
- if (acc::PrivatizeOp privatize =
- getPrivatizeOp(privateLocal, computeRegion)) {
- if (GPUParallelDimsAttr dims = privatize.getParDimsAttr()) {
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- }
- }
- }
- if (Operation *root = unwrapMemRefConversion(v).getDefiningOp()) {
- if (GPUParallelDimsAttr dims = getParDimsAttr(root)) {
- return llvm::any_of(dims.getArray(),
- [](auto d) { return d.isThreadX(); });
- }
- }
- return true;
+ acc::PrivateLocalOp privateLocal = getPrivateLocalForMemref(v);
+ GPUParallelDimsAttr dims =
+ privateLocal ? getPrivateParDims(privateLocal, computeRegion)
+ : GPUParallelDimsAttr();
+ if (!dims) {
+ if (Operation *root = unwrapMemRefConversion(v).getDefiningOp())
+ dims = getParDimsAttr(root);
+ }
+ return !dims ||
+ llvm::any_of(dims.getArray(), [](auto d) { return d.isThreadX(); });
};
- bool isPerThreadPrivate;
Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
bool isSharedStorage = isa_and_nonnull<memref::AllocOp>(rootOp) ||
isa_and_nonnull<acc::GPUSharedMemoryOp>(rootOp);
@@ -3598,26 +3574,21 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
isSharedStorage |=
addrSpace.getValue() == gpu::GPUDialect::getWorkgroupAddressSpace();
}
- if (memrefTy.hasStaticShape()) {
- isPerThreadPrivate =
- !isSharedStorage &&
- canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack) &&
- storageIsThreadXPrivate(op.getMemref());
- } else {
- isPerThreadPrivate = !isSharedStorage &&
- storageIsThreadXPrivate(op.getMemref()) &&
- llvm::any_of(op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) {
- return d.isThreadX();
- });
- }
+ bool isPerThreadPrivate =
+ !isSharedStorage && storageIsThreadXPrivate(op.getMemref()) &&
+ (memrefTy.hasStaticShape()
+ ? canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack)
+ : llvm::any_of(op.getParDims().getArray(),
+ [](mlir::acc::GPUParallelDimAttr d) {
+ return d.isThreadX();
+ }));
if (!isPerThreadPrivate) {
- // Block-shared accumulator: no-op when the accumulate spans a block dim or
- // sits in block context (threads already hold the block partial; the
- // atomic combine finishes across blocks). A thread-only shared reduction
- // with no block context, where several threads reduce into the same
- // element, is not yet supported.
- if (hasBlockDim || reductionHasBlockContext(op)) {
+ // Block-shared accumulator: no-op when the accumulate has block context,
+ // i.e. it spans a block dim or is nested in a block-mapped loop (threads
+ // already hold the block partial; the atomic combine finishes across
+ // blocks). A thread-only shared reduction with no block context, where
+ // several threads reduce into the same element, is not yet supported.
+ if (reductionHasBlockContext(op)) {
eraseDeadBounds();
} else {
(void)accSupport.emitNYI(
>From b9bbcc6dd27778c4b1885ba4abaede7c1f28b808 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 03:49:35 -0700
Subject: [PATCH 3/4] [mlir][OpenACC] Atomicize contended shared array
reduction updates
When a block-shared array accumulator is updated in place, rewrite
same-element RMW stores to atomic RMW so worker/vector partials are
not lost, while leaving partitioned (thread-varying) indices alone.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 141 ++++++++++++++++--
.../acc-cg-to-gpu-reduction-array-shared.mlir | 59 ++++++++
2 files changed, 189 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index a1bf283ca7e58..d048686dbacd8 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -3493,6 +3493,122 @@ void ACCCGToGPULowering::processAccumulateOp(acc::ReductionAccumulateOp op) {
}
}
+/// True when \p v takes a distinct value per thread: it derives from a thread
+/// id, either directly or through the bounds of an enclosing loop. By this
+/// point a thread-mapped loop carries its mapping in its bounds rather than in
+/// `acc.par_dims`, so the bounds are what must be inspected.
+static bool isThreadVarying(Value v, ArrayRef<Value> threadIds,
+ DenseSet<Value> &visited) {
+ if (!v || !visited.insert(v).second)
+ return false;
+ if (llvm::is_contained(threadIds, v))
+ return true;
+ if (auto arg = dyn_cast<BlockArgument>(v)) {
+ Operation *owner = arg.getOwner()->getParentOp();
+ unsigned dim = arg.getArgNumber();
+ if (auto loop = dyn_cast<scf::ParallelOp>(owner)) {
+ if (dim >= loop.getLowerBound().size())
+ return false;
+ return isThreadVarying(loop.getLowerBound()[dim], threadIds, visited) ||
+ isThreadVarying(loop.getStep()[dim], threadIds, visited);
+ }
+ if (auto loop = dyn_cast<scf::ForOp>(owner))
+ return dim == 0 &&
+ (isThreadVarying(loop.getLowerBound(), threadIds, visited) ||
+ isThreadVarying(loop.getStep(), threadIds, visited));
+ return false;
+ }
+ Operation *def = v.getDefiningOp();
+ if (!def)
+ return false;
+ if (isa<gpu::ThreadIdOp, gpu::LaneIdOp>(def))
+ return true;
+ return llvm::any_of(def->getOperands(), [&](Value o) {
+ return isThreadVarying(o, threadIds, visited);
+ });
+}
+
+/// Strips view and memory-space-cast chains to the underlying buffer.
+static Value accumulatorRoot(Value v) {
+ while (Operation *op = v.getDefiningOp()) {
+ if (auto cast = dyn_cast<memref::MemorySpaceCastOp>(op)) {
+ v = cast.getSource();
+ continue;
+ }
+ if (auto viewLike = dyn_cast<ViewLikeOpInterface>(op)) {
+ if (isa<MemRefType>(viewLike.getViewSource().getType())) {
+ v = viewLike.getViewSource();
+ continue;
+ }
+ }
+ break;
+ }
+ return v;
+}
+
+/// Matches `%l = load %m[%i]` / `%c = combine(%l, %x)` / `store %c, %m[%i]` on
+/// the accumulator \p accum and returns the contributed value `%x`.
+static Value matchAccumulatorUpdate(memref::StoreOp store, Value accum) {
+ if (accumulatorRoot(store.getMemRef()) != accum)
+ return {};
+ Operation *combine = store.getValueToStore().getDefiningOp();
+ if (!combine || combine->getNumOperands() != 2)
+ return {};
+ for (unsigned i = 0; i != 2; ++i) {
+ auto load = combine->getOperand(i).getDefiningOp<memref::LoadOp>();
+ if (!load || accumulatorRoot(load.getMemRef()) != accum)
+ continue;
+ if (!llvm::equal(load.getIndices(), store.getIndices()))
+ continue;
+ return combine->getOperand(1 - i);
+ }
+ return {};
+}
+
+/// A block-shared accumulator is updated in place by the loop body, so several
+/// threads may hit the same element. Make those updates atomic unless the
+/// element index provably varies across the participating threads.
+static void atomicizeSharedAccumulatorUpdates(Value accum,
+ arith::AtomicRMWKind kind,
+ ArrayRef<Value> threadIds,
+ RewriterBase &rewriter) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ SmallVector<memref::StoreOp> stores;
+ SmallVector<Value> worklist{accum};
+ DenseSet<Value> seen;
+ while (!worklist.empty()) {
+ Value cur = worklist.pop_back_val();
+ if (!seen.insert(cur).second)
+ continue;
+ for (Operation *user : cur.getUsers()) {
+ if (auto store = dyn_cast<memref::StoreOp>(user))
+ stores.push_back(store);
+ else if (isa<ViewLikeOpInterface, memref::MemorySpaceCastOp>(user))
+ llvm::append_range(worklist, user->getResults());
+ }
+ }
+
+ for (memref::StoreOp store : stores) {
+ Value contribution = matchAccumulatorUpdate(store, accum);
+ if (!contribution)
+ continue;
+ // A thread-varying index means each thread owns its element, so the
+ // existing plain update is already race-free.
+ if (llvm::any_of(store.getIndices(), [&](Value idx) {
+ DenseSet<Value> visited;
+ return isThreadVarying(idx, threadIds, visited);
+ }))
+ continue;
+ Operation *combine = store.getValueToStore().getDefiningOp();
+ rewriter.setInsertionPoint(store);
+ memref::AtomicRMWOp::create(rewriter, store.getLoc(), kind, contribution,
+ store.getMemRef(), store.getIndices());
+ rewriter.eraseOp(store);
+ if (combine && combine->use_empty())
+ rewriter.eraseOp(combine);
+ }
+}
+
void ACCCGToGPULowering::processAccumulateArrayOp(
acc::ReductionAccumulateArrayOp op) {
LLVM_DEBUG(llvm::dbgs() << "processing accumulate array op: " << *op << "\n");
@@ -3583,17 +3699,20 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
return d.isThreadX();
}));
if (!isPerThreadPrivate) {
- // Block-shared accumulator: no-op when the accumulate has block context,
- // i.e. it spans a block dim or is nested in a block-mapped loop (threads
- // already hold the block partial; the atomic combine finishes across
- // blocks). A thread-only shared reduction with no block context, where
- // several threads reduce into the same element, is not yet supported.
- if (reductionHasBlockContext(op)) {
- eraseDeadBounds();
- } else {
- (void)accSupport.emitNYI(
- loc, "reduction: shared-memory array reduction accumulate");
- }
+ // Block-shared accumulator: the body already updated it in place, so the
+ // block partial is complete and the atomic combine finishes across blocks.
+ // Threads that share an element must not race, so their in-place updates
+ // become atomic.
+ SmallVector<Value> threadIds;
+ for (gpu::Processor proc :
+ {gpu::Processor::ThreadX, gpu::Processor::ThreadY,
+ gpu::Processor::ThreadZ}) {
+ if (Value id = getGPUThreadIdFor(proc))
+ threadIds.push_back(id);
+ }
+ atomicizeSharedAccumulatorUpdates(accumulatorRoot(memref), kind, threadIds,
+ rewriter);
+ eraseDeadBounds();
return;
}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
index aeb4539be83eb..f72ff0164fc41 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
@@ -5,9 +5,12 @@
// all_reduce would scale the block partial by the thread count. The block
// partial is already in place and the atomic combine merges across blocks, so
// the accumulate lowers away to nothing (like the block-only case).
+// Here every thread updates element 0, so the in-place update must be atomic
+// or the threads race and partials are lost.
// CHECK-LABEL: func.func @array_reduction_shared
// CHECK: gpu.launch
+// CHECK: memref.atomic_rmw addi
// CHECK-NOT: gpu.all_reduce
// CHECK-NOT: acc.reduction_accumulate_array
// CHECK-NOT: acc.bounds
@@ -58,3 +61,59 @@ func.func @array_reduction_shared(%arg0: memref<8192xi32>) {
acc.copyout accPtr(%0 : memref<8192xi32>) to varPtr(%arg0 : memref<8192xi32>) {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
return
}
+
+// The same block-shared accumulator, but each thread updates its own element.
+// The update is already race-free, so it must stay a plain store.
+
+// CHECK-LABEL: func.func @array_reduction_shared_partitioned
+// CHECK: gpu.launch
+// CHECK-NOT: memref.atomic_rmw
+// CHECK-NOT: gpu.all_reduce
+// CHECK-NOT: acc.reduction_accumulate_array
+
+func.func @array_reduction_shared_partitioned(%arg0: memref<8192xi32>) {
+ %0 = acc.copyin varPtr(%arg0 : memref<8192xi32>) -> memref<8192xi32> {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
+ acc.kernel_environment dataOperands(%0 : memref<8192xi32>) {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_pw {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg2 = %0) : (memref<8192xi32>) {
+ %c8192 = arith.constant 8192 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %c1 = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %2 = acc.reduction_init %arg2 <add> : memref<8192xi32> {
+ %alloc = memref.alloc() : memref<8192xi32>
+ scf.parallel (%i) = (%c0) to (%c8192) step (%c1) {
+ memref.store %c0_i32, %alloc[%i] : memref<8192xi32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ acc.yield %alloc : memref<8192xi32>
+ }
+ scf.parallel (%bx_iv) = (%c0) to (%kbx) step (%c1) {
+ scf.parallel (%tx_iv) = (%c0) to (%c8192) step (%c1) {
+ %3 = memref.load %2[%tx_iv] : memref<8192xi32>
+ %4 = arith.addi %3, %c1_i32 : i32
+ memref.store %4, %2[%tx_iv] : memref<8192xi32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x]>}
+ %b = acc.bounds extent(%c8192 : index)
+ acc.reduction_accumulate_array %2 bounds(%b) <add> : memref<8192xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.reduction_combine_region %2 into %arg2 : memref<8192xi32> {
+ scf.for %i = %c0 to %c8192 step %c1 {
+ %3 = memref.load %2[%i] : memref<8192xi32>
+ %4 = memref.load %arg2[%i] : memref<8192xi32>
+ %5 = arith.addi %3, %4 : i32
+ memref.store %5, %arg2[%i] : memref<8192xi32>
+ }
+ }
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
+ acc.copyout accPtr(%0 : memref<8192xi32>) to varPtr(%arg0 : memref<8192xi32>) {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
+ return
+}
>From eac5c2496fd4500e61d9b6b80d898425871bf050 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 30 Jul 2026 04:05:36 -0700
Subject: [PATCH 4/4] [test][OpenACC] Cover worker accumulate on gang-scoped
array storage
The gang-storage/worker-accumulate case only exercises ACCCGToGPU, so it
belongs here rather than downstream. Every worker updates the same element,
which must lower to an atomic rather than a plain read-modify-write.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../acc-cg-to-gpu-reduction-array-shared.mlir | 49 +++++++++++++++++++
1 file changed, 49 insertions(+)
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
index f72ff0164fc41..dfd092003b4f4 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array-shared.mlir
@@ -117,3 +117,52 @@ func.func @array_reduction_shared_partitioned(%arg0: memref<8192xi32>) {
acc.copyout accPtr(%0 : memref<8192xi32>) to varPtr(%arg0 : memref<8192xi32>) {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
return
}
+
+// A gang-scoped array private_local (block-only storage dims) feeding a
+// worker-level accumulate must not take the per-thread all_reduce path. Every
+// worker updates element 0 of the gang-private slot, so the update must be
+// atomic.
+
+// CHECK-LABEL: func.func @array_reduction_gang_storage_thread_accum
+// CHECK: gpu.launch
+// CHECK: memref.atomic_rmw addi
+// CHECK-NOT: gpu.all_reduce
+// CHECK-NOT: acc.reduction_accumulate_array
+
+func.func @array_reduction_gang_storage_thread_accum(%arg0: memref<4xi32>) {
+ %0 = acc.copyin varPtr(%arg0 : memref<4xi32>) -> memref<4xi32> {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
+ acc.kernel_environment dataOperands(%0 : memref<4xi32>) {
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %c32 = arith.constant 32 : index
+ %bx = acc.par_width %c2 {par_dim = #acc.par_dim<block_x>}
+ %wy = acc.par_width %c4 {par_dim = #acc.par_dim<thread_y>}
+ %tx = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
+ %private = acc.privatize [#acc<par_dims[block_x]>] : () -> !acc.private_type<memref<4xi32>>
+ acc.compute_region launch(%kbx = %bx, %kwy = %wy, %ktx = %tx) ins(%arg2 = %0, %priv = %private) : (memref<4xi32>, !acc.private_type<memref<4xi32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %c4_idx = arith.constant 4 : index
+ scf.parallel (%bx_iv) = (%c0) to (%kbx) step (%c1) {
+ %local = acc.private_local %priv {acc.par_dims = #acc<par_dims[block_x]>} : (!acc.private_type<memref<4xi32>>) -> memref<4xi32>
+ scf.for %i = %c0 to %c4_idx step %c1 {
+ memref.store %c0_i32, %local[%i] : memref<4xi32>
+ }
+ scf.parallel (%wy_iv) = (%c0) to (%kwy) step (%c1) {
+ %3 = memref.load %local[%c0] : memref<4xi32>
+ %4 = arith.addi %3, %c1_i32 : i32
+ memref.store %4, %local[%c0] : memref<4xi32>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_y]>}
+ %b = acc.bounds extent(%c4_idx : index)
+ acc.reduction_accumulate_array %local bounds(%b) <add> : memref<4xi32> {par_dims = #acc<par_dims[thread_y]>}
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
+ acc.copyout accPtr(%0 : memref<4xi32>) to varPtr(%arg0 : memref<4xi32>) {dataClause = #acc<data_clause acc_reduction>, implicit = true, name = "r"}
+ return
+}
More information about the Mlir-commits
mailing list