[Mlir-commits] [mlir] b9d15ca - [mlir][acc] Decide reduction atomics by destination instead of by partitioning (#235099)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 09:01:55 PDT 2026
Author: Vijay Kandiah
Date: 2026-10-26T16:01:50Z
New Revision: b9d15ca71c729d89b54f4baa233dd850b10b1905
URL: https://github.com/llvm/llvm-project/commit/b9d15ca71c729d89b54f4baa233dd850b10b1905
DIFF: https://github.com/llvm/llvm-project/commit/b9d15ca71c729d89b54f4baa233dd850b10b1905.diff
LOG: [mlir][acc] Decide reduction atomics by destination instead of by partitioning (#235099)
A reduction combine became atomic only if the reduction's parallel dims
named a
block dimension, so the choice rested entirely on how the reduction was
partitioned. Partitioning answers whether a merge is needed at all.
Whether that
merge has to be atomic should depend instead on how many blocks can
reach the
destination, which is a property of the launch. A thread-only reduction
shows
the difference: it names no block dimension, yet every block still
combines its
own partial into the same shared slot, so the combine stayed non-atomic,
leaving
those updates open to a data race.
An array accumulate partitioned only across threads avoided emitting
that
unsynchronized merge by reporting not implemented whenever the launch
had a
block dimension. The test was on the dimension's presence and not its
width, so
it also rejected launches whose block dimension is one wide and
therefore
partitioned nothing.
Repro:
```mlir
// repro.mlir
func.func @thread_only_array_reduction() {
%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) {
%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
}
```
```
mlir-opt repro.mlir --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))"
```
This reported not implemented prior to this change:
```
error: not yet implemented: reduction: thread-only array reduction accumulate
```
### Fix
Decide the atomic from how many blocks can reach the destination. A
block
dimension in the reduction implies more than one, and so does a
thread-only
reduction in a multi-block launch; a thread-private destination is
reached by a
single thread either way. A launch width that is not a compile-time
constant is
counted conservatively as more than one block. The accumulate in turn
skips
dimensions launched one wide, since those partition nothing. With the
merge
across blocks now handled correctly by the combine, the accumulate has
no case
left that it cannot lower, so the diagnostic is removed.
The same command now lowers to a per-element reduction over the array:
```mlir
scf.for %arg12 = %c0_2 to %1 step %c1_3 {
%2 = memref.load %alloca[%arg12] : memref<8xi32>
%3 = gpu.all_reduce add %2 uniform {
} {acc.par_dims = #acc<par_dims[thread_x]>} : (i32) -> i32
memref.store %3, %alloca[%arg12] : memref<8xi32>
}
```
The combine change is not array-specific. It closes the same exposure
for scalar
thread-only reductions in multi-block launches, which were already
reachable and
lowered to a plain store into a shared slot.
One conservative choice is left in place in the combine: a reduction
whose own parallel
dims name a block dimension takes the atomic path even when that
dimension is
one wide. Narrowing that to the launch width should be a separate
follow-up.
Added:
mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-thread-only-multi-block.mlir
Modified:
mlir/include/mlir/Dialect/OpenACC/OpenACCCGOps.td
mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-array.mlir
Removed:
################################################################################
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