[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:37 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openacc
Author: Matsu (khaki3)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/212971.diff
1 Files Affected:
- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp (+69-30)
``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 0d30f54f12447..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())
@@ -777,8 +792,11 @@ 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()))
+ // 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()) &&
+ storageHasThreadX(privateLocal, computeRegion))
return std::nullopt;
ModuleOp module = computeRegion->getParentOfType<ModuleOp>();
FailureOr<bool> isCandidate = isPrivateLocalSharedMemoryCandidate(
@@ -1225,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.
@@ -2707,10 +2726,12 @@ 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) &&
+ if ((isThreadXPrivatize(privatizeOp) ||
+ (arrayAccum && storageHasThreadX(privateLocal, computeRegion))) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -2799,7 +2820,8 @@ void ACCCGToGPULowering::processPrivateLocal(
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
- if ((parDim.isThreadX() || arrayAccum) &&
+ if ((parDim.isThreadX() ||
+ (arrayAccum && storageHasThreadX(privateLocal, computeRegion))) &&
canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
@@ -3526,30 +3548,47 @@ 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.
- bool isPerThreadPrivate;
- if (memrefTy.hasStaticShape()) {
- Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
- isPerThreadPrivate =
- !isa_and_nonnull<memref::AllocOp>(rootOp) &&
- canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack);
- } else {
- isPerThreadPrivate = llvm::any_of(
- op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) { return d.isThreadX(); });
- }
+ // 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 {
+ 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(); });
+ };
+ 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();
+ }
+ 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 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 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(
``````````
</details>
https://github.com/llvm/llvm-project/pull/212971
More information about the Mlir-commits
mailing list