[Mlir-commits] [mlir] 930cd75 - [mlir][OpenACC] Support static multi-rank OpenACC array reduction accumulators (#210853)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 22 08:05:40 PDT 2026
Author: Matsu
Date: 2026-07-22T08:05:35-07:00
New Revision: 930cd75aa341776565af08cc098c9270e9a16623
URL: https://github.com/llvm/llvm-project/commit/930cd75aa341776565af08cc098c9270e9a16623
DIFF: https://github.com/llvm/llvm-project/commit/930cd75aa341776565af08cc098c9270e9a16623.diff
LOG: [mlir][OpenACC] Support static multi-rank OpenACC array reduction accumulators (#210853)
Example:
```fortran
!$acc parallel loop reduction(+:a)
do i = 1, n
a(i,:) = a(i,:) + input(i,:)
end do
```
In this code, the reduction accumulator can have rank greater than one,
but GPU lowering assumes rank one.
Fix: initialize rank-N accumulators with nested loops and delinearize
flattened indices before loading, reducing, and storing each element.
Added:
Modified:
mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 77fa82ba6bb4e..f0fd7ae96ad4f 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,15 @@ 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() > 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 +3371,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..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
@@ -134,3 +134,115 @@ 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>
+// 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
+ %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>
+// 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
+ %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 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_two_partial_bounds_strided_layout
+// CHECK: %[[LB:.*]] = arith.constant 5 : index
+// CHECK: %[[STEP:.*]] = arith.constant 2 : index
+// CHECK: %[[EXTENT:.*]] = arith.constant 3 : index
+// 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 %[[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
+}
More information about the Mlir-commits
mailing list