[Mlir-commits] [mlir] 195df3d - [mlir][OpenACC] Privatize worker reduction accumulators per worker (#210456)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 20 03:14:22 PDT 2026
Author: Matsu
Date: 2026-07-20T03:14:16-07:00
New Revision: 195df3d4654b43ecf35927e574a8061bfaecbb78
URL: https://github.com/llvm/llvm-project/commit/195df3d4654b43ecf35927e574a8061bfaecbb78
DIFF: https://github.com/llvm/llvm-project/commit/195df3d4654b43ecf35927e574a8061bfaecbb78.diff
LOG: [mlir][OpenACC] Privatize worker reduction accumulators per worker (#210456)
Example:
```fortran
!$acc parallel
!$acc loop gang reduction(+:sum)
do k = 1, p
!$acc loop worker reduction(+:sum)
do j = 1, n
!$acc loop vector reduction(+:sum)
do i = 1, m
sum = sum + a(i,j,k)
end do
end do
end do
```
In this code, the worker accumulator is both gang- and worker-scoped.
Shared
memory allocated only one copy, so workers raced while updating the same
value.
Fix: allocate one shared-memory accumulator per worker whenever worker
parallelism is present.
Added:
mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-private.mlir
Modified:
mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index cd2e732dbc817..5f4af6a5bcc5b 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -319,8 +319,8 @@ collectPrivateLocalParDims(PrivateLocalOp privateLocal,
static FailureOr<std::optional<int64_t>> getWorkerPrivateSharedMemoryNumCopies(
PrivateLocalOp privateLocal, ComputeRegionOp computeRegion,
- bool isWorkerPrivate, bool isGangPrivate, OpenACCSupport *support) {
- if (!isWorkerPrivate || isGangPrivate)
+ bool isWorkerPrivate, OpenACCSupport *support) {
+ if (!isWorkerPrivate)
return std::optional<int64_t>(1);
GPUParallelDimAttr threadY =
@@ -394,8 +394,8 @@ FailureOr<bool> isPrivateLocalSharedMemoryCandidate(
return false;
FailureOr<std::optional<int64_t>> numCopies =
- getWorkerPrivateSharedMemoryNumCopies(
- privateLocal, computeRegion, isWorkerPrivate, isGangPrivate, support);
+ getWorkerPrivateSharedMemoryNumCopies(privateLocal, computeRegion,
+ isWorkerPrivate, support);
if (failed(numCopies))
return failure();
return numCopies->has_value();
@@ -411,15 +411,12 @@ std::optional<int64_t> getPrivateLocalSharedMemoryUpperBoundBytes(
SmallVector<GPUParallelDimAttr> parDims =
collectPrivateLocalParDims(privateLocal, computeRegion);
- bool isGangPrivate =
- llvm::any_of(parDims, [&](auto parDim) { return policy.isGang(parDim); });
bool isWorkerPrivate = llvm::any_of(
parDims, [&](auto parDim) { return policy.isWorker(parDim); });
FailureOr<std::optional<int64_t>> numCopies =
- getWorkerPrivateSharedMemoryNumCopies(privateLocal, computeRegion,
- isWorkerPrivate, isGangPrivate,
- /*support=*/nullptr);
+ getWorkerPrivateSharedMemoryNumCopies(
+ privateLocal, computeRegion, isWorkerPrivate, /*support=*/nullptr);
if (failed(numCopies) || !numCopies->has_value())
return std::nullopt;
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-private.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-private.mlir
new file mode 100644
index 0000000000000..3866c8057f129
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-worker-reduction-private.mlir
@@ -0,0 +1,41 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" | FileCheck %s
+
+// CHECK-LABEL: func.func @worker_reduction_private
+// CHECK: acc.gpu_shared_memory
+// CHECK-SAME: num_copies = 4
+
+func.func @worker_reduction_private() {
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %c32 = arith.constant 32 : index
+ %block = acc.par_width %c2 {par_dim = #acc.par_dim<block_x>}
+ %worker = acc.par_width %c4 {par_dim = #acc.par_dim<thread_y>}
+ %vector = acc.par_width %c32 {par_dim = #acc.par_dim<thread_x>}
+ acc.kernel_environment {
+ %private = acc.privatize [#acc<par_dims[thread_y]>] : () -> !acc.private_type<memref<i32>>
+ acc.compute_region launch(%bx = %block, %wy = %worker, %vx = %vector)
+ ins(%private_arg = %private) : (!acc.private_type<memref<i32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ scf.parallel (%block_iv) = (%c0) to (%bx) step (%c1) {
+ scf.parallel (%k) = (%block_iv) to (%bx) step (%bx) {
+ %local = acc.private_local %private_arg
+ {acc.par_dims = #acc<par_dims[thread_y]>}
+ : (!acc.private_type<memref<i32>>) -> memref<i32>
+ memref.store %c0_i32, %local[] : memref<i32>
+ scf.parallel (%worker_iv) = (%c0) to (%wy) step (%c1) {
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_y]>}
+ %value = memref.load %local[] : memref<i32>
+ acc.reduction_accumulate %value to %local <add>
+ : i32 -> memref<i32> {par_dims = #acc<par_dims[thread_y]>}
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[sequential]>}
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x]>}
+ acc.yield
+ } {origin = "acc.parallel"}
+ }
+ return
+}
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 66e500c4d2100..69f529abf6fa6 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -42,11 +42,13 @@ class OpenACCUtilsCGTest : public ::testing::Test {
static ComputeRegionOp buildComputeRegionWithPrivateLocal(
MLIRContext &context, OpBuilder &b, Location loc, ModuleOp module,
GPUParallelDimsAttr privatizeParDims, ValueRange launchArgs,
- PrivateLocalOp &privateLocalOut, PrivatizeOp &privatizeOut) {
+ PrivateLocalOp &privateLocalOut, PrivatizeOp &privatizeOut,
+ MemRefType memTy = {}, bool addReductionAccumulator = false) {
IRRewriter rewriter(&context);
rewriter.setInsertionPointToStart(module.getBody());
- MemRefType memTy = MemRefType::get({4}, b.getI32Type());
+ if (!memTy)
+ memTy = MemRefType::get({4}, b.getI32Type());
Type privateTy = PrivateType::get(&context, memTy);
privatizeOut = PrivatizeOp::create(rewriter, loc, privateTy, ValueRange{},
privatizeParDims);
@@ -71,7 +73,20 @@ class OpenACCUtilsCGTest : public ::testing::Test {
setParDimsAttr(par, GPUParallelDimsAttr::get(&context, {parDim}));
srcBuilder.setInsertionPoint(par.getBody()->getTerminator());
}
- PrivateLocalOp::create(srcBuilder, loc, memTy, privArg);
+ PrivateLocalOp privateLocal =
+ PrivateLocalOp::create(srcBuilder, loc, memTy, privArg);
+ if (addReductionAccumulator) {
+ Value partial =
+ arith::ConstantIntOp::create(srcBuilder, loc, b.getI32Type(), 1);
+ SmallVector<GPUParallelDimAttr> reductionDims;
+ for (GPUParallelDimAttr parDim : privatizeParDims.getArray())
+ if (!parDim.isAnyBlock())
+ reductionDims.push_back(parDim);
+ ReductionAccumulateOp::create(
+ srcBuilder, loc, partial, privateLocal.getResult(),
+ ReductionOperator::AccAdd,
+ GPUParallelDimsAttr::get(&context, reductionDims));
+ }
IRMapping mapping;
auto cr = buildComputeRegion(
@@ -748,6 +763,35 @@ TEST_F(OpenACCUtilsCGTest,
EXPECT_TRUE(*isCandidate);
}
+TEST_F(OpenACCUtilsCGTest, getSharedMemoryBytesGangWorkerReductionAccumulator) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+ GPUParallelDimsAttr gangWorkerDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context),
+ GPUParallelDimAttr::threadYDim(&context)});
+ auto c2 = arith::ConstantIndexOp::create(b, loc, 2);
+ auto c4 = arith::ConstantIndexOp::create(b, loc, 4);
+ auto bx =
+ ParWidthOp::create(b, loc, c2, GPUParallelDimAttr::blockXDim(&context));
+ auto ty =
+ ParWidthOp::create(b, loc, c4, GPUParallelDimAttr::threadYDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ MemRefType scalarTy = MemRefType::get({}, b.getI32Type());
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, gangWorkerDims,
+ ValueRange{bx.getResult(), ty.getResult()}, privateLocal, privatize,
+ scalarTy, /*addReductionAccumulator=*/true);
+
+ DefaultACCToGPUMappingPolicy policy;
+ std::optional<int64_t> upperBound =
+ getPrivateLocalSharedMemoryUpperBoundBytes(privateLocal, cr, *module,
+ policy);
+ ASSERT_TRUE(upperBound.has_value());
+ EXPECT_EQ(*upperBound, 16);
+}
+
TEST_F(OpenACCUtilsCGTest,
isPrivateLocalSharedMemoryCandidateWorkerPrivateDynamicFails) {
OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
More information about the Mlir-commits
mailing list