[Mlir-commits] [mlir] [mlir][OpenACC] Support static and dynamic multi-rank OpenACC array reduction accumulators (PR #210853)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 22 06:50:44 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 01/11] [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 02/11] [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 03/11] [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 04/11] [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 05/11] [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 06/11] [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 07/11] [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"}
}
>From 08888a2115abad8c630b5fab31617fcc3e48821e Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 20:57:42 -0700
Subject: [PATCH 08/11] [mlir][OpenACC] Test partial multi-rank reduction
bounds
Cover delinearization with partial bounds and a non-identity memref layout.
---
.../acc-cg-to-gpu-reduction-array.mlir | 32 +++++++++++++++++++
1 file changed, 32 insertions(+)
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 56469f801f026..357f1e14d7a49 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
@@ -325,3 +325,35 @@ func.func @dynamic_rank_two_array_reduction_thread_yz(%m: index, %n: index) {
}
return
}
+
+// CHECK-LABEL: func.func @rank_two_partial_bounds_strided_layout
+// CHECK: %[[LB:.*]] = arith.constant 5 : index
+// CHECK: %[[STEP:.*]] = arith.constant 2 : index
+// CHECK: %[[EXTENT:.*]] = arith.constant 3 : index
+// CHECK: %[[LAYOUT_ALLOCA:.*]] = memref.alloca() : memref<3x4xi32, strided<[8, 2]>>
+// CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index
+// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index
+// CHECK: scf.for %[[LINEAR:.*]] = %[[LB]] to %[[UB]] step %[[STEP]]
+// CHECK: %[[COL:.*]] = arith.remui %[[LINEAR]], %{{.*}} : index
+// CHECK: %[[ROW_LINEAR:.*]] = arith.divui %[[LINEAR]], %{{.*}} : index
+// CHECK: %[[ROW:.*]] = arith.remui %[[ROW_LINEAR]], %{{.*}} : index
+// CHECK: memref.load %[[LAYOUT_ALLOCA]][%[[ROW]], %[[COL]]] : memref<3x4xi32, strided<[8, 2]>>
+func.func @rank_two_partial_bounds_strided_layout() {
+ %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) {
+ %lb = arith.constant 5 : index
+ %step = arith.constant 2 : index
+ %extent = arith.constant 3 : index
+ %local = memref.alloca() : memref<3x4xi32, strided<[8, 2]>>
+ %bounds = acc.bounds lowerbound(%lb : index) extent(%extent : index)
+ stride(%step : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ : memref<3x4xi32, strided<[8, 2]>>
+ {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
>From 35c5c495bad0d3e902e6b32cdb1a3365bd378b72 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Mon, 20 Jul 2026 22:44:28 -0700
Subject: [PATCH 09/11] [mlir][OpenACC] Preserve descriptor-backed reduction
provenance
Trace unique stores and loads so descriptor-backed thread-private accumulators are initialized and reduced correctly without re-reducing indexed shared storage.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 155 +++++++++++++++++-
.../acc-cg-to-gpu-reduction-array.mlir | 29 ++++
2 files changed, 181 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 84243de10830a..61ba40247d8b0 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -339,6 +339,119 @@ static bool getPassThroughResults(Operation *userOp, Value trackedOperand,
}
return false;
}
+
+ bool hasMemoryEffects = false;
+ if (auto effects = dyn_cast<MemoryEffectOpInterface>(userOp)) {
+ SmallVector<MemoryEffects::EffectInstance> instances;
+ effects.getEffects(instances);
+ hasMemoryEffects = !instances.empty();
+ }
+ if (userOp->getNumResults() == 1 &&
+ llvm::is_contained(userOp->getOperands(), trackedOperand) &&
+ !hasMemoryEffects) {
+ passThroughResults.push_back(userOp->getResult(0));
+ return true;
+ }
+ return false;
+}
+
+template <typename Effect>
+static bool hasEffectOnValue(Operation *op, Value value) {
+ auto effects = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!effects)
+ return false;
+ SmallVector<MemoryEffects::EffectInstance> instances;
+ effects.getEffects(instances);
+ return llvm::any_of(instances, [&](const auto &instance) {
+ return isa<Effect>(instance.getEffect()) && instance.getValue() == value;
+ });
+}
+
+static Value getReadAddress(Operation *op) {
+ auto effects = dyn_cast<MemoryEffectOpInterface>(op);
+ if (!effects)
+ return {};
+ SmallVector<MemoryEffects::EffectInstance> instances;
+ effects.getEffects(instances);
+ Value address;
+ for (const auto &instance : instances) {
+ Value value = instance.getValue();
+ if (!isa<MemoryEffects::Read>(instance.getEffect()) || !value)
+ continue;
+ if (address && address != value)
+ return {};
+ address = value;
+ }
+ return address;
+}
+
+/// Returns the unique pointer-like value stored to \p address before \p before.
+static Value getUniqueStoredPointerValue(Value address, Operation *before) {
+ Operation *storedBy = nullptr;
+ Value stored;
+ SmallVector<Value> worklist{address};
+ DenseSet<Value> seen;
+ while (!worklist.empty()) {
+ Value alias = worklist.pop_back_val();
+ if (!seen.insert(alias).second)
+ continue;
+ for (Operation *user : alias.getUsers()) {
+ if (hasEffectOnValue<MemoryEffects::Write>(user, alias)) {
+ if (storedBy || user->getBlock() != before->getBlock() ||
+ !user->isBeforeInBlock(before))
+ return {};
+ Value candidate;
+ for (Value operand : user->getOperands()) {
+ if (operand == alias)
+ continue;
+ if (candidate)
+ return {};
+ candidate = operand;
+ }
+ if (!candidate)
+ return {};
+ storedBy = user;
+ stored = candidate;
+ continue;
+ }
+ SmallVector<Value> through;
+ if (getPassThroughResults(user, alias, through))
+ worklist.append(through.begin(), through.end());
+ }
+ }
+ return stored;
+}
+
+/// Tracks a uniquely stored pointer-like value through its subsequent load.
+static bool
+getMemoryPassThroughResults(Operation *userOp, Value trackedOperand,
+ SmallVectorImpl<Value> &passThroughResults) {
+ if (hasEffectOnValue<MemoryEffects::Read>(userOp, trackedOperand) &&
+ userOp->getNumResults() == 1) {
+ passThroughResults.push_back(userOp->getResult(0));
+ return true;
+ }
+
+ if (!llvm::is_contained(userOp->getOperands(), trackedOperand))
+ return false;
+ auto effects = dyn_cast<MemoryEffectOpInterface>(userOp);
+ if (!effects)
+ return false;
+ SmallVector<MemoryEffects::EffectInstance> instances;
+ effects.getEffects(instances);
+ for (const auto &instance : instances) {
+ Value target = instance.getValue();
+ if (!isa<MemoryEffects::Write>(instance.getEffect()) || !target ||
+ target == trackedOperand)
+ continue;
+ unsigned writers = llvm::count_if(target.getUsers(), [&](Operation *user) {
+ return hasEffectOnValue<MemoryEffects::Write>(user, target);
+ });
+ if (writers != 1)
+ return false;
+ passThroughResults.push_back(target);
+ return true;
+ }
return false;
}
@@ -742,6 +855,8 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
SmallVector<Value> through;
if (getPassThroughResults(user, cur, through))
worklist.append(through.begin(), through.end());
+ else if (getMemoryPassThroughResults(user, cur, through))
+ worklist.append(through.begin(), through.end());
else if (isa<ViewLikeOpInterface>(user))
worklist.append(user->result_begin(), user->result_end());
}
@@ -790,6 +905,41 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
buildLoopNest(buildLoopNest, 0);
}
+/// True when \p memref resolves to storage owned by the current thread.
+static bool isPerThreadArrayReductionStorage(Value memref) {
+ DenseSet<Value> seen;
+ SmallVector<Value> worklist{memref};
+ while (!worklist.empty()) {
+ Value current = worklist.pop_back_val();
+ if (!current || !seen.insert(current).second)
+ continue;
+ Operation *def = current.getDefiningOp();
+ if (!def)
+ continue;
+ if (isa<memref::AllocaOp>(def))
+ return true;
+ if (isa<acc::UnwrapPrivateOp, acc::GPUSharedMemoryOp>(def))
+ return false;
+
+ if (Value address = getReadAddress(def)) {
+ Value stored = getUniqueStoredPointerValue(address, def);
+ if (!stored)
+ return false;
+ worklist.push_back(stored);
+ continue;
+ }
+
+ if (mlir::acc::GPUParallelDimsAttr parDims =
+ mlir::acc::getParDimsAttr(def)) {
+ return llvm::any_of(parDims.getArray(),
+ [](auto dim) { return dim.isThreadX(); });
+ }
+
+ worklist.append(def->getOperands().begin(), def->getOperands().end());
+ }
+ return false;
+}
+
std::optional<int64_t>
ACCCGToGPULowering::isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
MemRefType baseTy) {
@@ -3342,9 +3492,8 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
// Per-element gpu.all_reduce is only correct when each thread owns its own
// 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 isPerThreadPrivate = isa_and_nonnull<memref::AllocaOp>(rootOp);
+ // rather than its shape or the accumulate's par_dims.
+ bool isPerThreadPrivate = isPerThreadArrayReductionStorage(memref);
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
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 357f1e14d7a49..1573e0cb76f2a 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,35 @@ func.func @array_reduction_small_shared() {
return
}
+// A reduction private that exceeds the per-thread stack budget uses indexed
+// backing storage shared by thread_x lanes. It must not be reduced again.
+// CHECK-LABEL: func.func @array_reduction_large_indexed_private
+// CHECK: acc.unwrap_private
+// CHECK: memref.view
+// CHECK: memref.subview
+// CHECK-NOT: gpu.all_reduce
+// CHECK-NOT: acc.reduction_accumulate_array
+func.func @array_reduction_large_indexed_private() {
+ %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<8192xi32>>
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx)
+ ins(%arg0 = %private) : (!acc.private_type<memref<8192xi32>>) {
+ %c8192 = arith.constant 8192 : index
+ %local = acc.private_local %arg0
+ {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
+ : (!acc.private_type<memref<8192xi32>>) -> memref<8192xi32>
+ %bounds = acc.bounds extent(%c8192 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ : memref<8192xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// CHECK-LABEL: func.func @array_reduction_dynamic_shared
// CHECK: %[[DYN_SHARED:.*]] = acc.gpu_shared_memory
// CHECK: builtin.unrealized_conversion_cast %[[DYN_SHARED]]
>From 46f3341a39dfe96de4e1865dca58ffd94086f890 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Tue, 21 Jul 2026 06:12:10 -0700
Subject: [PATCH 10/11] [mlir][OpenACC] Harden reduction storage provenance
Use explicit storage-preserving interfaces and reject ambiguous aliases so descriptor-backed reductions are classified conservatively.
---
.../Dialect/OpenACC/OpenACCOpsInterfaces.td | 15 +++
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 115 +++++++++---------
.../acc-cg-to-gpu-reduction-array.mlir | 52 ++++++++
3 files changed, 126 insertions(+), 56 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
index eddd21a7e917b..33bd833ef9737 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
@@ -45,6 +45,21 @@ def PartialEntityAccessOpInterface : OpInterface<"PartialEntityAccessOpInterface
];
}
+def UnderlyingStorageOpInterface : OpInterface<"UnderlyingStorageOpInterface"> {
+ let cppNamespace = "::mlir::acc";
+
+ let description = [{
+ An interface for operations that derive a pointer-like result from one
+ underlying storage value. Other operands may describe metadata such as
+ bounds, offsets, or strides and are not storage sources.
+ }];
+
+ let methods = [
+ InterfaceMethod<"Get the underlying storage value", "::mlir::Value",
+ "getStorageSource", (ins)>,
+ ];
+}
+
def AddressOfGlobalOpInterface : OpInterface<"AddressOfGlobalOpInterface"> {
let cppNamespace = "::mlir::acc";
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 61ba40247d8b0..845b913e81e90 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -306,53 +306,29 @@ static bool sameEffectiveValue(Value x, int64_t y) {
return *conX == y;
}
-/// Continues tracking a memref through view-like and partial-access ops.
+/// Returns the operand that preserves the storage represented by \p op.
+static Value getStorageSource(Operation *op) {
+ if (ViewLikeOpInterface viewLike = dyn_cast<ViewLikeOpInterface>(op))
+ return viewLike.getViewSource();
+ if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op);
+ castOp && castOp->getNumOperands() == 1 && castOp->getNumResults() == 1)
+ return castOp->getOperand(0);
+ if (acc::PartialEntityAccessOpInterface partialAccess =
+ dyn_cast<acc::PartialEntityAccessOpInterface>(op))
+ return partialAccess.getBaseEntity();
+ if (acc::UnderlyingStorageOpInterface storage =
+ dyn_cast<acc::UnderlyingStorageOpInterface>(op))
+ return storage.getStorageSource();
+ return {};
+}
+
+/// Continues tracking a storage value through preserving operations.
static bool getPassThroughResults(Operation *userOp, Value trackedOperand,
SmallVectorImpl<Value> &passThroughResults) {
- if (ViewLikeOpInterface viewLikeOp = dyn_cast<ViewLikeOpInterface>(userOp)) {
- if (viewLikeOp.getViewSource() == trackedOperand) {
- passThroughResults.push_back(viewLikeOp.getViewDest());
- return true;
- }
+ if (getStorageSource(userOp) != 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.
- if (acc::PartialEntityAccessOpInterface partialAccess =
- dyn_cast<acc::PartialEntityAccessOpInterface>(userOp)) {
- if (partialAccess.getBaseEntity() == trackedOperand) {
- passThroughResults.append(userOp->result_begin(), userOp->result_end());
- return true;
- }
- return false;
- }
-
- bool hasMemoryEffects = false;
- if (auto effects = dyn_cast<MemoryEffectOpInterface>(userOp)) {
- SmallVector<MemoryEffects::EffectInstance> instances;
- effects.getEffects(instances);
- hasMemoryEffects = !instances.empty();
- }
- if (userOp->getNumResults() == 1 &&
- llvm::is_contained(userOp->getOperands(), trackedOperand) &&
- !hasMemoryEffects) {
- passThroughResults.push_back(userOp->getResult(0));
- return true;
- }
- return false;
+ passThroughResults.append(userOp->result_begin(), userOp->result_end());
+ return true;
}
template <typename Effect>
@@ -385,8 +361,39 @@ static Value getReadAddress(Operation *op) {
return address;
}
+static bool hasAllocaRoot(Value value) {
+ DenseSet<Value> seen;
+ Value current = value;
+ while (current && seen.insert(current).second) {
+ Operation *def = current.getDefiningOp();
+ if (!def)
+ return false;
+ if (isa<memref::AllocaOp>(def))
+ return true;
+ current = getStorageSource(def);
+ }
+ return false;
+}
+
/// Returns the unique pointer-like value stored to \p address before \p before.
static Value getUniqueStoredPointerValue(Value address, Operation *before) {
+ Operation *scope = before->getParentOfType<acc::ComputeRegionOp>();
+ DominanceInfo dominance(scope ? scope : before->getParentOp());
+ auto safelyPrecedes = [&](Operation *store) {
+ if (dominance.dominates(store, before))
+ return true;
+ auto predicate = store->getParentOfType<acc::PredicateRegionOp>();
+ auto privateLocal = address.getDefiningOp<acc::PrivateLocalOp>();
+ GPUParallelDimsAttr parDims;
+ if (privateLocal)
+ parDims = mlir::acc::getParDimsAttr(privateLocal);
+ bool isThreadPrivate =
+ parDims && llvm::any_of(parDims.getArray(),
+ [](auto dim) { return dim.isAnyThread(); });
+ return predicate && store->getParentOp() == predicate.getOperation() &&
+ dominance.dominates(predicate, before) &&
+ (isThreadPrivate || hasAllocaRoot(address));
+ };
Operation *storedBy = nullptr;
Value stored;
SmallVector<Value> worklist{address};
@@ -397,8 +404,7 @@ static Value getUniqueStoredPointerValue(Value address, Operation *before) {
continue;
for (Operation *user : alias.getUsers()) {
if (hasEffectOnValue<MemoryEffects::Write>(user, alias)) {
- if (storedBy || user->getBlock() != before->getBlock() ||
- !user->isBeforeInBlock(before))
+ if (storedBy || !safelyPrecedes(user))
return {};
Value candidate;
for (Value operand : user->getOperands()) {
@@ -417,6 +423,8 @@ static Value getUniqueStoredPointerValue(Value address, Operation *before) {
SmallVector<Value> through;
if (getPassThroughResults(user, alias, through))
worklist.append(through.begin(), through.end());
+ else if (!hasEffectOnValue<MemoryEffects::Read>(user, alias))
+ return {};
}
}
return stored;
@@ -857,8 +865,6 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
worklist.append(through.begin(), through.end());
else if (getMemoryPassThroughResults(user, cur, through))
worklist.append(through.begin(), through.end());
- else if (isa<ViewLikeOpInterface>(user))
- worklist.append(user->result_begin(), user->result_end());
}
}
return nullptr;
@@ -908,14 +914,11 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
/// True when \p memref resolves to storage owned by the current thread.
static bool isPerThreadArrayReductionStorage(Value memref) {
DenseSet<Value> seen;
- SmallVector<Value> worklist{memref};
- while (!worklist.empty()) {
- Value current = worklist.pop_back_val();
- if (!current || !seen.insert(current).second)
- continue;
+ Value current = memref;
+ while (current && seen.insert(current).second) {
Operation *def = current.getDefiningOp();
if (!def)
- continue;
+ return false;
if (isa<memref::AllocaOp>(def))
return true;
if (isa<acc::UnwrapPrivateOp, acc::GPUSharedMemoryOp>(def))
@@ -925,7 +928,7 @@ static bool isPerThreadArrayReductionStorage(Value memref) {
Value stored = getUniqueStoredPointerValue(address, def);
if (!stored)
return false;
- worklist.push_back(stored);
+ current = stored;
continue;
}
@@ -935,7 +938,7 @@ static bool isPerThreadArrayReductionStorage(Value memref) {
[](auto dim) { return dim.isThreadX(); });
}
- worklist.append(def->getOperands().begin(), def->getOperands().end());
+ current = getStorageSource(def);
}
return false;
}
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 1573e0cb76f2a..fd3bb1e5240f0 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,58 @@ func.func @array_reduction_small_shared() {
return
}
+// Dynamic allocation metadata must not be mistaken for storage provenance.
+// CHECK-LABEL: func.func @array_reduction_heap_alloca_extent
+// CHECK: memref.alloca() : memref<index>
+// CHECK: %[[EXTENT:.*]] = memref.load
+// CHECK: memref.alloc(%[[EXTENT]]) : memref<?xi32>
+// CHECK-NOT: gpu.all_reduce
+// CHECK: gpu.terminator
+func.func @array_reduction_heap_alloca_extent(%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(%extent = %n) : (index) {
+ %slot = memref.alloca() : memref<index>
+ memref.store %extent, %slot[] : memref<index>
+ %loaded_extent = memref.load %slot[] : memref<index>
+ %heap = memref.alloc(%loaded_extent) : memref<?xi32>
+ %bounds = acc.bounds extent(%loaded_extent : index)
+ acc.reduction_accumulate_array %heap bounds(%bounds) <add>
+ : memref<?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ memref.dealloc %heap : memref<?xi32>
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// Non-storage operands of pure operations must not affect classification.
+// CHECK-LABEL: func.func @array_reduction_shared_select
+// CHECK: arith.select
+// CHECK-NOT: gpu.all_reduce
+// CHECK: gpu.terminator
+func.func @array_reduction_shared_select(%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(%extent = %n) : (index) {
+ %true = arith.constant true
+ %condition_slot = memref.alloca() : memref<i1>
+ memref.store %true, %condition_slot[] : memref<i1>
+ %condition = memref.load %condition_slot[] : memref<i1>
+ %heap = memref.alloc(%extent) : memref<?xi32>
+ %selected = arith.select %condition, %heap, %heap : memref<?xi32>
+ %bounds = acc.bounds extent(%extent : index)
+ acc.reduction_accumulate_array %selected bounds(%bounds) <add>
+ : memref<?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ memref.dealloc %heap : memref<?xi32>
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// A reduction private that exceeds the per-thread stack budget uses indexed
// backing storage shared by thread_x lanes. It must not be reduced again.
// CHECK-LABEL: func.func @array_reduction_large_indexed_private
>From 63348cf99dcdd04d151fd312cb6e6c42d61273fe Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 22 Jul 2026 06:49:56 -0700
Subject: [PATCH 11/11] [mlir][OpenACC] Keep array reduction lowering
mechanical
---
.../Dialect/OpenACC/OpenACCOpsInterfaces.td | 15 -
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 297 ++++--------------
.../acc-cg-to-gpu-reduction-array.mlir | 244 ++------------
3 files changed, 80 insertions(+), 476 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
index 33bd833ef9737..eddd21a7e917b 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
@@ -45,21 +45,6 @@ def PartialEntityAccessOpInterface : OpInterface<"PartialEntityAccessOpInterface
];
}
-def UnderlyingStorageOpInterface : OpInterface<"UnderlyingStorageOpInterface"> {
- let cppNamespace = "::mlir::acc";
-
- let description = [{
- An interface for operations that derive a pointer-like result from one
- underlying storage value. Other operands may describe metadata such as
- bounds, offsets, or strides and are not storage sources.
- }];
-
- let methods = [
- InterfaceMethod<"Get the underlying storage value", "::mlir::Value",
- "getStorageSource", (ins)>,
- ];
-}
-
def AddressOfGlobalOpInterface : OpInterface<"AddressOfGlobalOpInterface"> {
let cppNamespace = "::mlir::acc";
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 845b913e81e90..f0fd7ae96ad4f 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -306,159 +306,27 @@ static bool sameEffectiveValue(Value x, int64_t y) {
return *conX == y;
}
-/// Returns the operand that preserves the storage represented by \p op.
-static Value getStorageSource(Operation *op) {
- if (ViewLikeOpInterface viewLike = dyn_cast<ViewLikeOpInterface>(op))
- return viewLike.getViewSource();
- if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op);
- castOp && castOp->getNumOperands() == 1 && castOp->getNumResults() == 1)
- return castOp->getOperand(0);
- if (acc::PartialEntityAccessOpInterface partialAccess =
- dyn_cast<acc::PartialEntityAccessOpInterface>(op))
- return partialAccess.getBaseEntity();
- if (acc::UnderlyingStorageOpInterface storage =
- dyn_cast<acc::UnderlyingStorageOpInterface>(op))
- return storage.getStorageSource();
- return {};
-}
-
-/// Continues tracking a storage value through preserving operations.
+/// Continues tracking a memref through view-like and partial-access ops.
static bool getPassThroughResults(Operation *userOp, Value trackedOperand,
SmallVectorImpl<Value> &passThroughResults) {
- if (getStorageSource(userOp) != trackedOperand)
- return false;
- passThroughResults.append(userOp->result_begin(), userOp->result_end());
- return true;
-}
-
-template <typename Effect>
-static bool hasEffectOnValue(Operation *op, Value value) {
- auto effects = dyn_cast<MemoryEffectOpInterface>(op);
- if (!effects)
- return false;
- SmallVector<MemoryEffects::EffectInstance> instances;
- effects.getEffects(instances);
- return llvm::any_of(instances, [&](const auto &instance) {
- return isa<Effect>(instance.getEffect()) && instance.getValue() == value;
- });
-}
-
-static Value getReadAddress(Operation *op) {
- auto effects = dyn_cast<MemoryEffectOpInterface>(op);
- if (!effects)
- return {};
- SmallVector<MemoryEffects::EffectInstance> instances;
- effects.getEffects(instances);
- Value address;
- for (const auto &instance : instances) {
- Value value = instance.getValue();
- if (!isa<MemoryEffects::Read>(instance.getEffect()) || !value)
- continue;
- if (address && address != value)
- return {};
- address = value;
- }
- return address;
-}
-
-static bool hasAllocaRoot(Value value) {
- DenseSet<Value> seen;
- Value current = value;
- while (current && seen.insert(current).second) {
- Operation *def = current.getDefiningOp();
- if (!def)
- return false;
- if (isa<memref::AllocaOp>(def))
+ if (ViewLikeOpInterface viewLikeOp = dyn_cast<ViewLikeOpInterface>(userOp)) {
+ if (viewLikeOp.getViewSource() == trackedOperand) {
+ passThroughResults.push_back(viewLikeOp.getViewDest());
return true;
- current = getStorageSource(def);
+ }
+ return false;
}
- return false;
-}
-/// Returns the unique pointer-like value stored to \p address before \p before.
-static Value getUniqueStoredPointerValue(Value address, Operation *before) {
- Operation *scope = before->getParentOfType<acc::ComputeRegionOp>();
- DominanceInfo dominance(scope ? scope : before->getParentOp());
- auto safelyPrecedes = [&](Operation *store) {
- if (dominance.dominates(store, before))
+ // 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.
+ if (acc::PartialEntityAccessOpInterface partialAccess =
+ dyn_cast<acc::PartialEntityAccessOpInterface>(userOp)) {
+ if (partialAccess.getBaseEntity() == trackedOperand) {
+ passThroughResults.append(userOp->result_begin(), userOp->result_end());
return true;
- auto predicate = store->getParentOfType<acc::PredicateRegionOp>();
- auto privateLocal = address.getDefiningOp<acc::PrivateLocalOp>();
- GPUParallelDimsAttr parDims;
- if (privateLocal)
- parDims = mlir::acc::getParDimsAttr(privateLocal);
- bool isThreadPrivate =
- parDims && llvm::any_of(parDims.getArray(),
- [](auto dim) { return dim.isAnyThread(); });
- return predicate && store->getParentOp() == predicate.getOperation() &&
- dominance.dominates(predicate, before) &&
- (isThreadPrivate || hasAllocaRoot(address));
- };
- Operation *storedBy = nullptr;
- Value stored;
- SmallVector<Value> worklist{address};
- DenseSet<Value> seen;
- while (!worklist.empty()) {
- Value alias = worklist.pop_back_val();
- if (!seen.insert(alias).second)
- continue;
- for (Operation *user : alias.getUsers()) {
- if (hasEffectOnValue<MemoryEffects::Write>(user, alias)) {
- if (storedBy || !safelyPrecedes(user))
- return {};
- Value candidate;
- for (Value operand : user->getOperands()) {
- if (operand == alias)
- continue;
- if (candidate)
- return {};
- candidate = operand;
- }
- if (!candidate)
- return {};
- storedBy = user;
- stored = candidate;
- continue;
- }
- SmallVector<Value> through;
- if (getPassThroughResults(user, alias, through))
- worklist.append(through.begin(), through.end());
- else if (!hasEffectOnValue<MemoryEffects::Read>(user, alias))
- return {};
}
- }
- return stored;
-}
-
-/// Tracks a uniquely stored pointer-like value through its subsequent load.
-static bool
-getMemoryPassThroughResults(Operation *userOp, Value trackedOperand,
- SmallVectorImpl<Value> &passThroughResults) {
- if (hasEffectOnValue<MemoryEffects::Read>(userOp, trackedOperand) &&
- userOp->getNumResults() == 1) {
- passThroughResults.push_back(userOp->getResult(0));
- return true;
- }
-
- if (!llvm::is_contained(userOp->getOperands(), trackedOperand))
- return false;
- auto effects = dyn_cast<MemoryEffectOpInterface>(userOp);
- if (!effects)
return false;
- SmallVector<MemoryEffects::EffectInstance> instances;
- effects.getEffects(instances);
- for (const auto &instance : instances) {
- Value target = instance.getValue();
- if (!isa<MemoryEffects::Write>(instance.getEffect()) || !target ||
- target == trackedOperand)
- continue;
- unsigned writers = llvm::count_if(target.getUsers(), [&](Operation *user) {
- return hasEffectOnValue<MemoryEffects::Write>(user, target);
- });
- if (writers != 1)
- return false;
- passThroughResults.push_back(target);
- return true;
}
return false;
}
@@ -473,14 +341,6 @@ 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;
@@ -841,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 thread_x
+/// feeds if it needs per-thread storage: its par_dims include a thread dim
/// and it has block context so the cross-thread all_reduce is well defined.
static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
SmallVector<Value> worklist{v};
@@ -853,38 +713,31 @@ static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
for (Operation *user : cur.getUsers()) {
if (acc::ReductionAccumulateArrayOp accArr =
dyn_cast<acc::ReductionAccumulateArrayOp>(user)) {
- bool hasThreadX = false;
+ bool hasThread = false;
for (auto pd : accArr.getParDims().getArray())
- hasThreadX |= pd.isThreadX();
- if (hasThreadX && reductionHasBlockContext(accArr))
+ hasThread |= pd.isAnyThread();
+ if (hasThread && reductionHasBlockContext(accArr))
return accArr;
continue;
}
SmallVector<Value> through;
if (getPassThroughResults(user, cur, through))
worklist.append(through.begin(), through.end());
- else if (getMemoryPassThroughResults(user, cur, through))
- worklist.append(through.begin(), through.end());
+ else if (isa<ViewLikeOpInterface>(user))
+ worklist.append(user->result_begin(), user->result_end());
}
}
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).
static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
MemRefType baseTy,
arith::AtomicRMWKind kind) {
- assert(baseTy.getRank() > 0 &&
- "per-thread array reduction accumulator must have positive rank");
+ 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);
@@ -896,11 +749,7 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
return;
}
- Value ub =
- baseTy.isDynamicDim(dim)
- ? memref::DimOp::create(b, loc, alloca, dim).getResult()
- : arith::ConstantIndexOp::create(b, loc, baseTy.getDimSize(dim))
- .getResult();
+ 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());
@@ -911,38 +760,6 @@ static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
buildLoopNest(buildLoopNest, 0);
}
-/// True when \p memref resolves to storage owned by the current thread.
-static bool isPerThreadArrayReductionStorage(Value memref) {
- DenseSet<Value> seen;
- Value current = memref;
- while (current && seen.insert(current).second) {
- Operation *def = current.getDefiningOp();
- if (!def)
- return false;
- if (isa<memref::AllocaOp>(def))
- return true;
- if (isa<acc::UnwrapPrivateOp, acc::GPUSharedMemoryOp>(def))
- return false;
-
- if (Value address = getReadAddress(def)) {
- Value stored = getUniqueStoredPointerValue(address, def);
- if (!stored)
- return false;
- current = stored;
- continue;
- }
-
- if (mlir::acc::GPUParallelDimsAttr parDims =
- mlir::acc::getParDimsAttr(def)) {
- return llvm::any_of(parDims.getArray(),
- [](auto dim) { return dim.isThreadX(); });
- }
-
- current = getStorageSource(def);
- }
- return false;
-}
-
std::optional<int64_t>
ACCCGToGPULowering::isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
MemRefType baseTy) {
@@ -2677,14 +2494,9 @@ void ACCCGToGPULowering::processPrivateLocal(
// the accumulate can reduce each element across threads.
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
- bool isThreadPrivate = isThreadXPrivatize(privatizeOp);
- bool canUseDynamicAlloca =
- canUseDynamicReductionAlloca(baseTy, privatizeOp, arrayAccum);
- if ((isThreadPrivate || arrayAccum) &&
- (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
- canUseDynamicAlloca)) {
- Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
- privatizeOp.getDynamicSizes());
+ if ((isThreadXPrivatize(privatizeOp) || arrayAccum) &&
+ canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
+ Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
FailureOr<arith::AtomicRMWKind> kind = getReductionKind(
arrayAccum.getReductionOperator(), baseTy.getElementType(), loc);
@@ -2771,13 +2583,9 @@ void ACCCGToGPULowering::processPrivateLocal(
acc::ReductionAccumulateArrayOp arrayAccum =
perThreadArrayReductionAccum(privateLocal.getResult());
for (mlir::acc::GPUParallelDimAttr parDim : parDimsPair.first) {
- bool canUseDynamicAlloca =
- canUseDynamicReductionAlloca(baseTy, privatizeOp, arrayAccum);
if ((parDim.isThreadX() || arrayAccum) &&
- (canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack) ||
- canUseDynamicAlloca)) {
- Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy,
- privatizeOp.getDynamicSizes());
+ canUseStackAlloca(baseTy, loc, options.maxThreadPrivateStack)) {
+ Value alloca = memref::AllocaOp::create(rewriter, loc, baseTy);
if (arrayAccum) {
FailureOr<arith::AtomicRMWKind> kind = getReductionKind(
arrayAccum.getReductionOperator(), baseTy.getElementType(), loc);
@@ -3451,10 +3259,12 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
(void)accSupport.emitNYI(loc, "reduction: non-MemRefTy accumulate array");
return;
}
- if (memrefTy.getRank() == 0) {
- (void)accSupport.emitNYI(loc, "reduction: rank-zero accumulate array");
+ 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))
@@ -3494,9 +3304,24 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
}
// Per-element gpu.all_reduce is only correct when each thread owns its own
- // accumulator copy. Classify the operand from its underlying allocation
- // rather than its shape or the accumulate's par_dims.
- bool isPerThreadPrivate = isPerThreadArrayReductionStorage(memref);
+ // 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(); });
+ }
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
@@ -3524,19 +3349,7 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
Value lb =
boundsOp.getLowerbound() ? toIndex(boundsOp.getLowerbound()) : zero;
- 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);
- }
- }
+ Value step = boundsOp.getStride() ? toIndex(boundsOp.getStride()) : one;
// 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.
@@ -3560,15 +3373,13 @@ 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 =
- memrefTy.isDynamicDim(dim)
- ? memref::DimOp::create(rewriter, loc, memref, dim).getResult()
- : arith::ConstantIndexOp::create(rewriter, loc,
- memrefTy.getDimSize(dim))
- .getResult();
+ Value dimSize = arith::ConstantIndexOp::create(
+ rewriter, loc, memrefTy.getDimSize(dim));
indices[dim] =
arith::RemUIOp::create(rewriter, loc, linearIndex, dimSize);
if (dim != 0)
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 fd3bb1e5240f0..ace560d869792 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,122 +84,14 @@ func.func @array_reduction_small_shared() {
return
}
-// Dynamic allocation metadata must not be mistaken for storage provenance.
-// CHECK-LABEL: func.func @array_reduction_heap_alloca_extent
-// CHECK: memref.alloca() : memref<index>
-// CHECK: %[[EXTENT:.*]] = memref.load
-// CHECK: memref.alloc(%[[EXTENT]]) : memref<?xi32>
-// CHECK-NOT: gpu.all_reduce
-// CHECK: gpu.terminator
-func.func @array_reduction_heap_alloca_extent(%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(%extent = %n) : (index) {
- %slot = memref.alloca() : memref<index>
- memref.store %extent, %slot[] : memref<index>
- %loaded_extent = memref.load %slot[] : memref<index>
- %heap = memref.alloc(%loaded_extent) : memref<?xi32>
- %bounds = acc.bounds extent(%loaded_extent : index)
- acc.reduction_accumulate_array %heap bounds(%bounds) <add>
- : memref<?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
- memref.dealloc %heap : memref<?xi32>
- acc.yield
- } {origin = "acc.parallel"}
- return
-}
-
-// Non-storage operands of pure operations must not affect classification.
-// CHECK-LABEL: func.func @array_reduction_shared_select
-// CHECK: arith.select
-// CHECK-NOT: gpu.all_reduce
-// CHECK: gpu.terminator
-func.func @array_reduction_shared_select(%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(%extent = %n) : (index) {
- %true = arith.constant true
- %condition_slot = memref.alloca() : memref<i1>
- memref.store %true, %condition_slot[] : memref<i1>
- %condition = memref.load %condition_slot[] : memref<i1>
- %heap = memref.alloc(%extent) : memref<?xi32>
- %selected = arith.select %condition, %heap, %heap : memref<?xi32>
- %bounds = acc.bounds extent(%extent : index)
- acc.reduction_accumulate_array %selected bounds(%bounds) <add>
- : memref<?xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
- memref.dealloc %heap : memref<?xi32>
- acc.yield
- } {origin = "acc.parallel"}
- return
-}
-
-// A reduction private that exceeds the per-thread stack budget uses indexed
-// backing storage shared by thread_x lanes. It must not be reduced again.
-// CHECK-LABEL: func.func @array_reduction_large_indexed_private
-// CHECK: acc.unwrap_private
-// CHECK: memref.view
-// CHECK: memref.subview
-// CHECK-NOT: gpu.all_reduce
-// CHECK-NOT: acc.reduction_accumulate_array
-func.func @array_reduction_large_indexed_private() {
- %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<8192xi32>>
- acc.compute_region launch(%kbx = %bx, %ktx = %tx)
- ins(%arg0 = %private) : (!acc.private_type<memref<8192xi32>>) {
- %c8192 = arith.constant 8192 : index
- %local = acc.private_local %arg0
- {acc.par_dims = #acc<par_dims[block_x, thread_x]>}
- : (!acc.private_type<memref<8192xi32>>) -> memref<8192xi32>
- %bounds = acc.bounds extent(%c8192 : index)
- acc.reduction_accumulate_array %local bounds(%bounds) <add>
- : memref<8192xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
- acc.yield
- } {origin = "acc.parallel"}
- return
-}
-
-// CHECK-LABEL: func.func @array_reduction_dynamic_shared
-// 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
- %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>>
- %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 %cast 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
-// CHECK: %[[BYTE_STEP:.*]] = arith.constant 8 : index
+// CHECK: %[[STEP:.*]] = arith.constant 2 : index
// CHECK: %[[EXTENT:.*]] = arith.constant 3 : 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: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index
// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index
-// CHECK: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[ELEMENT_STEP]]
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[UB]] step %[[STEP]]
func.func @array_reduction_strided_extent() {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
@@ -207,11 +99,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
- %c8 = arith.constant 8 : index
+ %c2 = arith.constant 2 : index
%c3 = arith.constant 3 : index
%local = memref.alloca() : memref<8xi32>
%bounds = acc.bounds lowerbound(%c1_b : index) extent(%c3 : index)
- stride(%c8 : index) {strideInBytes = true}
+ stride(%c2 : index)
acc.reduction_accumulate_array %local bounds(%bounds) <add>
: memref<8xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
acc.yield
@@ -219,14 +111,13 @@ func.func @array_reduction_strided_extent() {
return
}
-// 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.
+// 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.
//
// CHECK-LABEL: func.func @array_reduction_dynamic_par_dims
-// CHECK: memref.reinterpret_cast
-// CHECK-NOT: gpu.all_reduce
-// CHECK: gpu.terminator
+// CHECK: scf.for
+// CHECK: gpu.all_reduce add
func.func @array_reduction_dynamic_par_dims(%buf: memref<?xi32>, %n: index) {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
@@ -254,6 +145,8 @@ func.func @array_reduction_dynamic_par_dims(%buf: memref<?xi32>, %n: 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>
+// CHECK: %[[R2_RESULT:.*]] = gpu.all_reduce add
+// CHECK: memref.store %[[R2_RESULT]], %[[R2_ALLOCA]][%[[R2_I_IDX]], %[[R2_J_IDX]]]
func.func @rank_two_array_reduction() {
%c1 = arith.constant 1 : index
@@ -271,34 +164,6 @@ 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: acc.gpu_shared_memory
-// 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:.*]] =
@@ -312,6 +177,8 @@ func.func @rank_two_array_reduction_thread_y() {
// 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>
+// CHECK: %[[R3_RESULT:.*]] = gpu.all_reduce add
+// CHECK: memref.store %[[R3_RESULT]], %[[R3_ALLOCA]][%[[R3_I_IDX]], %[[R3_J_IDX]], %[[R3_K_IDX]]]
func.func @rank_three_array_reduction() {
%c1 = arith.constant 1 : index
@@ -329,81 +196,22 @@ func.func @rank_three_array_reduction() {
return
}
+// Unsupported dynamic multi-rank accumulators must not generate invalid
+// indexing operations.
// CHECK-LABEL: func.func @dynamic_rank_two_array_reduction
// CHECK: gpu.launch
-// 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: %[[DYN_DESC:.*]] = builtin.unrealized_conversion_cast %[[DYN_ALLOCA]]
-// CHECK: scf.for %[[DYN_LINEAR:.*]] =
-// CHECK: memref.dim %[[DYN_DESC]]
-// CHECK: arith.remui %[[DYN_LINEAR]]
-// CHECK: arith.divui %[[DYN_LINEAR]]
-// 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
- %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]>]
- : (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>
- %descriptor = builtin.unrealized_conversion_cast %local
- : memref<?x?xi32> to memref<?x?xi32>
- %bounds = acc.bounds extent(%ext : index)
- acc.reduction_accumulate_array %descriptor bounds(%bounds) <add>
- : memref<?x?xi32>
- {par_dims = #acc<par_dims[block_x, thread_x]>}
- acc.yield
- } {origin = "acc.parallel"}
- }
- return
-}
-
-// 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_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-NOT: memref.load
// CHECK: gpu.terminator
-func.func @dynamic_rank_two_array_reduction_thread_yz(%m: index, %n: index) {
+func.func @dynamic_rank_two_array_reduction(%local: memref<?x?xi32>, %extent: 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>}
- %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, %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, 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, thread_z]>}
- acc.yield
- } {origin = "acc.parallel"}
- }
+ %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
}
@@ -411,14 +219,14 @@ func.func @dynamic_rank_two_array_reduction_thread_yz(%m: index, %n: index) {
// CHECK: %[[LB:.*]] = arith.constant 5 : index
// CHECK: %[[STEP:.*]] = arith.constant 2 : index
// CHECK: %[[EXTENT:.*]] = arith.constant 3 : index
-// CHECK: %[[LAYOUT_ALLOCA:.*]] = memref.alloca() : memref<3x4xi32, strided<[8, 2]>>
+// CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<3x4xi32, strided<[8, 2]>>
// CHECK: %[[SPAN:.*]] = arith.muli %[[EXTENT]], %[[STEP]] : index
// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SPAN]] : index
// CHECK: scf.for %[[LINEAR:.*]] = %[[LB]] to %[[UB]] step %[[STEP]]
// CHECK: %[[COL:.*]] = arith.remui %[[LINEAR]], %{{.*}} : index
// CHECK: %[[ROW_LINEAR:.*]] = arith.divui %[[LINEAR]], %{{.*}} : index
// CHECK: %[[ROW:.*]] = arith.remui %[[ROW_LINEAR]], %{{.*}} : index
-// CHECK: memref.load %[[LAYOUT_ALLOCA]][%[[ROW]], %[[COL]]] : memref<3x4xi32, strided<[8, 2]>>
+// CHECK: memref.load %[[ALLOCA]][%[[ROW]], %[[COL]]] : memref<3x4xi32, strided<[8, 2]>>
func.func @rank_two_partial_bounds_strided_layout() {
%c1 = arith.constant 1 : index
%c128 = arith.constant 128 : index
More information about the Mlir-commits
mailing list