[Mlir-commits] [mlir] f4b5fb1 - [mlir][OpenACC] Support partial ThreadX reduction launches (#211665)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 23 14:54:55 PDT 2026
Author: Matsu
Date: 2026-07-23T14:54:50-07:00
New Revision: f4b5fb1569eb944ebe18e52c83e6ad9b0239f53c
URL: https://github.com/llvm/llvm-project/commit/f4b5fb1569eb944ebe18e52c83e6ad9b0239f53c
DIFF: https://github.com/llvm/llvm-project/commit/f4b5fb1569eb944ebe18e52c83e6ad9b0239f53c.diff
LOG: [mlir][OpenACC] Support partial ThreadX reduction launches (#211665)
Example:
```fortran
!$acc parallel loop collapse(2) num_gangs(10) reduction(+:a)
do i = 1, n
do j = 1, n
do k = 1, m
a(k) = a(k) + x(k,i) * y(k,j)
end do
end do
end do
```
In this code, ACCCG pads ThreadX to the subgroup size because
`gpu.all_reduce` uses subgroup shuffles. This unnecessarily increases
private reduction storage when only one Y/Z row is active.
Fix: preserve partial ThreadX launches for a single row, while retaining
alignment when Y/Z rows could share a subgroup.
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 618e45a803146..3bb78cd834b3e 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -1032,12 +1032,16 @@ LogicalResult ACCCGToGPULowering::rewrite() {
// - Subgroup reductions (gpu.all_reduce) require full subgroups
// - Per-row workgroup barriers require blockDim.x aligned to subgroupSize
bool isShuffleEnabled = 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.isThreadX() || parDim.isThreadY()) {
+ if (parDim.isThreadY() ||
+ (alignThreadXReduction && parDim.isThreadX())) {
// Shuffle are enabled. Need to adjust the ThreadX length.
isShuffleEnabled = true;
return WalkResult::interrupt();
@@ -1055,7 +1059,8 @@ LogicalResult ACCCGToGPULowering::rewrite() {
ArrayRef<mlir::acc::GPUParallelDimAttr> parDims =
mlir::acc::getParDimsAttr(allReduce).getArray();
for (auto parDim : parDims) {
- if (parDim.isThreadX() || parDim.isThreadY()) {
+ if (parDim.isThreadY() ||
+ (alignThreadXReduction && parDim.isThreadX())) {
isShuffleEnabled = true;
return WalkResult::interrupt();
}
@@ -1073,6 +1078,7 @@ LogicalResult ACCCGToGPULowering::rewrite() {
Value curBlockDimX = launch.getBlockSizeX();
Value curBlockDimY = launch.getBlockSizeY();
+ Value curBlockDimZ = launch.getBlockSizeZ();
// Emit a report on changing parallelism.
accSupport.emitRemark(computeRegion, [&]() {
@@ -1099,38 +1105,47 @@ LogicalResult ACCCGToGPULowering::rewrite() {
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 == 1). If blockDim.y > 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.
+ // 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 && *constBlockDimX > 1 &&
- *constBlockDimX < subgroupSize && *constBlockDimY == 1) {
+ if (constBlockDimX && constBlockDimY && constBlockDimZ &&
+ *constBlockDimX > 1 && *constBlockDimX < subgroupSize &&
+ *constBlockDimY == 1 && *constBlockDimZ == 1) {
skipAlign = true;
}
- // Update both the ThreadX length and the number of ThreadY.
- // When the original blockDim.x and blockDim.y are compile-time
+ // 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
// that the GpuKernelOutliningPass can set `known_block_size` on the
// outlined gpu.func.
- Value newBlockDimX, newBlockDimY;
- if (constBlockDimX && constBlockDimY) {
+ Value newBlockDimX, newBlockDimY, newBlockDimZ;
+ if (constBlockDimX && constBlockDimY && constBlockDimZ) {
int64_t bdx = *constBlockDimX;
int64_t bdy = *constBlockDimY;
+ int64_t bdz = *constBlockDimZ;
int64_t alignedBdx =
((bdx + subgroupAlignMask) / subgroupSize) * subgroupSize;
- int64_t numThreads = bdx * bdy;
- int64_t newBdy = std::max<int64_t>(1, numThreads / alignedBdx);
+ int64_t numXYThreads = bdx * bdy;
+ int64_t numThreads = numXYThreads * bdz;
+ int64_t newBdy = std::max<int64_t>(1, numXYThreads / alignedBdx);
+ int64_t newBdz =
+ std::max<int64_t>(1, numThreads / (alignedBdx * newBdy));
newBlockDimX =
arith::ConstantIndexOp::create(rewriter, loc, alignedBdx);
newBlockDimY = arith::ConstantIndexOp::create(rewriter, loc, newBdy);
+ newBlockDimZ = arith::ConstantIndexOp::create(rewriter, loc, newBdz);
} else {
- // numThreads = blockDim.x * blockDim.y
- Value numThreads =
+ // numXYThreads = blockDim.x * blockDim.y
+ Value numXYThreads =
arith::MulIOp::create(rewriter, loc, curBlockDimX, curBlockDimY);
+ Value numThreads =
+ arith::MulIOp::create(rewriter, loc, numXYThreads, curBlockDimZ);
// blockDim.x = ((blockDim.x + mask) / subgroupSize) * subgroupSize
Value cstMask =
arith::ConstantIndexOp::create(rewriter, loc, subgroupAlignMask);
@@ -1142,16 +1157,23 @@ LogicalResult ACCCGToGPULowering::rewrite() {
arith::DivUIOp::create(rewriter, loc, padded, cstSubgroupSize);
newBlockDimX = arith::MulIOp::create(rewriter, loc, subgroupsRequired,
cstSubgroupSize);
- // blockDim.y = max(1, numThreads / blockDim.x)
+ // blockDim.y = max(1, numXYThreads / blockDim.x)
Value quotient =
- arith::DivUIOp::create(rewriter, loc, numThreads, newBlockDimX);
+ arith::DivUIOp::create(rewriter, loc, numXYThreads, newBlockDimX);
Value cst1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
newBlockDimY = arith::MaxUIOp::create(rewriter, loc, cst1, quotient);
+ // blockDim.z = max(1, numThreads / (blockDim.x * blockDim.y))
+ Value newNumXYThreads =
+ arith::MulIOp::create(rewriter, loc, newBlockDimX, newBlockDimY);
+ quotient =
+ arith::DivUIOp::create(rewriter, loc, numThreads, newNumXYThreads);
+ newBlockDimZ = arith::MaxUIOp::create(rewriter, loc, cst1, quotient);
}
if (!skipAlign) {
launch.getBlockSizeXMutable().assign(newBlockDimX);
launch.getBlockSizeYMutable().assign(newBlockDimY);
+ launch.getBlockSizeZMutable().assign(newBlockDimZ);
}
}
}
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 422e172e1ec5b..93354321865a7 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
@@ -255,3 +255,95 @@ func.func @rank_two_partial_bounds_strided_layout() {
} {origin = "acc.parallel"}
return
}
+
+// CHECK-LABEL: func.func @partial_thread_x_reduction
+// CHECK: %[[C16:.*]] = arith.constant 16 : index
+// CHECK-NOT: arith.constant 31 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C16]],
+// CHECK: gpu.all_reduce add
+func.func @partial_thread_x_reduction() {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c16 {par_dim = #acc.par_dim<thread_x>}
+ acc.compute_region launch(%kbx = %bx, %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>
+ : memref<2xi32> {par_dims = #acc<par_dims[block_x, thread_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// CHECK-LABEL: func.func @thread_y_reduction_still_aligned
+// CHECK: %[[C32:.*]] = arith.constant 32 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32]],
+func.func @thread_y_reduction_still_aligned() {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %ty = acc.par_width %c16 {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.
+//
+// CHECK-LABEL: func.func @thread_x_reduction_with_thread_y_width
+// CHECK: %[[C32_ROWS:.*]] = arith.constant 32 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C32_ROWS]],
+func.func @thread_x_reduction_with_thread_y_width() {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c16 = arith.constant 16 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c16 {par_dim = #acc.par_dim<thread_x>}
+ %ty = acc.par_width %c2 {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_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
+
+// ThreadZ rows have the same subgroup-packing constraint as ThreadY rows.
+//
+// CHECK-LABEL: func.func @thread_x_reduction_with_thread_z_width
+// CHECK: %[[C352_Z_ROWS:.*]] = arith.constant 352 : index
+// CHECK: %[[C2_Z_ROWS:.*]] = arith.constant 2 : index
+// CHECK: gpu.launch
+// CHECK-SAME: threads({{.*}}) in (%{{.*}} = %[[C352_Z_ROWS]], %{{.*}} = %{{.*}}, %{{.*}} = %[[C2_Z_ROWS]])
+func.func @thread_x_reduction_with_thread_z_width() {
+ %c1 = arith.constant 1 : index
+ %c3 = arith.constant 3 : index
+ %c341 = arith.constant 341 : index
+ %bx = acc.par_width %c1 {par_dim = #acc.par_dim<block_x>}
+ %tx = acc.par_width %c341 {par_dim = #acc.par_dim<thread_x>}
+ %tz = acc.par_width %c3 {par_dim = #acc.par_dim<thread_z>}
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx, %ktz = %tz) {
+ %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_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ return
+}
More information about the Mlir-commits
mailing list