[Mlir-commits] [mlir] [mlir][acc] Keep worker-only launches unaligned in ACCCGToGPU (PR #217382)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 19 20:50:56 PDT 2026
https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/217382
>From 9aa73cf45010c73e574000fa31865aa4b0b0f5c0 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 19 Aug 2026 04:05:02 -0700
Subject: [PATCH 1/3] [mlir][acc] Keep worker-only launches unaligned in
ACCCGToGPU
Example:
```fortran
!$acc parallel loop gang worker reduction(+:s) num_gangs(1) num_workers(16)
```
This loop declares no vector level, so the launch is (1, 16, 1). Subgroup
alignment padded blockDim.x to 32 and recomputed blockDim.y as
max(1, 16/32) = 1, leaving one worker to run every iteration while 32
lanes repeated the same work.
Fix: skip the alignment when the launch is (1, N, 1) with 1 < N <=
subgroupSize, where the workers already are the subgroup lanes. Such a
launch also needs no per-row barrier, since a ThreadX row is one thread,
and no ThreadX predication, since no lane is padded in. N is capped at
subgroupSize because the worker-indexed shared reduction buffer holds
subgroupSize entries, and N == 1 keeps padding because it has no worker
parallelism to preserve.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 81 +++++++++++++------
.../acc-cg-to-gpu-reduction-array.mlir | 37 ++++++++-
2 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 0dda217250e90..bef4a1b6a5612 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -656,6 +656,20 @@ class ACCCGToGPULowering {
llvm::DenseMap<Type, Value> privatizeBroadcastCache;
int64_t staticBlockDimX = 1024;
+ int64_t staticBlockDimY = 1024;
+ int64_t staticBlockDimZ = 1024;
+
+ /// True when the launch is known to be (1, N, 1) with 1 < N <= subgroupSize:
+ /// every worker owns exactly one thread. Such a launch is left unaligned, so
+ /// no ThreadX lane is ever padded in and a ThreadX row holds one thread. N is
+ /// capped by the subgroup size because worker-indexed shared buffers hold
+ /// subgroupSize entries, and N == 1 is excluded because a launch with no
+ /// parallelism to preserve is better off padded into one full subgroup.
+ bool isSingleThreadWorkerLaunch() const {
+ return staticBlockDimX == 1 && staticBlockDimZ == 1 &&
+ staticBlockDimY > 1 && staticBlockDimY <= options.subgroupSize;
+ }
+
acc::DefaultACCToGPUMappingPolicy defaultPolicy;
SharedMemoryBudget sharedMemBudget;
SmallVector<std::string> sharedMemPrivateVarNames;
@@ -906,7 +920,13 @@ LogicalResult ACCCGToGPULowering::rewrite() {
if (matchPattern(blockDimX, m_ConstantInt(&bdxVal)))
staticBlockDimX = bdxVal.getSExtValue();
Value blockDimY = launchArgument(gpu::Processor::ThreadY);
+ APInt bdyVal;
+ if (matchPattern(blockDimY, m_ConstantInt(&bdyVal)))
+ staticBlockDimY = bdyVal.getSExtValue();
Value blockDimZ = launchArgument(gpu::Processor::ThreadZ);
+ APInt bdzVal;
+ if (matchPattern(blockDimZ, m_ConstantInt(&bdzVal)))
+ staticBlockDimZ = bdzVal.getSExtValue();
Value gridDimX = launchArgument(gpu::Processor::BlockX);
Value gridDimY = launchArgument(gpu::Processor::BlockY);
Value gridDimZ = launchArgument(gpu::Processor::BlockZ);
@@ -1091,7 +1111,30 @@ LogicalResult ACCCGToGPULowering::rewrite() {
});
}
- if (isShuffleEnabled || hasThreadYBarrier) {
+ std::optional<int64_t> constBlockDimX =
+ getConstantIntValue(launch.getBlockSizeX());
+ std::optional<int64_t> constBlockDimY =
+ getConstantIntValue(launch.getBlockSizeY());
+ std::optional<int64_t> constBlockDimZ =
+ getConstantIntValue(launch.getBlockSizeZ());
+
+ // Skip subgroup alignment only when the total thread count is already
+ // below a subgroup (constant blockDim.x in 2..subgroupSize-1 and
+ // constant blockDim.y/z == 1). If blockDim.y/z > 1 or is unknown,
+ // padding blockDim.x to a subgroup is still required so subgroups don't
+ // cross row boundaries for row-local shuffle/ThreadY-barrier reductions.
+ bool skipAlign = constBlockDimX && constBlockDimY && constBlockDimZ &&
+ *constBlockDimX > 1 && *constBlockDimX < subgroupSize &&
+ *constBlockDimY == 1 && *constBlockDimZ == 1;
+
+ // A (1, N, 1) launch needs no padding either: the workers are already the
+ // subgroup lanes, so shuffle reductions run on linearized thread IDs and a
+ // per-row barrier covers a single thread. Padding would instead fold the
+ // workers into ThreadX lanes and serialize the worker parallelism.
+ if (isSingleThreadWorkerLaunch())
+ skipAlign = true;
+
+ if ((isShuffleEnabled || hasThreadYBarrier) && !skipAlign) {
rewriter.setInsertionPoint(launch);
Value curBlockDimX = launch.getBlockSizeX();
@@ -1121,22 +1164,6 @@ LogicalResult ACCCGToGPULowering::rewrite() {
.str();
});
- std::optional<int64_t> constBlockDimX = getConstantIntValue(curBlockDimX);
- std::optional<int64_t> constBlockDimY = getConstantIntValue(curBlockDimY);
- std::optional<int64_t> constBlockDimZ = getConstantIntValue(curBlockDimZ);
-
- // Skip subgroup alignment only when the total thread count is already
- // below a subgroup (constant blockDim.x in 2..subgroupSize-1 and
- // constant blockDim.y/z == 1). If blockDim.y/z > 1 or is unknown,
- // padding blockDim.x to a subgroup is still required so subgroups don't
- // cross row boundaries for row-local shuffle/ThreadY-barrier reductions.
- bool skipAlign = false;
- if (constBlockDimX && constBlockDimY && constBlockDimZ &&
- *constBlockDimX > 1 && *constBlockDimX < subgroupSize &&
- *constBlockDimY == 1 && *constBlockDimZ == 1) {
- skipAlign = true;
- }
-
// Update the ThreadX length and the numbers of ThreadY and ThreadZ.
// When the original block dimensions are compile-time
// constants, compute the adjusted dimensions as constants directly so
@@ -1188,11 +1215,9 @@ LogicalResult ACCCGToGPULowering::rewrite() {
newBlockDimZ = arith::MaxUIOp::create(rewriter, loc, cst1, quotient);
}
- if (!skipAlign) {
- launch.getBlockSizeXMutable().assign(newBlockDimX);
- launch.getBlockSizeYMutable().assign(newBlockDimY);
- launch.getBlockSizeZMutable().assign(newBlockDimZ);
- }
+ launch.getBlockSizeXMutable().assign(newBlockDimX);
+ launch.getBlockSizeYMutable().assign(newBlockDimY);
+ launch.getBlockSizeZMutable().assign(newBlockDimZ);
}
}
@@ -1649,6 +1674,12 @@ void ACCCGToGPULowering::createBarrier(
}
void ACCCGToGPULowering::createPerRowBarrier(Location loc) {
+ // A row of a (1, N, 1) launch is a single thread, so there is nothing to
+ // synchronize. Returning before setting hasThreadYBarrier keeps the launch
+ // unaligned; a subgroup barrier here would instead synchronize N workers.
+ if (isSingleThreadWorkerLaunch())
+ return;
+
hasThreadYBarrier = true;
if (staticBlockDimX <= options.subgroupSize) {
@@ -2161,7 +2192,8 @@ void ACCCGToGPULowering::processPredicateRegion(
// Exception: if this region contains a thread-level (vector or worker)
// routine call, all ThreadX threads must reach the call so the routine's
// workgroup-wide barriers (e.g. shared memory alloca sync) converge.
- if (hasThreadYReduction) {
+ // A single-thread worker launch is left unaligned, so it has no such lanes.
+ if (hasThreadYReduction && !isSingleThreadWorkerLaunch()) {
MLIRContext *ctx = computeRegion->getContext();
mlir::acc::GPUParallelDimAttr threadXParDim =
mlir::acc::GPUParallelDimAttr::threadXDim(ctx);
@@ -3268,7 +3300,8 @@ void ACCCGToGPULowering::createGPUAllReduceOp(
}
// Subgroup alignment may introduce extra ThreadX lanes even when ThreadX is
// not part of the reduction. Predicate on ThreadX so only one lane stores.
- if (!hasThreadX)
+ // A single-thread worker launch is left unaligned, so it has no such lanes.
+ if (!hasThreadX && !isSingleThreadWorkerLaunch())
inactiveParDims.push_back(mlir::acc::GPUParallelDimAttr::threadXDim(ctx));
Value predicate = emitPredicate(loc, inactiveParDims);
// Predication is only needed when the store target is visible to
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 64eaec673dd1c..12c45e9727d9c 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
@@ -278,11 +278,16 @@ func.func @partial_thread_x_reduction() {
return
}
-// CHECK-LABEL: func.func @thread_y_reduction_still_aligned
-// CHECK: %[[C32:.*]] = arith.constant 32 : index
+// A worker-only launch keeps its (1, N, 1) shape: every worker owns one thread,
+// so the workers are already the subgroup lanes. Padding ThreadX would fold the
+// workers into lanes and serialize them.
+//
+// CHECK-LABEL: func.func @thread_y_reduction_single_thread_rows
+// CHECK: %[[C16_ROWS:.*]] = arith.constant 16 : index
+// CHECK-NOT: arith.constant 32 : index
// CHECK: gpu.launch
-// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32]],
-func.func @thread_y_reduction_still_aligned() {
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %[[C16_ROWS]],
+func.func @thread_y_reduction_single_thread_rows() {
%c1 = arith.constant 1 : index
%c16 = arith.constant 16 : index
%bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
@@ -298,6 +303,30 @@ func.func @thread_y_reduction_still_aligned() {
return
}
+// More workers than a subgroup still get aligned: the worker-indexed shared
+// reduction buffer only holds subgroupSize entries.
+//
+// CHECK-LABEL: func.func @thread_y_reduction_more_workers_than_subgroup
+// CHECK: %[[C32_WIDE:.*]] = arith.constant 32 : index
+// CHECK: %[[C2_WIDE:.*]] = arith.constant 2 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_WIDE]], %{{.*}} = %[[C2_WIDE]],
+func.func @thread_y_reduction_more_workers_than_subgroup() {
+ %c1 = arith.constant 1 : index
+ %c64 = arith.constant 64 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %ty = acc.par_width %c64 {par_dim = #acc.par_dim<thread_y>}
+ acc.compute_region launch(%kbx = %bx, %kty = %ty) {
+ %c0_i32 = arith.constant 0 : i32
+ %local = memref.alloca() : memref<i32>
+ acc.reduction_accumulate %c0_i32 to %local <add>
+ : i32 -> memref<i32>
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// A ThreadX-only reduction still needs aligned rows when ThreadY is greater
// than one, because a physical subgroup must not contain multiple logical rows.
//
>From ae58f30f0a2ae40eb6c7476ff785e26adeeef4d9 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 19 Aug 2026 18:25:48 -0700
Subject: [PATCH 2/3] [mlir][acc] Keep any worker-only launch unaligned, not
just single-lane rows
A launch whose thread-level reductions are all worker ones was padded to a
subgroup-wide blockDim.x, and blockDim.y was divided by the same factor to
keep the thread count. num_workers(16) vector_length(8) thus ran as
(32, 4, 1): 4 workers instead of 16, with the folded-away workers turned into
ThreadX lanes redoing each other's work.
Worker reductions combine their partials in the lowest blockDim.y threads of
the block, so a row can sit anywhere inside a subgroup. Classify the thread
dimensions the reductions span, and skip the alignment when only the workers
ask for it and the row width divides the subgroup size. A per-row barrier in
such a launch covers a row narrower than a subgroup, so emit a lane-masked
barrier over exactly that range rather than a subgroup-wide one, which would
tie together the other rows sharing the subgroup.
Thread-level array accumulates are left out: they combine through atomics on
a shared array instead of worker shuffles, and leaving their rows unaligned
drops updates in some geometries.
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 157 +++++++++++++-----
.../acc-cg-to-gpu-reduction-array.mlir | 58 +++++++
2 files changed, 173 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index bef4a1b6a5612..24f5a9f399ea6 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -424,6 +424,36 @@ static void emitGPUBarrierSubgroup(OpBuilder &builder, Location loc) {
gpu::BarrierScope::Subgroup);
}
+/// Emits a barrier over the lanes of this thread's ThreadY row, for a row that
+/// is narrower than a subgroup. \p rowWidth (blockDim.x) must divide
+/// \p subgroupSize, so the rows tile each subgroup and the row holding lane L
+/// covers the lanes [L - L % rowWidth, L - L % rowWidth + rowWidth). A
+/// subgroup-wide barrier cannot be used here: it would also tie together the
+/// other rows sharing the subgroup, which run different workers.
+static void emitGPUBarrierRow(OpBuilder &builder, Location loc,
+ int64_t rowWidth, int64_t subgroupSize) {
+ assert(rowWidth > 0 && subgroupSize % rowWidth == 0 &&
+ "row must tile the subgroup");
+ Type i32Ty = builder.getI32Type();
+ int64_t rowsPerSubgroup = subgroupSize / rowWidth;
+
+ // laneBase = (tid.y % rowsPerSubgroup) * rowWidth. Both operands are powers
+ // of two here, since rowWidth divides the subgroup size.
+ Value threadY = gpu::ThreadIdOp::create(builder, loc, builder.getIndexType(),
+ gpu::Dimension::y);
+ Value threadY32 = arith::IndexCastOp::create(builder, loc, i32Ty, threadY);
+ Value rowMask =
+ arith::ConstantIntOp::create(builder, loc, i32Ty, rowsPerSubgroup - 1);
+ Value rowIndex = arith::AndIOp::create(builder, loc, threadY32, rowMask);
+ Value width = arith::ConstantIntOp::create(builder, loc, i32Ty, rowWidth);
+ Value laneBase = arith::MulIOp::create(builder, loc, rowIndex, width);
+
+ Value laneBits = arith::ConstantIntOp::create(builder, loc, i32Ty,
+ (int64_t{1} << rowWidth) - 1);
+ Value mask = arith::ShLIOp::create(builder, loc, laneBits, laneBase);
+ NVVM::SyncWarpOp::create(builder, loc, mask);
+}
+
/// Lowers a single `acc.compute_region` to GPU dialect IR.
class ACCCGToGPULowering {
public:
@@ -647,6 +677,10 @@ class ACCCGToGPULowering {
llvm::DenseMap<gpu::Processor, Value> dimensionMap;
// True if ThreadY reduction exists, which triggers subgroup alignment
bool hasThreadYReduction = false;
+ // True if a ThreadX reduction exists, whose shuffles stay inside a row
+ bool hasThreadXReduction = false;
+ // True if an array reduction spans any thread dimension
+ bool hasThreadLevelArrayReduction = false;
// True if any ThreadX routine call exists in the kernel
bool hasThreadLevelRoutineCall = false;
// True when a per-row ThreadY barrier is emitted
@@ -670,6 +704,25 @@ class ACCCGToGPULowering {
staticBlockDimY > 1 && staticBlockDimY <= options.subgroupSize;
}
+ /// True when the workers are the only reason this launch would need its rows
+ /// aligned to subgroup boundaries: the thread-level reductions are all worker
+ /// ones, and the launch is (X, N, 1) with 1 < N <= subgroupSize and X
+ /// dividing the subgroup size. Worker reductions combine their per-worker
+ /// partials in the lowest N threads of the block rather than within a row, so
+ /// a row may sit anywhere inside a subgroup; a row that is narrower than a
+ /// subgroup is then synchronized with a lane-masked barrier. The bounds on N
+ /// and ThreadZ come from the per-worker partial buffer, which holds
+ /// subgroupSize entries indexed by ThreadY alone. Thread-level array
+ /// accumulates are left out: they combine through atomics on one shared array
+ /// instead of worker shuffles, and leaving their rows unaligned drops updates
+ /// in some geometries.
+ bool isWorkerOnlyShuffleLaunch() const {
+ return hasThreadYReduction && !hasThreadXReduction &&
+ !hasThreadLevelArrayReduction && staticBlockDimZ == 1 &&
+ staticBlockDimY > 1 && staticBlockDimY <= options.subgroupSize &&
+ options.subgroupSize % staticBlockDimX == 0;
+ }
+
acc::DefaultACCToGPUMappingPolicy defaultPolicy;
SharedMemoryBudget sharedMemBudget;
SmallVector<std::string> sharedMemPrivateVarNames;
@@ -855,15 +908,20 @@ LogicalResult ACCCGToGPULowering::rewrite() {
// Pre-compute if thread-level reductions exist. ThreadY reduction generates
// shuffles which require subgroup alignment (blockDim.x = subgroupSize),
- // meaning ThreadX lanes exist even without explicit ThreadX parallelism.
- computeRegion->walk([&](acc::ReductionAccumulateOp op) -> WalkResult {
+ // meaning ThreadX lanes exist even without explicit ThreadX parallelism. The
+ // ThreadX dimension is tracked too, to tell a launch whose reductions are all
+ // worker ones from one that also shuffles inside a row.
+ computeRegion->walk([&](acc::ReductionAccumulateOp op) {
for (auto parDim : op.getParDimsAttr().getArray()) {
- if (parDim.isThreadY()) {
- hasThreadYReduction = true;
- return WalkResult::interrupt();
- }
+ hasThreadXReduction |= parDim.isThreadX();
+ hasThreadYReduction |= parDim.isThreadY();
+ }
+ });
+ computeRegion->walk([&](acc::ReductionAccumulateArrayOp op) {
+ for (auto parDim : op.getParDimsAttr().getArray()) {
+ if (!parDim.isAnyBlock())
+ hasThreadLevelArrayReduction = true;
}
- return WalkResult::advance();
});
// Pre-compute if any thread-level (vector or worker) routine call exists.
@@ -1069,47 +1127,42 @@ LogicalResult ACCCGToGPULowering::rewrite() {
// because:
// - Subgroup reductions (gpu.all_reduce) require full subgroups
// - Per-row workgroup barriers require blockDim.x aligned to subgroupSize
- bool isShuffleEnabled = false;
+ // Tracked apart because only a ThreadX reduction needs its row to start at
+ // a subgroup boundary: it shuffles within the row, so a subgroup that
+ // crosses into the next row would mix rows. A worker reduction combines the
+ // per-worker partials on linearized thread IDs instead, which is
+ // independent of where the rows sit.
+ bool needsThreadXAlign = false;
+ bool needsThreadYAlign = false;
bool alignThreadXReduction =
getConstantIntValue(launch.getBlockSizeY()) != 1 ||
getConstantIntValue(launch.getBlockSizeZ()) != 1;
- launch.walk([&](gpu::AllReduceOp allReduce) -> WalkResult {
- ArrayRef<mlir::acc::GPUParallelDimAttr> parDims =
- mlir::acc::getParDimsAttr(allReduce).getArray();
- for (auto parDim : parDims) {
- if (parDim.isThreadY() ||
- (alignThreadXReduction && parDim.isThreadX())) {
- // Shuffle are enabled. Need to adjust the ThreadX length.
- isShuffleEnabled = true;
- return WalkResult::interrupt();
- }
+ // A reduction over both dimensions counts as a ThreadX one: it still
+ // shuffles ThreadX lanes within the row.
+ auto classifyAllReduce = [&](gpu::AllReduceOp allReduce) {
+ bool hasThreadX = false;
+ bool hasThreadY = false;
+ for (auto parDim : mlir::acc::getParDimsAttr(allReduce).getArray()) {
+ hasThreadX |= parDim.isThreadX();
+ hasThreadY |= parDim.isThreadY();
}
- return WalkResult::advance();
+ if (hasThreadX && alignThreadXReduction)
+ needsThreadXAlign = true;
+ else if (hasThreadY)
+ needsThreadYAlign = true;
+ };
+
+ launch.walk(classifyAllReduce);
+ // Also check if called routines have thread-level reductions.
+ launch.walk([&](func::CallOp callOp) {
+ if (gpu::GPUFuncOp callee =
+ callOp->getParentOfType<ModuleOp>().lookupSymbol<gpu::GPUFuncOp>(
+ callOp.getCallee()))
+ callee.walk(classifyAllReduce);
});
- // Also check if called routines have ThreadY reductions
- if (!isShuffleEnabled) {
- launch.walk([&](func::CallOp callOp) -> WalkResult {
- if (gpu::GPUFuncOp callee =
- callOp->getParentOfType<ModuleOp>()
- .lookupSymbol<gpu::GPUFuncOp>(callOp.getCallee())) {
- callee.walk([&](gpu::AllReduceOp allReduce) -> WalkResult {
- ArrayRef<mlir::acc::GPUParallelDimAttr> parDims =
- mlir::acc::getParDimsAttr(allReduce).getArray();
- for (auto parDim : parDims) {
- if (parDim.isThreadY() ||
- (alignThreadXReduction && parDim.isThreadX())) {
- isShuffleEnabled = true;
- return WalkResult::interrupt();
- }
- }
- return WalkResult::advance();
- });
- }
- return isShuffleEnabled ? WalkResult::interrupt()
- : WalkResult::advance();
- });
- }
+
+ bool isShuffleEnabled = needsThreadXAlign || needsThreadYAlign;
std::optional<int64_t> constBlockDimX =
getConstantIntValue(launch.getBlockSizeX());
@@ -1134,6 +1187,17 @@ LogicalResult ACCCGToGPULowering::rewrite() {
if (isSingleThreadWorkerLaunch())
skipAlign = true;
+ // When the workers are the only reason to align, keep the launch as it is.
+ // Worker reductions combine their partials in the lowest blockDim.y threads
+ // of the block, which sit in the first subgroup whatever blockDim.x is, so
+ // no row has to start at a subgroup boundary, and the barriers over such a
+ // row are lane-masked. Padding blockDim.x here would instead divide
+ // blockDim.y by the same factor and leave the folded-away workers as
+ // ThreadX lanes redoing each other's work.
+ if (needsThreadYAlign && !needsThreadXAlign && !hasThreadYBarrier &&
+ isWorkerOnlyShuffleLaunch())
+ skipAlign = true;
+
if ((isShuffleEnabled || hasThreadYBarrier) && !skipAlign) {
rewriter.setInsertionPoint(launch);
@@ -1680,6 +1744,15 @@ void ACCCGToGPULowering::createPerRowBarrier(Location loc) {
if (isSingleThreadWorkerLaunch())
return;
+ // A worker-only launch keeps its rows unaligned, so a row narrower than a
+ // subgroup covers only part of one. Synchronize exactly that lane range and
+ // leave hasThreadYBarrier alone, so the launch is not padded on account of
+ // this barrier.
+ if (isWorkerOnlyShuffleLaunch() && staticBlockDimX < options.subgroupSize) {
+ emitGPUBarrierRow(rewriter, loc, staticBlockDimX, options.subgroupSize);
+ return;
+ }
+
hasThreadYBarrier = true;
if (staticBlockDimX <= options.subgroupSize) {
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 12c45e9727d9c..f9c6288b33052 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
@@ -303,6 +303,64 @@ func.func @thread_y_reduction_single_thread_rows() {
return
}
+// A worker-only launch with several threads per worker keeps its shape too.
+// Worker partials are combined in the lowest ThreadY threads of the block, so
+// the rows do not have to start at subgroup boundaries.
+//
+// CHECK-LABEL: func.func @thread_y_reduction_narrow_rows
+// CHECK: %[[C8_NARROW:.*]] = arith.constant 8 : index
+// CHECK: %[[C16_NARROW:.*]] = arith.constant 16 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C8_NARROW]], %{{.*}} = %[[C16_NARROW]],
+func.func @thread_y_reduction_narrow_rows() {
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c8 {par_dim = #acc.par_dim<thread_x>}
+ %ty = acc.par_width %c16 {par_dim = #acc.par_dim<thread_y>}
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx, %kty = %ty) {
+ %c0_i32 = arith.constant 0 : i32
+ %local = memref.alloca() : memref<i32>
+ acc.reduction_accumulate %c0_i32 to %local <add>
+ : i32 -> memref<i32>
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// A ThreadX reduction in the same region shuffles within a row, so the rows are
+// aligned again: blockDim.x is padded to a subgroup and blockDim.y divided by
+// the same factor.
+//
+// CHECK-LABEL: func.func @thread_y_reduction_with_thread_x_reduction
+// CHECK: %[[C32_BOTH:.*]] = arith.constant 32 : index
+// CHECK: %[[C4_BOTH:.*]] = arith.constant 4 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_BOTH]], %{{.*}} = %[[C4_BOTH]],
+func.func @thread_y_reduction_with_thread_x_reduction() {
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c8 {par_dim = #acc.par_dim<thread_x>}
+ %ty = acc.par_width %c16 {par_dim = #acc.par_dim<thread_y>}
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx, %kty = %ty) {
+ %c0_i32 = arith.constant 0 : i32
+ %worker = memref.alloca() : memref<i32>
+ %vector = memref.alloca() : memref<i32>
+ acc.reduction_accumulate %c0_i32 to %worker <add>
+ : i32 -> memref<i32>
+ {par_dims = #acc<par_dims[block_x, thread_y]>}
+ acc.reduction_accumulate %c0_i32 to %vector <add>
+ : i32 -> memref<i32>
+ {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
// More workers than a subgroup still get aligned: the worker-indexed shared
// reduction buffer only holds subgroupSize entries.
//
>From cfe0ecd9a3a25f5699408995dede17833cfed867 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 19 Aug 2026 20:42:31 -0700
Subject: [PATCH 3/3] [mlir][acc] Shorten the worker-only launch comments
---
.../Dialect/OpenACC/Transforms/ACCCGToGPU.cpp | 73 ++++++-------------
.../acc-cg-to-gpu-reduction-array.mlir | 10 +--
2 files changed, 26 insertions(+), 57 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index 24f5a9f399ea6..7f7f76ad63e91 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -424,12 +424,9 @@ static void emitGPUBarrierSubgroup(OpBuilder &builder, Location loc) {
gpu::BarrierScope::Subgroup);
}
-/// Emits a barrier over the lanes of this thread's ThreadY row, for a row that
-/// is narrower than a subgroup. \p rowWidth (blockDim.x) must divide
-/// \p subgroupSize, so the rows tile each subgroup and the row holding lane L
-/// covers the lanes [L - L % rowWidth, L - L % rowWidth + rowWidth). A
-/// subgroup-wide barrier cannot be used here: it would also tie together the
-/// other rows sharing the subgroup, which run different workers.
+/// Emits a barrier over the lanes of this thread's ThreadY row. \p rowWidth
+/// (blockDim.x) must divide \p subgroupSize so the rows tile a subgroup. A
+/// subgroup-wide barrier would also tie together the other rows sharing it.
static void emitGPUBarrierRow(OpBuilder &builder, Location loc,
int64_t rowWidth, int64_t subgroupSize) {
assert(rowWidth > 0 && subgroupSize % rowWidth == 0 &&
@@ -693,29 +690,18 @@ class ACCCGToGPULowering {
int64_t staticBlockDimY = 1024;
int64_t staticBlockDimZ = 1024;
- /// True when the launch is known to be (1, N, 1) with 1 < N <= subgroupSize:
- /// every worker owns exactly one thread. Such a launch is left unaligned, so
- /// no ThreadX lane is ever padded in and a ThreadX row holds one thread. N is
- /// capped by the subgroup size because worker-indexed shared buffers hold
- /// subgroupSize entries, and N == 1 is excluded because a launch with no
- /// parallelism to preserve is better off padded into one full subgroup.
+ /// True when the launch is (1, N, 1) with 1 < N <= subgroupSize, so every
+ /// worker owns one thread. N is capped by the worker-indexed shared buffers,
+ /// which hold subgroupSize entries.
bool isSingleThreadWorkerLaunch() const {
return staticBlockDimX == 1 && staticBlockDimZ == 1 &&
staticBlockDimY > 1 && staticBlockDimY <= options.subgroupSize;
}
- /// True when the workers are the only reason this launch would need its rows
- /// aligned to subgroup boundaries: the thread-level reductions are all worker
- /// ones, and the launch is (X, N, 1) with 1 < N <= subgroupSize and X
- /// dividing the subgroup size. Worker reductions combine their per-worker
- /// partials in the lowest N threads of the block rather than within a row, so
- /// a row may sit anywhere inside a subgroup; a row that is narrower than a
- /// subgroup is then synchronized with a lane-masked barrier. The bounds on N
- /// and ThreadZ come from the per-worker partial buffer, which holds
- /// subgroupSize entries indexed by ThreadY alone. Thread-level array
- /// accumulates are left out: they combine through atomics on one shared array
- /// instead of worker shuffles, and leaving their rows unaligned drops updates
- /// in some geometries.
+ /// True when only worker reductions would need the rows aligned. They combine
+ /// their partials in the lowest N threads of the block, not within a row, so
+ /// a row may sit anywhere in a subgroup. Array accumulates use atomics
+ /// instead and still need alignment.
bool isWorkerOnlyShuffleLaunch() const {
return hasThreadYReduction && !hasThreadXReduction &&
!hasThreadLevelArrayReduction && staticBlockDimZ == 1 &&
@@ -908,9 +894,8 @@ LogicalResult ACCCGToGPULowering::rewrite() {
// Pre-compute if thread-level reductions exist. ThreadY reduction generates
// shuffles which require subgroup alignment (blockDim.x = subgroupSize),
- // meaning ThreadX lanes exist even without explicit ThreadX parallelism. The
- // ThreadX dimension is tracked too, to tell a launch whose reductions are all
- // worker ones from one that also shuffles inside a row.
+ // meaning ThreadX lanes exist even without explicit ThreadX parallelism.
+ // ThreadX is tracked too, to tell worker-only reductions from row shuffles.
computeRegion->walk([&](acc::ReductionAccumulateOp op) {
for (auto parDim : op.getParDimsAttr().getArray()) {
hasThreadXReduction |= parDim.isThreadX();
@@ -1127,11 +1112,8 @@ LogicalResult ACCCGToGPULowering::rewrite() {
// because:
// - Subgroup reductions (gpu.all_reduce) require full subgroups
// - Per-row workgroup barriers require blockDim.x aligned to subgroupSize
- // Tracked apart because only a ThreadX reduction needs its row to start at
- // a subgroup boundary: it shuffles within the row, so a subgroup that
- // crosses into the next row would mix rows. A worker reduction combines the
- // per-worker partials on linearized thread IDs instead, which is
- // independent of where the rows sit.
+ // Tracked apart because only a ThreadX reduction shuffles within the row,
+ // so only it needs the row to start at a subgroup boundary.
bool needsThreadXAlign = false;
bool needsThreadYAlign = false;
bool alignThreadXReduction =
@@ -1180,20 +1162,11 @@ LogicalResult ACCCGToGPULowering::rewrite() {
*constBlockDimX > 1 && *constBlockDimX < subgroupSize &&
*constBlockDimY == 1 && *constBlockDimZ == 1;
- // A (1, N, 1) launch needs no padding either: the workers are already the
- // subgroup lanes, so shuffle reductions run on linearized thread IDs and a
- // per-row barrier covers a single thread. Padding would instead fold the
- // workers into ThreadX lanes and serialize the worker parallelism.
+ // Padding a worker-only launch divides blockDim.y by the same factor and
+ // leaves the folded-away workers as ThreadX lanes redoing each other's
+ // work, so keep its shape instead.
if (isSingleThreadWorkerLaunch())
skipAlign = true;
-
- // When the workers are the only reason to align, keep the launch as it is.
- // Worker reductions combine their partials in the lowest blockDim.y threads
- // of the block, which sit in the first subgroup whatever blockDim.x is, so
- // no row has to start at a subgroup boundary, and the barriers over such a
- // row are lane-masked. Padding blockDim.x here would instead divide
- // blockDim.y by the same factor and leave the folded-away workers as
- // ThreadX lanes redoing each other's work.
if (needsThreadYAlign && !needsThreadXAlign && !hasThreadYBarrier &&
isWorkerOnlyShuffleLaunch())
skipAlign = true;
@@ -1738,16 +1711,14 @@ void ACCCGToGPULowering::createBarrier(
}
void ACCCGToGPULowering::createPerRowBarrier(Location loc) {
- // A row of a (1, N, 1) launch is a single thread, so there is nothing to
- // synchronize. Returning before setting hasThreadYBarrier keeps the launch
- // unaligned; a subgroup barrier here would instead synchronize N workers.
+ // A row of a (1, N, 1) launch is a single thread; a subgroup barrier would
+ // instead synchronize N workers.
if (isSingleThreadWorkerLaunch())
return;
- // A worker-only launch keeps its rows unaligned, so a row narrower than a
- // subgroup covers only part of one. Synchronize exactly that lane range and
- // leave hasThreadYBarrier alone, so the launch is not padded on account of
- // this barrier.
+ // A worker-only row narrower than a subgroup covers only part of one, so
+ // synchronize just its lanes. Both paths leave hasThreadYBarrier unset, so
+ // the launch is not padded on account of this barrier.
if (isWorkerOnlyShuffleLaunch() && staticBlockDimX < options.subgroupSize) {
emitGPUBarrierRow(rewriter, loc, staticBlockDimX, options.subgroupSize);
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 f9c6288b33052..32785bc3cc448 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
@@ -303,9 +303,8 @@ func.func @thread_y_reduction_single_thread_rows() {
return
}
-// A worker-only launch with several threads per worker keeps its shape too.
-// Worker partials are combined in the lowest ThreadY threads of the block, so
-// the rows do not have to start at subgroup boundaries.
+// A worker-only launch with several threads per worker keeps its shape too:
+// the partials are combined in the lowest ThreadY threads of the block.
//
// CHECK-LABEL: func.func @thread_y_reduction_narrow_rows
// CHECK: %[[C8_NARROW:.*]] = arith.constant 8 : index
@@ -330,9 +329,8 @@ func.func @thread_y_reduction_narrow_rows() {
return
}
-// A ThreadX reduction in the same region shuffles within a row, so the rows are
-// aligned again: blockDim.x is padded to a subgroup and blockDim.y divided by
-// the same factor.
+// A ThreadX reduction in the same region shuffles within a row, so the rows
+// are aligned again and blockDim.y is divided by the same factor.
//
// CHECK-LABEL: func.func @thread_y_reduction_with_thread_x_reduction
// CHECK: %[[C32_BOTH:.*]] = arith.constant 32 : index
More information about the Mlir-commits
mailing list