[Mlir-commits] [mlir] [mlir][acc] Decide reduction atomics by destination instead of by partitioning (PR #235099)
Vijay Kandiah
llvmlistbot at llvm.org
Mon Oct 26 08:46:27 PDT 2026
https://github.com/VijayKandiah updated https://github.com/llvm/llvm-project/pull/235099
>From ebead7bc927485a14ed254825a58cc07a7b793d5 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Fri, 23 Oct 2026 13:59:39 -0700
Subject: [PATCH] [mlir][acc] Decide reduction atomics by destination instead
of by partitioning
---
.../mlir/Dialect/OpenACC/OpenACCCGOps.td | 4 +
mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp | 9 +++
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 65 ++++++++-------
.../acc-cg-to-gpu-reduction-array.mlir | 81 +++++++++++++++++++
...gpu-reduction-thread-only-multi-block.mlir | 46 +++++++++++
5 files changed, 174 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-thread-only-multi-block.mlir
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
index 077c1b7a0db172..0136721b26f7a7 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
@@ -656,6 +656,10 @@ def OpenACC_ComputeRegionOp
/// Check whether all parallel dimensions have width 1.
bool isEffectivelySerial();
+ /// Check whether the launch covers more than one thread block. A width that
+ /// is not known at compile time is conservatively counted as more than one.
+ bool launchesMultipleBlocks();
+
/// Get the block argument representing the width for a given dimension.
::mlir::BlockArgument parDimToWidth(
::mlir::acc::GPUParallelDimAttr parDim);
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
index dd2aae7e121ebe..e24e1330b0b3a5 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
@@ -713,6 +713,15 @@ bool ComputeRegionOp::isEffectivelySerial() {
checkDim(GPUParallelDimAttr::blockZDim(ctx));
}
+bool ComputeRegionOp::launchesMultipleBlocks() {
+ return llvm::any_of(getLaunchParDims(), [&](GPUParallelDimAttr dim) {
+ if (!dim.isAnyBlock())
+ return false;
+ std::optional<uint64_t> width = getKnownConstantLaunchArg(dim);
+ return !width || *width != 1;
+ });
+}
+
BlockArgument ComputeRegionOp::parDimToWidth(GPUParallelDimAttr parDim) {
for (auto [pos, launchArg] : llvm::enumerate(getLaunchArgs())) {
auto parOp = launchArg.getDefiningOp<ParWidthOp>();
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index db017f6a1852a9..06d1cb3a1e0625 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -3994,32 +3994,28 @@ void ACCCGToGPULowering::processAccumulateArrayOp(
rewriter.eraseOp(boundsOp);
};
+ auto isUnitExtent = [&](mlir::acc::GPUParallelDimAttr parDim) {
+ std::optional<uint64_t> width =
+ computeRegion.getKnownConstantLaunchArg(parDim);
+ return width && *width == 1;
+ };
+
+ // A dimension launched one wide partitions nothing, so skip it.
bool hasThreadDim = false;
- bool hasBlockDim = false;
for (auto pd : op.getParDims().getArray()) {
+ if (isUnitExtent(pd))
+ continue;
hasThreadDim |= pd.isAnyThread();
- hasBlockDim |= pd.isAnyBlock();
}
- // Block-only (gang) reduction: each element is produced by one gang, so the
- // per-gang copy already holds the result and the combine does the rest.
- if (hasBlockDim && !hasThreadDim) {
+ // Without a thread dimension no two threads share an accumulator copy, so
+ // each copy already holds a complete total and nothing has to be merged
+ // here. Merging across gangs, if any, is left to the combine.
+ if (!hasThreadDim) {
eraseDeadBounds();
return;
}
- // A thread-only accumulate merges with a within-block all_reduce, which is
- // complete only in one block: block context, or a launch with no block dim.
- // Multi-block thread-only still grid-strides across blocks, so stays NYI.
- bool regionLaunchesBlocks = llvm::any_of(
- computeRegion.getLaunchParDims(),
- [](mlir::acc::GPUParallelDimAttr d) { return d.isAnyBlock(); });
- if (!reductionHasBlockContext(op) && regionLaunchesBlocks) {
- (void)accSupport.emitNYI(
- loc, "reduction: thread-only array reduction accumulate");
- return;
- }
-
// 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/heap allocation is block-shared regardless of
@@ -4179,20 +4175,27 @@ void ACCCGToGPULowering::processReductionCombineOp(acc::ReductionCombineOp op) {
SmallVector<mlir::acc::GPUParallelDimAttr> parDims =
getReductionCombineParDims(op);
- for (auto parDim : parDims) {
- if (parDim.isAnyBlock() && !destIsPerThreadPrivate) {
- // Block reduction directly stores to the accumulator using atomic.
- // The predication (tid.x == 0 when subgroup-aligned) is already handled
- // by the parent predicate_region processing.
- // Reloading a grid-shared slot races with other blocks; record
- // it and replace with the block-reduced register value in the fixup.
- auto srcLoad =
- memref::LoadOp::create(rewriter, loc, srcMemref, ValueRange{});
- pendingCombineReloads.push_back({srcMemref, srcLoad});
- constructAtomicAccumulation(loc, destMemref, /*indices=*/{}, srcLoad,
- kind);
- return;
- }
+ // An atomic is needed exactly when more than one block can reach the
+ // destination. A block dimension in the reduction's parallel dims implies
+ // that; so does a thread-only reduction in a multi-block launch, where every
+ // block combines its own partial into the same shared slot.
+ bool destIsGridShared =
+ !destIsPerThreadPrivate &&
+ (llvm::any_of(parDims,
+ [](mlir::acc::GPUParallelDimAttr parDim) {
+ return parDim.isAnyBlock();
+ }) ||
+ computeRegion.launchesMultipleBlocks());
+ if (destIsGridShared) {
+ // The predication (tid.x == 0 when subgroup-aligned) is already handled
+ // by the parent predicate_region processing.
+ // Reloading a grid-shared slot races with other blocks; record it and
+ // replace with the block-reduced register value in the fixup.
+ auto srcLoad =
+ memref::LoadOp::create(rewriter, loc, srcMemref, ValueRange{});
+ pendingCombineReloads.push_back({srcMemref, srcLoad});
+ constructAtomicAccumulation(loc, destMemref, /*indices=*/{}, srcLoad, kind);
+ return;
}
// Atomic construction is not needed; lower this operation to typical
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 0e0a9366da26f3..ed07b09138b637 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
@@ -453,3 +453,84 @@ func.func @thread_only_array_reduction_single_block() {
} <{origin = "acc.parallel"}>
return
}
+
+// A block dim launched one wide is one block, so a thread-only accumulate is
+// just as complete as it is with no block dim at all.
+//
+// CHECK-LABEL: func.func @thread_only_array_reduction_unit_block
+// CHECK: gpu.launch
+// CHECK-NOT: acc.reduction_accumulate_array
+// CHECK: scf.for %[[UB_IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK: %[[UB_ELT:.*]] = memref.load %[[UB_ALLOCA:.*]][%[[UB_IV]]] : memref<8xi32>
+// CHECK: %[[UB_RED:.*]] = gpu.all_reduce add %[[UB_ELT]]
+// CHECK: memref.store %[[UB_RED]], %[[UB_ALLOCA]][%[[UB_IV]]] : memref<8xi32>
+// CHECK: }
+func.func @thread_only_array_reduction_unit_block() {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_pw 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) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c0_i32 = arith.constant 0 : i32
+ %local = memref.alloca() : memref<8xi32>
+ scf.for %i = %c0 to %c8 step %c1 {
+ memref.store %c0_i32, %local[%i] : memref<8xi32>
+ }
+ %bounds = acc.bounds extent(%c8 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ par_dims(#acc<par_dims[thread_x]>) : memref<8xi32>
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
+
+// A dimension launched one wide partitions nothing, so an accumulate over
+// block_x and a one-wide thread_y reduces to the block-only case: the per-gang
+// copy already holds the result and no cross-thread merge is emitted.
+//
+// CHECK-LABEL: func.func @unit_thread_dim_array_reduction_is_block_only
+// CHECK-NOT: acc.reduction_accumulate_array
+// CHECK-NOT: gpu.all_reduce
+// CHECK-NOT: memref.atomic_rmw
+func.func @unit_thread_dim_array_reduction_is_block_only() {
+ %c1_pw = arith.constant 1 : index
+ %c8_pw = arith.constant 8 : index
+ %bx = acc.par_width %c8_pw par_dim(#acc.par_dim<block_x>)
+ %ty = acc.par_width %c1_pw par_dim(#acc.par_dim<thread_y>)
+ acc.compute_region launch(%kbx = %bx, %kty = %ty) {
+ %c2 = arith.constant 2 : index
+ %local = memref.alloca() : memref<2xi32>
+ %bounds = acc.bounds extent(%c2 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ par_dims(#acc<par_dims[block_x, thread_y]>) : memref<2xi32>
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
+
+// When every dimension the accumulate names is launched one wide, nothing
+// partitions the reduction at all and the accumulator already holds the result.
+//
+// CHECK-LABEL: func.func @all_unit_dims_array_reduction_is_nop
+// CHECK-NOT: acc.reduction_accumulate_array
+// CHECK-NOT: gpu.all_reduce
+// CHECK-NOT: memref.atomic_rmw
+func.func @all_unit_dims_array_reduction_is_nop() {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_pw par_dim(#acc.par_dim<block_x>)
+ %ty = acc.par_width %c1_pw par_dim(#acc.par_dim<thread_y>)
+ %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>)
+ acc.compute_region launch(%kbx = %bx, %kty = %ty, %ktx = %tx) {
+ %c2 = arith.constant 2 : index
+ %local = memref.alloca() : memref<2xi32>
+ %bounds = acc.bounds extent(%c2 : index)
+ acc.reduction_accumulate_array %local bounds(%bounds) <add>
+ par_dims(#acc<par_dims[block_x, thread_y]>) : memref<2xi32>
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-thread-only-multi-block.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-thread-only-multi-block.mlir
new file mode 100644
index 00000000000000..8f451bc5229104
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-thread-only-multi-block.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" --split-input-file | FileCheck %s
+
+// A reduction partitioned only across threads still leaves one partial per
+// block when the launch spans several blocks, and every block combines its
+// partial into the same shared slot. The combine must therefore be atomic even
+// though no block dimension appears in the reduction's parallel dims.
+
+// CHECK-LABEL: func.func @thread_only_combine_multi_block
+// CHECK: acc.atomic.update
+
+func.func @thread_only_combine_multi_block(%arg_slot: memref<i32>, %arg_res: memref<i32>) {
+ %c8_pw = arith.constant 8 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c8_pw 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(%a_slot = %arg_slot, %a_res = %arg_res) : (memref<i32>, memref<i32>) {
+ acc.predicate_region {
+ acc.reduction_combine %a_slot into %a_res <add> par_dims(#acc<par_dims[thread_x]>) : memref<i32>
+ }
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
+
+// -----
+
+// A single-block launch has only one partial, so the same thread-only combine
+// needs no atomic and stays a plain load/combine/store.
+
+// CHECK-LABEL: func.func @thread_only_combine_single_block
+// CHECK-NOT: acc.atomic.update
+// CHECK: memref.store
+
+func.func @thread_only_combine_single_block(%arg_slot: memref<i32>, %arg_res: memref<i32>) {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_pw 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(%a_slot = %arg_slot, %a_res = %arg_res) : (memref<i32>, memref<i32>) {
+ acc.predicate_region {
+ acc.reduction_combine %a_slot into %a_res <add> par_dims(#acc<par_dims[thread_x]>) : memref<i32>
+ }
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
More information about the Mlir-commits
mailing list