[Mlir-commits] [mlir] [mlir][OpenACC] Support multi-rank array reductions (PR #210853)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 20 20:34:09 PDT 2026
https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/210853
>From 2d005d7cf11a32e49f59fafddd6ccc1d39991825 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 17:05:29 -0700
Subject: [PATCH 1/7] [mlir][OpenACC] Support multi-rank array reductions
Initialize static rank-N per-thread accumulators and delinearize flattened reduction indices before accessing them.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 61 +++++++++---
.../acc-cg-to-gpu-reduction-array.mlir | 94 +++++++++++++++++++
2 files changed, 142 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 77fa82ba6bb4e..44ad1f3154174 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -736,17 +736,28 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
MemRefType baseTy,
arith::AtomicRMWKind kind) {
- assert(baseTy.getRank() == 1 && baseTy.hasStaticShape() &&
- "per-thread array reduction accumulator must be static rank-1");
+ assert(baseTy.getRank() > 0 && baseTy.hasStaticShape() &&
+ "per-thread array reduction accumulator must be static ranked");
Value ident = createIdentityValue(b, loc, baseTy.getElementType(), kind,
/*useOnlyFiniteValue=*/true);
Value lb = arith::ConstantIndexOp::create(b, loc, 0);
- Value ub = arith::ConstantIndexOp::create(b, loc, baseTy.getShape()[0]);
Value step = arith::ConstantIndexOp::create(b, loc, 1);
- auto forOp = scf::ForOp::create(b, loc, lb, ub, step);
- OpBuilder::InsertionGuard g(b);
- b.setInsertionPoint(forOp.getBody()->getTerminator());
- memref::StoreOp::create(b, loc, ident, alloca, forOp.getInductionVar());
+ SmallVector<Value> indices;
+ auto buildLoopNest = [&](auto &&self, unsigned dim) -> void {
+ if (dim == baseTy.getRank()) {
+ memref::StoreOp::create(b, loc, ident, alloca, indices);
+ return;
+ }
+
+ Value ub = arith::ConstantIndexOp::create(b, loc, baseTy.getShape()[dim]);
+ auto forOp = scf::ForOp::create(b, loc, lb, ub, step);
+ OpBuilder::InsertionGuard g(b);
+ b.setInsertionPoint(forOp.getBody()->getTerminator());
+ indices.push_back(forOp.getInductionVar());
+ self(self, dim + 1);
+ indices.pop_back();
+ };
+ buildLoopNest(buildLoopNest, 0);
}
std::optional<int64_t>
@@ -3244,10 +3255,19 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
Value memref = mapping.lookupOrDefault(op.getMemref());
MemRefType memrefTy = dyn_cast<MemRefType>(memref.getType());
- if (!memref)
+ if (!memrefTy) {
(void)accSupport.emitNYI(loc, "reduction: non-MemRefTy accumulate array");
- if (memrefTy.getRank() != 1)
- (void)accSupport.emitNYI(loc, "reduction: multi-rank accumulate array");
+ return;
+ }
+ if (memrefTy.getRank() == 0) {
+ (void)accSupport.emitNYI(loc, "reduction: rank-zero accumulate array");
+ return;
+ }
+ if (memrefTy.getRank() > 1 && !memrefTy.hasStaticShape()) {
+ (void)accSupport.emitNYI(loc,
+ "reduction: dynamic multi-rank accumulate array");
+ return;
+ }
FailureOr<arith::AtomicRMWKind> kindOr = getReductionKind(
op.getReductionOperator(), memrefTy.getElementType(), loc);
@@ -3355,9 +3375,24 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(forOp.getBody()->getTerminator());
Value iv = forOp.getInductionVar();
- Value elem = memref::LoadOp::create(rewriter, loc, memref, ValueRange{iv});
- createGPUAllReduceOp(loc, elem, memref, kind, op.getParDims(),
- ValueRange{iv},
+ SmallVector<Value> indices{iv};
+ if (memrefTy.getRank() > 1) {
+ assert(memrefTy.hasStaticShape() &&
+ "multi-rank array reduction accumulator must be static");
+ indices.resize(memrefTy.getRank());
+ Value linearIndex = iv;
+ for (int64_t dim = memrefTy.getRank() - 1; dim >= 0; --dim) {
+ Value dimSize = arith::ConstantIndexOp::create(
+ rewriter, loc, memrefTy.getDimSize(dim));
+ indices[dim] =
+ arith::RemUIOp::create(rewriter, loc, linearIndex, dimSize);
+ if (dim != 0)
+ linearIndex =
+ arith::DivUIOp::create(rewriter, loc, linearIndex, dimSize);
+ }
+ }
+ Value elem = memref::LoadOp::create(rewriter, loc, memref, indices);
+ createGPUAllReduceOp(loc, elem, memref, kind, op.getParDims(), indices,
/*isPerThreadPrivateTarget=*/true);
}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index bfb829c901af3..a008c84108f05 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -134,3 +134,97 @@ func.func @array_reduction_dynamic_par_dims(%buf: memref<?xi32>, %n: index) {
} {origin = "acc.parallel"}
return
}
+
+// CHECK-LABEL: func.func @rank_two_array_reduction
+// CHECK: %[[R2_ALLOCA:.*]] = memref.alloca() : memref<2x3xi32>
+// CHECK: scf.for %[[R2_I:.*]] =
+// CHECK: scf.for %[[R2_J:.*]] =
+// CHECK: memref.store %{{.*}}, %[[R2_ALLOCA]][%[[R2_I]], %[[R2_J]]] : memref<2x3xi32>
+// CHECK: scf.for %[[R2_LINEAR:.*]] =
+// CHECK: %[[R2_J_IDX:.*]] = arith.remui %[[R2_LINEAR]], %{{.*}} : index
+// CHECK: %[[R2_ROW:.*]] = arith.divui %[[R2_LINEAR]], %{{.*}} : index
+// CHECK: %[[R2_I_IDX:.*]] = arith.remui %[[R2_ROW]], %{{.*}} : index
+// CHECK: memref.load %[[R2_ALLOCA]][%[[R2_I_IDX]], %[[R2_J_IDX]]] : memref<2x3xi32>
+
+func.func @rank_two_array_reduction() {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+ %private = acc.privatize [#acc<par_dims[block_x, thread_x]>] : () -> !acc.private_type<memref<2x3xi32>>
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %private) : (!acc.private_type<memref<2x3xi32>>) {
+ %c6 = arith.constant 6 : index
+ %local = acc.private_local %arg0 {acc.par_dims = #acc<par_dims[block_x, thread_x]>} : (!acc.private_type<memref<2x3xi32>>) -> memref<2x3xi32>
+ %bounds = acc.bounds extent(%c6 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add> : memref<2x3xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// CHECK-LABEL: func.func @rank_three_array_reduction
+// CHECK: %[[R3_ALLOCA:.*]] = memref.alloca() : memref<2x2x2xi32>
+// CHECK: scf.for %[[R3_I:.*]] =
+// CHECK: scf.for %[[R3_J:.*]] =
+// CHECK: scf.for %[[R3_K:.*]] =
+// CHECK: memref.store %{{.*}}, %[[R3_ALLOCA]][%[[R3_I]], %[[R3_J]], %[[R3_K]]] : memref<2x2x2xi32>
+// CHECK: scf.for %[[R3_LINEAR:.*]] =
+// CHECK: %[[R3_K_IDX:.*]] = arith.remui %[[R3_LINEAR]], %{{.*}} : index
+// CHECK: %[[R3_PLANE:.*]] = arith.divui %[[R3_LINEAR]], %{{.*}} : index
+// CHECK: %[[R3_J_IDX:.*]] = arith.remui %[[R3_PLANE]], %{{.*}} : index
+// CHECK: %[[R3_ROW:.*]] = arith.divui %[[R3_PLANE]], %{{.*}} : index
+// CHECK: %[[R3_I_IDX:.*]] = arith.remui %[[R3_ROW]], %{{.*}} : index
+// CHECK: memref.load %[[R3_ALLOCA]][%[[R3_I_IDX]], %[[R3_J_IDX]], %[[R3_K_IDX]]] : memref<2x2x2xi32>
+
+func.func @rank_three_array_reduction() {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+ %private = acc.privatize [#acc<par_dims[block_x, thread_x]>] : () -> !acc.private_type<memref<2x2x2xi32>>
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%arg0 = %private) : (!acc.private_type<memref<2x2x2xi32>>) {
+ %c8 = arith.constant 8 : index
+ %local = acc.private_local %arg0 {acc.par_dims = #acc<par_dims[block_x, thread_x]>} : (!acc.private_type<memref<2x2x2xi32>>) -> memref<2x2x2xi32>
+ %bounds = acc.bounds extent(%c8 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add> : memref<2x2x2xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// Unsupported dynamic multi-rank and rank-zero accumulators must not generate
+// invalid indexing operations.
+// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction
+// CHECK: gpu.launch
+// CHECK-NOT: memref.load
+// CHECK: gpu.terminator
+func.func @dynamic_rank_two_array_reduction(%local: memref<?x?xi32>, %extent: index) {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {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(%arg0 = %local, %ext = %extent) : (memref<?x?xi32>, index) {
+ %bounds = acc.bounds extent(%ext : index)
+ acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<?x?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// CHECK-LABEL: func.func @rank_zero_array_reduction
+// CHECK: gpu.launch
+// CHECK-NOT: memref.load
+// CHECK: gpu.terminator
+func.func @rank_zero_array_reduction(%local: memref<i32>) {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {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(%arg0 = %local) : (memref<i32>) {
+ %extent = arith.constant 1 : index
+ %bounds = acc.bounds extent(%extent : index)
+ acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<i32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
>From 1a70c1679ce1c49a2e54fc8ed501ba06fa3d7234 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 17:43:57 -0700
Subject: [PATCH 2/7] [mlir][acc] Use dynamic stack storage for thread-private
arrays
Avoid grid-sized backing allocations for dynamic per-thread storage while preserving shared lowering for oversized static arrays.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 48 ++++++++++++-------
...acc-cg-to-gpu-privatize-threadprivate.mlir | 29 +++++++++++
.../acc-cg-to-gpu-reduction-array.mlir | 41 ++++++++++++----
3 files changed, 90 insertions(+), 28 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 44ad1f3154174..7135ca9177153 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -736,8 +736,8 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
MemRefType baseTy,
arith::AtomicRMWKind kind) {
- assert(baseTy.getRank() > 0 && baseTy.hasStaticShape() &&
- "per-thread array reduction accumulator must be static ranked");
+ assert(baseTy.getRank() > 0 &&
+ "per-thread array reduction accumulator must have positive rank");
Value ident = createIdentityValue(b, loc, baseTy.getElementType(), kind,
/*useOnlyFiniteValue=*/true);
Value lb = arith::ConstantIndexOp::create(b, loc, 0);
@@ -749,7 +749,11 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
return;
}
- Value ub = arith::ConstantIndexOp::create(b, loc, baseTy.getShape()[dim]);
+ Value ub =
+ baseTy.isDynamicDim(dim)
+ ? memref::DimOp::create(b, loc, alloca, dim).getResult()
+ : arith::ConstantIndexOp::create(b, loc, baseTy.getDimSize(dim))
+ .getResult();
auto forOp = scf::ForOp::create(b, loc, lb, ub, step);
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(forOp.getBody()->getTerminator());
@@ -2494,9 +2498,15 @@ void ACCCGToGPULowering::processPrivateLocal(
// the accumulate can reduce each element across threads.
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
- if ((isThreadXPrivatize(privatizeOp) || arrayAccum) &&
- canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
- Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
+ bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
+ bool canUseDynamicAlloca =
+ isThreadPrivate && baseTy.getRank() > 0 && !baseTy.hasStaticShape() &&
+ baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
+ if ((isThreadPrivate || arrayAccum) &&
+ (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
+ canUseDynamicAlloca)) {
+ Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
+ privatizeOp.getDynamicSizes());
if (arrayAccum) {
FailureOr<arith::AtomicRMWKind> kind = getReductionKind(
arrayAccum.getReductionOperator(), baseTy.getElementType(), loc);
@@ -2583,9 +2593,15 @@ void ACCCGToGPULowering::processPrivateLocal(
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
+ bool canUseDynamicAlloca =
+ parDim.isThreadX() && baseTy.getRank() > 0 &&
+ !baseTy.hasStaticShape() &&
+ baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
if ((parDim.isThreadX() || arrayAccum) &&
- canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
- Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
+ (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
+ canUseDynamicAlloca)) {
+ Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
+ privatizeOp.getDynamicSizes());
if (arrayAccum) {
FailureOr<arith::AtomicRMWKind> kind = getReductionKind(
arrayAccum.getReductionOperator(), baseTy.getElementType(), loc);
@@ -3263,12 +3279,6 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
(void)accSupport.emitNYI(loc, "reduction: rank-zero accumulate array");
return;
}
- if (memrefTy.getRank() > 1 && !memrefTy.hasStaticShape()) {
- (void)accSupport.emitNYI(loc,
- "reduction: dynamic multi-rank accumulate array");
- return;
- }
-
FailureOr<arith::AtomicRMWKind> kindOr = getReductionKind(
op.getReductionOperator(), memrefTy.getElementType(), loc);
if (failed(kindOr))
@@ -3377,13 +3387,15 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
Value iv = forOp.getInductionVar();
SmallVector<Value> indices{iv};
if (memrefTy.getRank() > 1) {
- assert(memrefTy.hasStaticShape() &&
- "multi-rank array reduction accumulator must be static");
indices.resize(memrefTy.getRank());
Value linearIndex = iv;
for (int64_t dim = memrefTy.getRank() - 1; dim >= 0; --dim) {
- Value dimSize = arith::ConstantIndexOp::create(
- rewriter, loc, memrefTy.getDimSize(dim));
+ Value dimSize =
+ memrefTy.isDynamicDim(dim)
+ ? memref::DimOp::create(rewriter, loc, memref, dim).getResult()
+ : arith::ConstantIndexOp::create(rewriter, loc,
+ memrefTy.getDimSize(dim))
+ .getResult();
indices[dim] =
arith::RemUIOp::create(rewriter, loc, linearIndex, dimSize);
if (dim != 0)
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
index 281eae4b5b0f7..b186720512991 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
@@ -23,3 +23,32 @@ func.func @threadprivate(%host: memref<i32>) {
} {origin = "acc.parallel"}
return
}
+
+// Dynamic thread-private storage is allocated in the launch body. A
+// grid-wide backing allocation would scale with the logical block count.
+//
+// CHECK-LABEL: func.func @dynamic_threadprivate
+// CHECK: acc.privatize(%{{.*}}) [#acc<par_dims[thread_x]>] : (index) -> !acc.private_type<memref<?xi32>>
+// CHECK: gpu.launch
+// CHECK: memref.alloca(%{{.*}}) : memref<?xi32>
+// CHECK-NOT: acc.unwrap_private
+
+func.func @dynamic_threadprivate(%n: index) {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+ %private = acc.privatize(%n) [#acc<par_dims[thread_x]>]
+ : (index) -> !acc.private_type<memref<?xi32>>
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx)
+ ins(%arg = %private, %extent = %n)
+ : (!acc.private_type<memref<?xi32>>, index) {
+ %c0 = arith.constant 0 : index
+ %c0_i32 = arith.constant 0 : i32
+ %local = acc.private_local %arg
+ : (!acc.private_type<memref<?xi32>>) -> memref<?xi32>
+ memref.store %c0_i32, %local[%c0] : memref<?xi32>
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index a008c84108f05..f287ae054a846 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -192,25 +192,46 @@ func.func @rank_three_array_reduction() {
return
}
-// Unsupported dynamic multi-rank and rank-zero accumulators must not generate
-// invalid indexing operations.
// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction
// CHECK: gpu.launch
-// CHECK-NOT: memref.load
-// CHECK: gpu.terminator
-func.func @dynamic_rank_two_array_reduction(%local: memref<?x?xi32>, %extent: index) {
+// CHECK-NOT: acc.unwrap_private
+// CHECK: %[[DYN_ALLOCA:.*]] = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xi32>
+// CHECK: scf.for %[[DYN_I:.*]] =
+// CHECK: scf.for %[[DYN_J:.*]] =
+// CHECK: memref.store %{{.*}}, %[[DYN_ALLOCA]][%[[DYN_I]], %[[DYN_J]]] : memref<?x?xi32>
+// CHECK: scf.for %[[DYN_LINEAR:.*]] =
+// CHECK: memref.dim %[[DYN_ALLOCA]]
+// CHECK: arith.remui %[[DYN_LINEAR]]
+// CHECK: arith.divui %[[DYN_LINEAR]]
+// CHECK: memref.dim %[[DYN_ALLOCA]]
+// CHECK: memref.load %[[DYN_ALLOCA]][%{{.*}}, %{{.*}}] : memref<?x?xi32>
+// CHECK: gpu.all_reduce add
+func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
+ %extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {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(%arg0 = %local, %ext = %extent) : (memref<?x?xi32>, index) {
- %bounds = acc.bounds extent(%ext : index)
- acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<?x?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
- acc.yield
- } {origin = "acc.parallel"}
+ %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_x]>]
+ : (index, index) -> !acc.private_type<memref<?x?xi32>>
+ acc.kernel_environment {
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx)
+ ins(%arg = %private, %ext = %extent)
+ : (!acc.private_type<memref<?x?xi32>>, index) {
+ %local = acc.private_local %arg
+ {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
+ : (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
+ %bounds = acc.bounds extent(%ext : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ : memref<?x?xi32>
+ {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
return
}
+// Rank-zero accumulators must not generate invalid indexing operations.
// CHECK-LABEL: func.func @rank_zero_array_reduction
// CHECK: gpu.launch
// CHECK-NOT: memref.load
>From 94cd48df3ac887a669b4f7c4cd871226d7014566 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 18:04:45 -0700
Subject: [PATCH 3/7] [mlir][OpenACC] Handle dynamic reductions on all thread
dimensions
Classify dynamic accumulators using every GPU thread dimension so rank-N reductions consistently receive per-thread storage.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 7 +++--
.../acc-cg-to-gpu-reduction-array.mlir | 29 ++++---------------
2 files changed, 9 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 7135ca9177153..e353f70939699 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -2500,7 +2500,8 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
bool canUseDynamicAlloca =
- isThreadPrivate && baseTy.getRank() > 0 && !baseTy.hasStaticShape() &&
+ (isThreadPrivate || arrayAccum) && baseTy.getRank() > 0 &&
+ !baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
if ((isThreadPrivate || arrayAccum) &&
(canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
@@ -2594,7 +2595,7 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
bool canUseDynamicAlloca =
- parDim.isThreadX() && baseTy.getRank() > 0 &&
+ (parDim.isThreadX() || arrayAccum) && baseTy.getRank() > 0 &&
!baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
if ((parDim.isThreadX() || arrayAccum) &&
@@ -3334,7 +3335,7 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
} else {
isPerThreadPrivate = llvm::any_of(
op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) { return d.isThreadX(); });
+ [](mlir::acc::GPUParallelDimAttr d) { return d.isAnyThread(); });
}
if (!isPerThreadPrivate) {
// Block-shared accumulator: no-op only when the accumulate spans a block
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index f287ae054a846..be91cc6762da6 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -211,41 +211,22 @@ func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
%c128 = arith.constant 128 : index
%extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
- %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
- %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_x]>]
+ %ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
+ %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
: (index, index) -> !acc.private_type<memref<?x?xi32>>
acc.kernel_environment {
- acc.compute_region launch(%kbx = %bx, %ktx = %tx)
+ acc.compute_region launch(%kbx = %bx, %kty = %ty)
ins(%arg = %private, %ext = %extent)
: (!acc.private_type<memref<?x?xi32>>, index) {
%local = acc.private_local %arg
- {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
+ {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<?x?xi32>
- {par_dims = #acc<par_dims[block_x, thread_x]>}
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
acc.yield
} {origin = "acc.parallel"}
}
return
}
-
-// Rank-zero accumulators must not generate invalid indexing operations.
-// CHECK-LABEL: func.func @rank_zero_array_reduction
-// CHECK: gpu.launch
-// CHECK-NOT: memref.load
-// CHECK: gpu.terminator
-func.func @rank_zero_array_reduction(%local: memref<i32>) {
- %c1 = arith.constant 1 : index
- %c128 = arith.constant 128 : index
- %bx = acc.par_width %c1 {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(%arg0 = %local) : (memref<i32>) {
- %extent = arith.constant 1 : index
- %bounds = acc.bounds extent(%extent : index)
- acc.reduction_accumulate_array %arg0 bounds(%bounds) <add> : memref<i32> {par_dims = #acc<par_dims[block_x, thread_x]>}
- acc.yield
- } {origin = "acc.parallel"}
- return
-}
>From caf6ec726262faae429717edb8cadc1b6cd00abc Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 18:20:44 -0700
Subject: [PATCH 4/7] [mlir][OpenACC] Find array reductions through descriptor
chains
Trace accumulate operands back to private storage so descriptor construction cannot bypass per-thread allocation and initialization.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 23 ++++++++++++
.../acc-cg-to-gpu-reduction-array.mlir | 35 +++++++++++++++++++
2 files changed, 58 insertions(+)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index e353f70939699..cc3b8910ad765 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -727,6 +727,29 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
worklist.append(user->result_begin(), user->result_end());
}
}
+
+ // Descriptor construction may obscure the forward use chain. Fall back to
+ // tracing each accumulate operand back to this private local.
+ acc::PrivateLocalOp privateLocal = v.getDefiningOp<acc::PrivateLocalOp>();
+ if (!privateLocal)
+ return nullptr;
+ acc::ComputeRegionOp computeRegion =
+ privateLocal->getParentOfType<acc::ComputeRegionOp>();
+ if (!computeRegion)
+ return nullptr;
+ acc::ReductionAccumulateArrayOp result;
+ computeRegion.walk([&](acc::ReductionAccumulateArrayOp accArr) {
+ if (getPrivateLocalForMemref(accArr.getMemref()) != privateLocal)
+ return WalkResult::advance();
+ bool hasThread = llvm::any_of(accArr.getParDims().getArray(),
+ [](auto pd) { return pd.isAnyThread(); });
+ if (!hasThread || !reductionHasBlockContext(accArr))
+ return WalkResult::advance();
+ result = accArr;
+ return WalkResult::interrupt();
+ });
+ if (result)
+ return result;
return nullptr;
}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index be91cc6762da6..ae9ea1e4c95bb 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -230,3 +230,38 @@ func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
}
return
}
+
+// CHECK-LABEL: func.func @descriptor_indirect_dynamic_thread_y
+// CHECK: gpu.launch
+// CHECK-NOT: acc.unwrap_private
+// CHECK: %[[INDIRECT_ALLOCA:.*]] = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xi32>
+// CHECK: scf.for %[[INDIRECT_I:.*]] =
+// CHECK: scf.for %[[INDIRECT_J:.*]] =
+// CHECK: memref.store %{{.*}}, %[[INDIRECT_ALLOCA]][%[[INDIRECT_I]], %[[INDIRECT_J]]] : memref<?x?xi32>
+// CHECK: gpu.all_reduce add
+func.func @descriptor_indirect_dynamic_thread_y(%m: index, %n: index) {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %extent = arith.muli %m, %n : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
+ %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
+ : (index, index) -> !acc.private_type<memref<?x?xi32>>
+ acc.kernel_environment {
+ acc.compute_region launch(%kbx = %bx, %kty = %ty)
+ ins(%arg = %private, %ext = %extent)
+ : (!acc.private_type<memref<?x?xi32>>, index) {
+ %local = acc.private_local %arg
+ {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
+ : (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
+ %descriptor_value = builtin.unrealized_conversion_cast %local
+ : memref<?x?xi32> to memref<?x?xi32>
+ %bounds = acc.bounds extent(%ext : index)
+ acc.reduction_accumulate_array %descriptor_value bounds(%bounds) <add>
+ : memref<?x?xi32>
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
+ return
+}
>From 3790839be0295d1bd7f4ca3bbd151aac1edd5a9a Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 18:34:24 -0700
Subject: [PATCH 5/7] [mlir][OpenACC] Allocate dynamic privates for any thread
dimension
Use privatization parallel dimensions so descriptor indirection cannot force dynamic thread-private backing arrays into grid-sized storage.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 48 +++++++------------
.../acc-cg-to-gpu-reduction-array.mlir | 34 ++++++++++---
2 files changed, 46 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index cc3b8910ad765..99315cad34e1c 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -383,6 +383,14 @@ static bool isThreadXPrivatize(PrivatizeOp privatize) {
return false;
}
+/// True when \p privatize has any thread-level parallelism.
+static bool isThreadPrivatize(PrivatizeOp privatize) {
+ if (GPUParallelDimsAttr parDimsAttr = privatize.getParDimsAttr())
+ return llvm::any_of(parDimsAttr.getArray(),
+ [](GPUParallelDimAttr d) { return d.isAnyThread(); });
+ return false;
+}
+
/// Emits a workgroup-wide GPU barrier.
static void emitGPUBarrierWorkgroup(OpBuilder &builder, Location loc) {
gpu::BarrierOp::create(builder, loc);
@@ -728,28 +736,6 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
}
}
- // Descriptor construction may obscure the forward use chain. Fall back to
- // tracing each accumulate operand back to this private local.
- acc::PrivateLocalOp privateLocal = v.getDefiningOp<acc::PrivateLocalOp>();
- if (!privateLocal)
- return nullptr;
- acc::ComputeRegionOp computeRegion =
- privateLocal->getParentOfType<acc::ComputeRegionOp>();
- if (!computeRegion)
- return nullptr;
- acc::ReductionAccumulateArrayOp result;
- computeRegion.walk([&](acc::ReductionAccumulateArrayOp accArr) {
- if (getPrivateLocalForMemref(accArr.getMemref()) != privateLocal)
- return WalkResult::advance();
- bool hasThread = llvm::any_of(accArr.getParDims().getArray(),
- [](auto pd) { return pd.isAnyThread(); });
- if (!hasThread || !reductionHasBlockContext(accArr))
- return WalkResult::advance();
- result = accArr;
- return WalkResult::interrupt();
- });
- if (result)
- return result;
return nullptr;
}
@@ -2523,12 +2509,13 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
bool canUseDynamicAlloca =
- (isThreadPrivate || arrayAccum) && baseTy.getRank() > 0 &&
+ isThreadPrivatize(privatizeOp) && baseTy.getRank() > 0 &&
!baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
- if ((isThreadPrivate || arrayAccum) &&
- (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
- canUseDynamicAlloca)) {
+ bool canUseStaticAlloca =
+ (isThreadPrivate || arrayAccum) &&
+ canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack);
+ if (canUseStaticAlloca || canUseDynamicAlloca) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
privatizeOp.getDynamicSizes());
if (arrayAccum) {
@@ -2618,12 +2605,13 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
bool canUseDynamicAlloca =
- (parDim.isThreadX() || arrayAccum) && baseTy.getRank() > 0 &&
+ isThreadPrivatize(privatizeOp) && baseTy.getRank() > 0 &&
!baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
- if ((parDim.isThreadX() || arrayAccum) &&
- (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
- canUseDynamicAlloca)) {
+ bool canUseStaticAlloca =
+ (parDim.isThreadX() || arrayAccum) &&
+ canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack);
+ if (canUseStaticAlloca || canUseDynamicAlloca) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
privatizeOp.getDynamicSizes());
if (arrayAccum) {
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index ae9ea1e4c95bb..b9fdf6ecabcb6 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -245,17 +245,39 @@ func.func @descriptor_indirect_dynamic_thread_y(%m: index, %n: index) {
%extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
%ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
- %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
+ %private_data = acc.privatize(%m, %n)
+ [#acc<par_dims[block_x, thread_y]>]
: (index, index) -> !acc.private_type<memref<?x?xi32>>
+ %private_descriptor = acc.privatize [#acc<par_dims[block_x, thread_y]>]
+ : () -> !acc.private_type<memref<1xindex>>
acc.kernel_environment {
acc.compute_region launch(%kbx = %bx, %kty = %ty)
- ins(%arg = %private, %ext = %extent)
- : (!acc.private_type<memref<?x?xi32>>, index) {
- %local = acc.private_local %arg
+ ins(%arg_data = %private_data, %arg_descriptor = %private_descriptor,
+ %ext = %extent)
+ : (!acc.private_type<memref<?x?xi32>>,
+ !acc.private_type<memref<1xindex>>, index) {
+ %c0 = arith.constant 0 : index
+ %c1_k = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ %local_data = acc.private_local %arg_data
{acc.par_dims = #acc<par_dims[block_x, thread_y]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
- %descriptor_value = builtin.unrealized_conversion_cast %local
- : memref<?x?xi32> to memref<?x?xi32>
+ %local_descriptor = acc.private_local %arg_descriptor
+ {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
+ : (!acc.private_type<memref<1xindex>>) -> memref<1xindex>
+ %m_dim = memref.dim %local_data, %c0 : memref<?x?xi32>
+ %n_dim = memref.dim %local_data, %c1_k : memref<?x?xi32>
+ scf.for %i = %c0 to %m_dim step %c1_k {
+ scf.for %j = %c0 to %n_dim step %c1_k {
+ memref.store %c0_i32, %local_data[%i, %j] : memref<?x?xi32>
+ }
+ }
+ %descriptor = builtin.unrealized_conversion_cast %local_data
+ : memref<?x?xi32> to index
+ memref.store %descriptor, %local_descriptor[%c0] : memref<1xindex>
+ %loaded_descriptor = memref.load %local_descriptor[%c0] : memref<1xindex>
+ %descriptor_value = builtin.unrealized_conversion_cast %loaded_descriptor
+ : index to memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
acc.reduction_accumulate_array %descriptor_value bounds(%bounds) <add>
: memref<?x?xi32>
>From affc0c39172f713075ac1912eb631d85b4ce6b56 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 19:10:24 -0700
Subject: [PATCH 6/7] [mlir][OpenACC] Preserve worker-private array storage
Restrict generated stack reductions to thread_x and keep thread_y/thread_z arrays indexed or shared, including explicit dynamic shared memory.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 50 ++++----
.../acc-cg-to-gpu-reduction-array.mlir | 107 +++++++++++-------
2 files changed, 86 insertions(+), 71 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 99315cad34e1c..48e7d13730eeb 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -383,14 +383,6 @@ static bool isThreadXPrivatize(PrivatizeOp privatize) {
return false;
}
-/// True when \p privatize has any thread-level parallelism.
-static bool isThreadPrivatize(PrivatizeOp privatize) {
- if (GPUParallelDimsAttr parDimsAttr = privatize.getParDimsAttr())
- return llvm::any_of(parDimsAttr.getArray(),
- [](GPUParallelDimAttr d) { return d.isAnyThread(); });
- return false;
-}
-
/// Emits a workgroup-wide GPU barrier.
static void emitGPUBarrierWorkgroup(OpBuilder &builder, Location loc) {
gpu::BarrierOp::create(builder, loc);
@@ -709,7 +701,7 @@ static bool reductionHasBlockContext(acc::ReductionAccumulateArrayOp accArr) {
}
/// Returns the array reduction accumulate (through cast/view ops) that \p v
-/// feeds if it needs per-thread storage: its par_dims include a thread dim
+/// feeds if it needs per-thread storage: its par_dims include thread_x
/// and it has block context so the cross-thread all_reduce is well defined.
static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
SmallVector<Value> worklist{v};
@@ -721,10 +713,10 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
for (Operation *user : cur.getUsers()) {
if (acc::ReductionAccumulateArrayOp accArr =
dyn_cast<acc::ReductionAccumulateArrayOp>(user)) {
- bool hasThread = false;
+ bool hasThreadX = false;
for (auto pd : accArr.getParDims().getArray())
- hasThread |= pd.isAnyThread();
- if (hasThread && reductionHasBlockContext(accArr))
+ hasThreadX |= pd.isThreadX();
+ if (hasThreadX && reductionHasBlockContext(accArr))
return accArr;
continue;
}
@@ -735,7 +727,6 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
worklist.append(user->result_begin(), user->result_end());
}
}
-
return nullptr;
}
@@ -2509,13 +2500,11 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
bool canUseDynamicAlloca =
- isThreadPrivatize(privatizeOp) && baseTy.getRank() > 0 &&
- !baseTy.hasStaticShape() &&
+ isThreadPrivate && baseTy.getRank() > 0 && !baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
- bool canUseStaticAlloca =
- (isThreadPrivate || arrayAccum) &&
- canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack);
- if (canUseStaticAlloca || canUseDynamicAlloca) {
+ if ((isThreadPrivate || arrayAccum) &&
+ (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
+ canUseDynamicAlloca)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
privatizeOp.getDynamicSizes());
if (arrayAccum) {
@@ -2605,13 +2594,12 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
bool canUseDynamicAlloca =
- isThreadPrivatize(privatizeOp) && baseTy.getRank() > 0 &&
+ parDim.isThreadX() && baseTy.getRank() > 0 &&
!baseTy.hasStaticShape() &&
baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
- bool canUseStaticAlloca =
- (parDim.isThreadX() || arrayAccum) &&
- canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack);
- if (canUseStaticAlloca || canUseDynamicAlloca) {
+ if ((parDim.isThreadX() || arrayAccum) &&
+ (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
+ canUseDynamicAlloca)) {
Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
privatizeOp.getDynamicSizes());
if (arrayAccum) {
@@ -3337,16 +3325,20 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
// 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.
+ Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
+ bool isExplicitlyShared =
+ isa_and_nonnull<memref::AllocOp, acc::GPUSharedMemoryOp>(rootOp);
bool isPerThreadPrivate;
if (memrefTy.hasStaticShape()) {
- Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
isPerThreadPrivate =
- !isa_and_nonnull<memref::AllocOp>(rootOp) &&
+ !isExplicitlyShared &&
canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack);
} else {
- isPerThreadPrivate = llvm::any_of(
- op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) { return d.isAnyThread(); });
+ isPerThreadPrivate = !isExplicitlyShared &&
+ 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
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index b9fdf6ecabcb6..269e0a1287046 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -84,6 +84,27 @@ func.func @array_reduction_small_shared() {
return
}
+// CHECK-LABEL: func.func @array_reduction_dynamic_shared
+// CHECK: acc.gpu_shared_memory
+// CHECK-NOT: gpu.all_reduce
+func.func @array_reduction_dynamic_shared(%n: index) {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {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(%ext = %n) : (index) {
+ %shared = acc.gpu_shared_memory(%ext)
+ {num_copies = 1 : i64, static_upper_bound_bytes = 4096 : i64}
+ : (index) -> memref<?xi32, #gpu.address_space<workgroup>>
+ %bounds = acc.bounds extent(%ext : index)
+ acc.reduction_accumulate_array %shared bounds(%bounds) <add>
+ : memref<?xi32, #gpu.address_space<workgroup>>
+ {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// CHECK-LABEL: func.func @array_reduction_strided_extent
// CHECK: gpu.launch
// CHECK: %[[LB:.*]] = arith.constant 1 : index
@@ -162,6 +183,33 @@ func.func @rank_two_array_reduction() {
return
}
+// CHECK-LABEL: func.func @rank_two_array_reduction_thread_y
+// CHECK: gpu.launch
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: gpu.all_reduce
+// CHECK: gpu.terminator
+func.func @rank_two_array_reduction_thread_y() {
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
+ %private = acc.privatize [#acc<par_dims[block_x, thread_y]>]
+ : () -> !acc.private_type<memref<2x3xi32>>
+ acc.compute_region launch(%kbx = %bx, %kty = %ty)
+ ins(%arg = %private) : (!acc.private_type<memref<2x3xi32>>) {
+ %c6 = arith.constant 6 : index
+ %local = acc.private_local %arg
+ {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
+ : (!acc.private_type<memref<2x3xi32>>) -> memref<2x3xi32>
+ %bounds = acc.bounds extent(%c6 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ : memref<2x3xi32>
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// CHECK-LABEL: func.func @rank_three_array_reduction
// CHECK: %[[R3_ALLOCA:.*]] = memref.alloca() : memref<2x2x2xi32>
// CHECK: scf.for %[[R3_I:.*]] =
@@ -211,75 +259,50 @@ func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
%c128 = arith.constant 128 : index
%extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
- %ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
- %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
+ %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+ %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_x]>]
: (index, index) -> !acc.private_type<memref<?x?xi32>>
acc.kernel_environment {
- acc.compute_region launch(%kbx = %bx, %kty = %ty)
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx)
ins(%arg = %private, %ext = %extent)
: (!acc.private_type<memref<?x?xi32>>, index) {
%local = acc.private_local %arg
- {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
+ {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<?x?xi32>
- {par_dims = #acc<par_dims[block_x, thread_y]>}
+ {par_dims = #acc<par_dims[block_x, thread_x]>}
acc.yield
} {origin = "acc.parallel"}
}
return
}
-// CHECK-LABEL: func.func @descriptor_indirect_dynamic_thread_y
+// Dynamic worker-level storage must remain indexed by thread_y instead of
+// becoming one stack allocation per physical CUDA thread.
+// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction_thread_y
// CHECK: gpu.launch
-// CHECK-NOT: acc.unwrap_private
-// CHECK: %[[INDIRECT_ALLOCA:.*]] = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xi32>
-// CHECK: scf.for %[[INDIRECT_I:.*]] =
-// CHECK: scf.for %[[INDIRECT_J:.*]] =
-// CHECK: memref.store %{{.*}}, %[[INDIRECT_ALLOCA]][%[[INDIRECT_I]], %[[INDIRECT_J]]] : memref<?x?xi32>
-// CHECK: gpu.all_reduce add
-func.func @descriptor_indirect_dynamic_thread_y(%m: index, %n: index) {
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: gpu.all_reduce
+// CHECK: gpu.terminator
+func.func @dynamic_rank_two_array_reduction_thread_y(%m: index, %n: index) {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
%extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
%ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
- %private_data = acc.privatize(%m, %n)
- [#acc<par_dims[block_x, thread_y]>]
+ %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
: (index, index) -> !acc.private_type<memref<?x?xi32>>
- %private_descriptor = acc.privatize [#acc<par_dims[block_x, thread_y]>]
- : () -> !acc.private_type<memref<1xindex>>
acc.kernel_environment {
acc.compute_region launch(%kbx = %bx, %kty = %ty)
- ins(%arg_data = %private_data, %arg_descriptor = %private_descriptor,
- %ext = %extent)
- : (!acc.private_type<memref<?x?xi32>>,
- !acc.private_type<memref<1xindex>>, index) {
- %c0 = arith.constant 0 : index
- %c1_k = arith.constant 1 : index
- %c0_i32 = arith.constant 0 : i32
- %local_data = acc.private_local %arg_data
+ ins(%arg = %private, %ext = %extent)
+ : (!acc.private_type<memref<?x?xi32>>, index) {
+ %local = acc.private_local %arg
{acc.par_dims = #acc<par_dims[block_x, thread_y]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
- %local_descriptor = acc.private_local %arg_descriptor
- {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
- : (!acc.private_type<memref<1xindex>>) -> memref<1xindex>
- %m_dim = memref.dim %local_data, %c0 : memref<?x?xi32>
- %n_dim = memref.dim %local_data, %c1_k : memref<?x?xi32>
- scf.for %i = %c0 to %m_dim step %c1_k {
- scf.for %j = %c0 to %n_dim step %c1_k {
- memref.store %c0_i32, %local_data[%i, %j] : memref<?x?xi32>
- }
- }
- %descriptor = builtin.unrealized_conversion_cast %local_data
- : memref<?x?xi32> to index
- memref.store %descriptor, %local_descriptor[%c0] : memref<1xindex>
- %loaded_descriptor = memref.load %local_descriptor[%c0] : memref<1xindex>
- %descriptor_value = builtin.unrealized_conversion_cast %loaded_descriptor
- : index to memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
- acc.reduction_accumulate_array %descriptor_value bounds(%bounds) <add>
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<?x?xi32>
{par_dims = #acc<par_dims[block_x, thread_y]>}
acc.yield
>From 3bba5f0a8722c8926daa522949668e420f5ea9aa Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 20:33:43 -0700
Subject: [PATCH 7/7] [mlir][OpenACC] Track array reduction storage provenance
Classify accumulators from their backing allocation through casts so shared storage is not reduced again and private storage is initialized correctly. Keep dynamic stack allocation reduction-specific and honor byte strides.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 71 ++++++++++++-------
...acc-cg-to-gpu-privatize-threadprivate.mlir | 29 --------
.../acc-cg-to-gpu-reduction-array.mlir | 61 ++++++++++------
3 files changed, 82 insertions(+), 79 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 48e7d13730eeb..84243de10830a 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -317,6 +317,17 @@ static bool getPassThroughResults(Operation *userOp, Value trackedOperand,
return false;
}
+ if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(userOp)) {
+ if (castOp->getNumOperands() == 1 && castOp->getNumResults() == 1 &&
+ castOp->getOperand(0) == trackedOperand &&
+ isa<MemRefType>(castOp->getOperand(0).getType()) &&
+ isa<MemRefType>(castOp->getResult(0).getType())) {
+ passThroughResults.push_back(castOp->getResult(0));
+ return true;
+ }
+ return false;
+ }
+
// Partial-entity accesses (e.g. array element or field access) forward the
// base entity through to their results, so treat them as pass-through when
// the base entity is the value being tracked.
@@ -341,6 +352,14 @@ static Value unwrapMemRefConversion(Value v) {
continue;
}
}
+ if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op);
+ castOp && castOp->getNumOperands() == 1 &&
+ castOp->getNumResults() == 1 &&
+ isa<MemRefType>(castOp->getOperand(0).getType()) &&
+ isa<MemRefType>(castOp->getResult(0).getType())) {
+ v = castOp->getOperand(0);
+ continue;
+ }
break;
}
return v;
@@ -730,6 +749,13 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
return nullptr;
}
+static bool
+canUseDynamicReductionAlloca(MemRefType baseTy, acc::PrivatizeOp privatizeOp,
+ acc::ReductionAccumulateArrayOp arrayAccum) {
+ return arrayAccum && !baseTy.hasStaticShape() &&
+ baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
+}
+
/// Store the reduction identity to every element of a freshly allocated
/// per-thread array accumulator so all lanes start from identity (the original
/// init loop may only run on one lane).
@@ -2500,8 +2526,7 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
bool canUseDynamicAlloca =
- isThreadPrivate && baseTy.getRank() > 0 && !baseTy.hasStaticShape() &&
- baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
+ canUseDynamicReductionAlloca(baseTy, privatizeOp, arrayAccum);
if ((isThreadPrivate || arrayAccum) &&
(canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
canUseDynamicAlloca)) {
@@ -2594,9 +2619,7 @@ void ACCCGToGPULowering::processPrivateLocal(
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
bool canUseDynamicAlloca =
- parDim.isThreadX() && baseTy.getRank() > 0 &&
- !baseTy.hasStaticShape() &&
- baseTy.getNumDynamicDims() == privatizeOp.getDynamicSizes().size();
+ canUseDynamicReductionAlloca(baseTy, privatizeOp, arrayAccum);
if ((parDim.isThreadX() || arrayAccum) &&
(canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
canUseDynamicAlloca)) {
@@ -3318,28 +3341,10 @@ 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.
+ // accumulator copy. Classify the operand from its underlying allocation
+ // rather than its shape or par_dims: only memref.alloca is thread-private.
Operation *rootOp = unwrapMemRefConversion(memref).getDefiningOp();
- bool isExplicitlyShared =
- isa_and_nonnull<memref::AllocOp, acc::GPUSharedMemoryOp>(rootOp);
- bool isPerThreadPrivate;
- if (memrefTy.hasStaticShape()) {
- isPerThreadPrivate =
- !isExplicitlyShared &&
- canUseStackAlloca(memrefTy, loc, options.maxThreadPrivateStack);
- } else {
- isPerThreadPrivate = !isExplicitlyShared &&
- llvm::any_of(op.getParDims().getArray(),
- [](mlir::acc::GPUParallelDimAttr d) {
- return d.isThreadX();
- });
- }
+ bool isPerThreadPrivate = isa_and_nonnull<memref::AllocaOp>(rootOp);
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
@@ -3367,7 +3372,19 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
Value lb =
boundsOp.getLowerbound() ? toIndex(boundsOp.getLowerbound()) : zero;
- Value step = boundsOp.getStride() ? toIndex(boundsOp.getStride()) : one;
+ Value step = one;
+ if (Value stride = boundsOp.getStride()) {
+ step = toIndex(stride);
+ if (boundsOp.getStrideInBytes()) {
+ int64_t elementSize =
+ getElementSizeInBytes(loc, memrefTy.getElementType());
+ if (elementSize == 0)
+ return;
+ Value elementSizeValue =
+ arith::ConstantIndexOp::create(rewriter, loc, elementSize);
+ step = arith::DivUIOp::create(rewriter, loc, step, elementSizeValue);
+ }
+ }
// Exclusive upper bound. `extent` counts elements, so the span is
// `extent * step` (for the common unit-stride case step is 1); fall back to
// the inclusive upperbound when no extent is given.
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
index b186720512991..281eae4b5b0f7 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-privatize-threadprivate.mlir
@@ -23,32 +23,3 @@ func.func @threadprivate(%host: memref<i32>) {
} {origin = "acc.parallel"}
return
}
-
-// Dynamic thread-private storage is allocated in the launch body. A
-// grid-wide backing allocation would scale with the logical block count.
-//
-// CHECK-LABEL: func.func @dynamic_threadprivate
-// CHECK: acc.privatize(%{{.*}}) [#acc<par_dims[thread_x]>] : (index) -> !acc.private_type<memref<?xi32>>
-// CHECK: gpu.launch
-// CHECK: memref.alloca(%{{.*}}) : memref<?xi32>
-// CHECK-NOT: acc.unwrap_private
-
-func.func @dynamic_threadprivate(%n: index) {
- %c1 = arith.constant 1 : index
- %c128 = arith.constant 128 : index
- %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
- %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
- %private = acc.privatize(%n) [#acc<par_dims[thread_x]>]
- : (index) -> !acc.private_type<memref<?xi32>>
- acc.compute_region launch(%kbx = %bx, %ktx = %tx)
- ins(%arg = %private, %extent = %n)
- : (!acc.private_type<memref<?xi32>>, index) {
- %c0 = arith.constant 0 : index
- %c0_i32 = arith.constant 0 : i32
- %local = acc.private_local %arg
- : (!acc.private_type<memref<?xi32>>) -> memref<?xi32>
- memref.store %c0_i32, %local[%c0] : memref<?xi32>
- acc.yield
- } {origin = "acc.parallel"}
- return
-}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
index 269e0a1287046..56469f801f026 100644
--- a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
@@ -85,7 +85,8 @@ func.func @array_reduction_small_shared() {
}
// CHECK-LABEL: func.func @array_reduction_dynamic_shared
-// CHECK: acc.gpu_shared_memory
+// CHECK: %[[DYN_SHARED:.*]] = acc.gpu_shared_memory
+// CHECK: builtin.unrealized_conversion_cast %[[DYN_SHARED]]
// CHECK-NOT: gpu.all_reduce
func.func @array_reduction_dynamic_shared(%n: index) {
%c1 = arith.constant 1 : index
@@ -96,8 +97,11 @@ func.func @array_reduction_dynamic_shared(%n: index) {
%shared = acc.gpu_shared_memory(%ext)
{num_copies = 1 : i64, static_upper_bound_bytes = 4096 : i64}
: (index) -> memref<?xi32, #gpu.address_space<workgroup>>
+ %cast = builtin.unrealized_conversion_cast %shared
+ : memref<?xi32, #gpu.address_space<workgroup>>
+ to memref<?xi32, #gpu.address_space<workgroup>>
%bounds = acc.bounds extent(%ext : index)
- acc.reduction_accumulate_array %shared bounds(%bounds) <add>
+ acc.reduction_accumulate_array %cast bounds(%bounds) <add>
: memref<?xi32, #gpu.address_space<workgroup>>
{par_dims = #acc<par_dims[block_x, thread_x]>}
acc.yield
@@ -108,11 +112,13 @@ func.func @array_reduction_dynamic_shared(%n: index) {
// CHECK-LABEL: func.func @array_reduction_strided_extent
// CHECK: gpu.launch
// CHECK: %[[LB:.*]] = arith.constant 1 : index
-// CHECK: %[[STEP:.*]] = arith.constant 2 : index
+// CHECK: %[[BYTE_STEP:.*]] = arith.constant 8 : index
// CHECK: %[[EXTENT:.*]] = arith.constant 3 : index
-// CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index
+// CHECK: %[[ELEMENT_SIZE:.*]] = arith.constant 4 : index
+// CHECK: %[[ELEMENT_STEP:.*]] = arith.divui %[[BYTE_STEP]], %[[ELEMENT_SIZE]] : index
+// CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[ELEMENT_STEP]] : index
// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index
-// CHECK: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[STEP]]
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[ELEMENT_STEP]]
func.func @array_reduction_strided_extent() {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
@@ -120,11 +126,11 @@ func.func @array_reduction_strided_extent() {
%tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
acc.compute_region launch(%kbx = %bx, %ktx = %tx) {
%c1_b = arith.constant 1 : index
- %c2 = arith.constant 2 : index
+ %c8 = arith.constant 8 : index
%c3 = arith.constant 3 : index
%local = memref.alloca() : memref<8xi32>
%bounds = acc.bounds lowerbound(%c1_b : index) extent(%c3 : index)
- stride(%c2 : index)
+ stride(%c8 : index) {strideInBytes = true}
acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<8xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
acc.yield
@@ -132,13 +138,14 @@ func.func @array_reduction_strided_extent() {
return
}
-// A dynamically-shaped accumulator (a strided view whose type conveys no size)
-// is classified per-thread from par_dims: a thread dimension means per-thread
-// storage, so lowering emits the per-element gpu.all_reduce.
+// A dynamically-shaped view over an external buffer is shared even when the
+// accumulate includes thread_x. Storage provenance, not par_dims, controls
+// whether gpu.all_reduce is valid.
//
// CHECK-LABEL: func.func @array_reduction_dynamic_par_dims
-// CHECK: scf.for
-// CHECK: gpu.all_reduce add
+// CHECK: memref.reinterpret_cast
+// CHECK-NOT: gpu.all_reduce
+// CHECK: gpu.terminator
func.func @array_reduction_dynamic_par_dims(%buf: memref<?xi32>, %n: index) {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
@@ -186,6 +193,7 @@ func.func @rank_two_array_reduction() {
// CHECK-LABEL: func.func @rank_two_array_reduction_thread_y
// CHECK: gpu.launch
// CHECK-NOT: memref.alloca
+// CHECK: acc.gpu_shared_memory
// CHECK-NOT: gpu.all_reduce
// CHECK: gpu.terminator
func.func @rank_two_array_reduction_thread_y() {
@@ -247,12 +255,13 @@ func.func @rank_three_array_reduction() {
// CHECK: scf.for %[[DYN_I:.*]] =
// CHECK: scf.for %[[DYN_J:.*]] =
// CHECK: memref.store %{{.*}}, %[[DYN_ALLOCA]][%[[DYN_I]], %[[DYN_J]]] : memref<?x?xi32>
+// CHECK: %[[DYN_DESC:.*]] = builtin.unrealized_conversion_cast %[[DYN_ALLOCA]]
// CHECK: scf.for %[[DYN_LINEAR:.*]] =
-// CHECK: memref.dim %[[DYN_ALLOCA]]
+// CHECK: memref.dim %[[DYN_DESC]]
// CHECK: arith.remui %[[DYN_LINEAR]]
// CHECK: arith.divui %[[DYN_LINEAR]]
-// CHECK: memref.dim %[[DYN_ALLOCA]]
-// CHECK: memref.load %[[DYN_ALLOCA]][%{{.*}}, %{{.*}}] : memref<?x?xi32>
+// CHECK: memref.dim %[[DYN_DESC]]
+// CHECK: memref.load %[[DYN_DESC]][%{{.*}}, %{{.*}}] : memref<?x?xi32>
// CHECK: gpu.all_reduce add
func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
%c1 = arith.constant 1 : index
@@ -269,8 +278,10 @@ func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
%local = acc.private_local %arg
{acc.par_dims = #acc<par_dims[block_x, thread_x]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
+ %descriptor = builtin.unrealized_conversion_cast %local
+ : memref<?x?xi32> to memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
- acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ acc.reduction_accumulate_array %descriptor bounds(%bounds) <add>
: memref<?x?xi32>
{par_dims = #acc<par_dims[block_x, thread_x]>}
acc.yield
@@ -279,32 +290,36 @@ func.func @dynamic_rank_two_array_reduction(%m: index, %n: index) {
return
}
-// Dynamic worker-level storage must remain indexed by thread_y instead of
+// Dynamic worker-level storage must remain indexed by thread_y/thread_z instead
// becoming one stack allocation per physical CUDA thread.
-// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction_thread_y
+// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction_thread_yz
// CHECK: gpu.launch
// CHECK-NOT: memref.alloca
// CHECK-NOT: gpu.all_reduce
+// CHECK: %[[WORKER_STORAGE:.*]] = memref.view %{{.*}}[%{{.*}}][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?xi8> to memref<?x?x?x?x?xi32>
+// CHECK: memref.subview %[[WORKER_STORAGE]][%{{.*}}, %{{.*}}, %{{.*}}, 0, 0] [1, 1, 1, %{{.*}}, %{{.*}}] [1, 1, 1, 1, 1]
// CHECK: gpu.terminator
-func.func @dynamic_rank_two_array_reduction_thread_y(%m: index, %n: index) {
+func.func @dynamic_rank_two_array_reduction_thread_yz(%m: index, %n: index) {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
%extent = arith.muli %m, %n : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
%ty = acc.par_width %c128 {par_dim = #acc.par_dim<thread_y>}
- %private = acc.privatize(%m, %n) [#acc<par_dims[block_x, thread_y]>]
+ %tz = acc.par_width %c128 {par_dim = #acc.par_dim<thread_z>}
+ %private = acc.privatize(%m, %n)
+ [#acc<par_dims[block_x, thread_y, thread_z]>]
: (index, index) -> !acc.private_type<memref<?x?xi32>>
acc.kernel_environment {
- acc.compute_region launch(%kbx = %bx, %kty = %ty)
+ acc.compute_region launch(%kbx = %bx, %kty = %ty, %ktz = %tz)
ins(%arg = %private, %ext = %extent)
: (!acc.private_type<memref<?x?xi32>>, index) {
%local = acc.private_local %arg
- {acc.par_dims = #acc<par_dims[block_x, thread_y]>}
+ {acc.par_dims = #acc<par_dims[block_x, thread_y, thread_z]>}
: (!acc.private_type<memref<?x?xi32>>) -> memref<?x?xi32>
%bounds = acc.bounds extent(%ext : index)
acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<?x?xi32>
- {par_dims = #acc<par_dims[block_x, thread_y]>}
+ {par_dims = #acc<par_dims[block_x, thread_y, thread_z]>}
acc.yield
} {origin = "acc.parallel"}
}
More information about the Mlir-commits
mailing list