[Mlir-commits] [mlir] [mlir][OpenACC] Fix nested array reduction storage dims (PR #212971)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 30 02:29:04 PDT 2026
https://github.com/khaki3 created https://github.com/llvm/llvm-project/pull/212971
Example:
```fortran
!$acc parallel loop gang reduction(+:a)
do i = 1, N
!$acc loop worker reduction(+:a)
do j = 1, M
a(i) = a(i) + b(j,i)
end do
end do
```
In this code, the worker accumulate is `thread_y` while the array temp is gang-scoped. Classifying it as per-thread made `acc.reduction_accumulate_array` emit `gpu.all_reduce` on shared storage and overcount.
Fix: treat only `thread_x` storage as per-thread for array accumulate; for gang-/worker-scoped and shared-memory temps, skip the accumulate (no `gpu.all_reduce`) when block context already holds the partial.
>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/2] [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 0d30f54f124470..ff5faf35492002 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/2] [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 ff5faf35492002..a1bf283ca7e58c 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(
More information about the Mlir-commits
mailing list