[Mlir-commits] [mlir] 9f00966 - [mlir][OpenACC] Record gang(static:) chunk size on lowered loops (#227037)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 28 15:46:11 PDT 2026
Author: Delaram Talaashrafi
Date: 2026-09-28T18:46:04-04:00
New Revision: 9f009662f3fba1d8955e5e79261c60c2f98d24b6
URL: https://github.com/llvm/llvm-project/commit/9f009662f3fba1d8955e5e79261c60c2f98d24b6
DIFF: https://github.com/llvm/llvm-project/commit/9f009662f3fba1d8955e5e79261c60c2f98d24b6.diff
LOG: [mlir][OpenACC] Record gang(static:) chunk size on lowered loops (#227037)
Lower OpenACC `gang(static:)` through compute lowering as
`acc.chunk_size`. A constant size is stored directly; `static:*` and a
non-constant size are recorded as -1. `gang(static:)` also maps to gang
dimension 1.
Added:
Modified:
mlir/include/mlir/Dialect/OpenACC/OpenACCCGAttributes.td
mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGAttributes.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGAttributes.td
index c401f5203ec99..c16960f604bc1 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCCGAttributes.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCCGAttributes.td
@@ -7,8 +7,8 @@
//===----------------------------------------------------------------------===//
//
// Defines OpenACC codegen attributes (GPU parallel dimensions used for
-// privatization, barrier management, and loop work-sharing, plus active
-// parallel dimensions for predication analysis).
+// privatization, barrier management, and loop work-sharing, active
+// parallel dimensions for predication analysis, and loop chunk size).
//
//===----------------------------------------------------------------------===//
@@ -103,6 +103,15 @@ def OpenACC_GPUBlockRedundantAttr
}];
}
+def OpenACC_ChunkSizeAttr : OpenACC_Attr<"ChunkSize", "chunk_size"> {
+ let summary = "Static gang chunk size for an OpenACC loop.";
+ let description = [{
+ Records the constant chunk size from an OpenACC `gang(static:N)` clause.
+ }];
+ let parameters = (ins "int64_t":$chunk_size);
+ let assemblyFormat = "`<` $chunk_size `>`";
+}
+
def OpenACC_ActiveParDimsAttr
: OpenACC_Attr<"ActiveParDims", "active_par_dims"> {
let summary = "Active GPU parallel dimensions for an OpenACC operation.";
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
index bb0af06c01421..3a8ed91a47b2d 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
@@ -118,6 +118,21 @@ bool hasGPUBlockRedundantAttr(Operation *op);
/// Mark \p op with the `acc.gpu_block_redundant` attribute.
void setGPUBlockRedundantAttr(Operation *op);
+/// Obtain the `acc.chunk_size` attribute carried by \p op, if any.
+ChunkSizeAttr getChunkSizeAttr(Operation *op);
+
+/// Return whether \p op carries an `acc.chunk_size` attribute.
+bool hasChunkSizeAttr(Operation *op);
+
+/// Set the `acc.chunk_size` attribute on \p op.
+void setChunkSizeAttr(Operation *op, ChunkSizeAttr attr);
+
+/// Set the `acc.chunk_size` attribute on \p op from a constant size.
+void setChunkSizeAttr(Operation *op, int64_t chunkSize);
+
+/// Return the chunk size carried by \p op, if any.
+std::optional<int64_t> getChunkSize(Operation *op);
+
/// Create a gang dim 1 GPUParallelDimsAttr based on the mapping policy.
inline GPUParallelDimsAttr
getGangDim1ParDimsAttr(MLIRContext *ctx, ACCToGPUMappingPolicy &policy) {
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
index 832de220289ea..65d69f447cb2c 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
@@ -226,12 +226,28 @@ getParallelDimensions(LoopOp loopOp, const ACCToGPUMappingPolicy &policy,
insertParDim(parDims, policy.gangDim(ctx, gangLevel));
}
} else if (loopOp.hasGang(deviceType) ||
- loopOp.getGangValue(GangArgType::Num, deviceType)) {
+ loopOp.getGangValue(GangArgType::Num, deviceType) ||
+ loopOp.getGangValue(GangArgType::Static, deviceType)) {
insertParDim(parDims, policy.gangDim(ctx, ParLevel::gang_dim1));
}
return parDims;
}
+/// If \p loopOp has a `gang(static:)` value, return the corresponding
+/// `acc.chunk_size` attribute. A constant size is stored directly; a
+/// non-constant size is recorded as -1, the same encoding as `static:*`.
+/// Returns a null attribute when the clause is absent.
+static ChunkSizeAttr getChunkSizeFromStatic(LoopOp loopOp,
+ DeviceType deviceType) {
+ DeviceType loopDeviceType = getGangWorkerVectorDeviceType(loopOp, deviceType);
+ Value gangStatic = loopOp.getGangValue(GangArgType::Static, loopDeviceType);
+ if (!gangStatic)
+ return {};
+ std::optional<int64_t> chunkSize = getConstantIntValue(gangStatic);
+ // Non-constant static is treated like `static:*` (implementation-chosen).
+ return ChunkSizeAttr::get(loopOp->getContext(), chunkSize.value_or(-1));
+}
+
/// Build `acc.compute_region` launch operands: one sequential `acc.par_width`
/// for `acc.serial`, for `acc.parallel` / `acc.kernels` when every num_gangs
/// operand and num_workers / vector_length are the constant 1, and otherwise
@@ -355,6 +371,9 @@ class ACCLoopConversion : public OpRewritePattern<LoopOp> {
auto parDimsAttr =
GPUParallelDimsAttr::get(loopOp->getContext(), parDims);
setParDimsAttr(forOp, parDimsAttr);
+ if (ChunkSizeAttr chunkSize =
+ getChunkSizeFromStatic(loopOp, deviceType))
+ setChunkSizeAttr(forOp, chunkSize);
}
rewriter.replaceOp(loopOp, forOp);
} else if (!isOpInComputeRegion(loopOp) &&
@@ -380,8 +399,10 @@ class ACCLoopConversion : public OpRewritePattern<LoopOp> {
auto parDimsAttr =
GPUParallelDimsAttr::get(loopOp->getContext(), parDims);
setParDimsAttr(parallelOp, parDimsAttr);
+ if (ChunkSizeAttr chunkSize =
+ getChunkSizeFromStatic(loopOp, deviceType))
+ setChunkSizeAttr(parallelOp, chunkSize);
}
-
rewriter.replaceOp(loopOp, parallelOp);
}
return success();
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 6d302b86e475a..2ab44e447adfb 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -223,6 +223,26 @@ void setGPUBlockRedundantAttr(Operation *op) {
GPUBlockRedundantAttr::get(op->getContext()));
}
+ChunkSizeAttr getChunkSizeAttr(Operation *op) {
+ return op->getDiscardableAttrOfType<ChunkSizeAttr>(ChunkSizeAttr::name);
+}
+
+bool hasChunkSizeAttr(Operation *op) { return getChunkSizeAttr(op) != nullptr; }
+
+void setChunkSizeAttr(Operation *op, ChunkSizeAttr attr) {
+ op->setDiscardableAttr(ChunkSizeAttr::name, attr);
+}
+
+void setChunkSizeAttr(Operation *op, int64_t chunkSize) {
+ setChunkSizeAttr(op, ChunkSizeAttr::get(op->getContext(), chunkSize));
+}
+
+std::optional<int64_t> getChunkSize(Operation *op) {
+ if (ChunkSizeAttr attr = getChunkSizeAttr(op))
+ return attr.getChunkSize();
+ return std::nullopt;
+}
+
void copyParDimsAttr(Operation *from, Operation *to) {
assert(hasParDimsAttr(from) &&
"expected parallel dimensions attribute to already be set");
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
index 4da7ca8868655..1a529e8f102e8 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -255,3 +255,58 @@ func.func @orphan_loop_unroll_annotation(%buf: memref<8xi32>) {
} independent attributes {llvm.loop_annotation = #llvm.loop_annotation<unroll = <disable = false, full = true>>}
return
}
+
+// -----
+
+// Independent loop with gang(static:N): chunk size and gang par_dims on scf.parallel.
+// CHECK-LABEL: func.func @parallel_loop_gang_static
+func.func @parallel_loop_gang_static(%buf: memref<1xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1_i32 = arith.constant 1 : i32
+ %c32_i32 = arith.constant 32 : i32
+ %c10_i32 = arith.constant 10 : i32
+ %c100_i32 = arith.constant 100 : i32
+
+ %dev = acc.copyin varPtr(%buf : memref<1xi32>) -> memref<1xi32>
+ // CHECK-NOT: acc.parallel
+ // CHECK: acc.kernel_environment
+ // CHECK: acc.compute_region
+ // CHECK: scf.parallel
+ // CHECK: acc.chunk_size = #acc.chunk_size<32>
+ // CHECK-SAME: acc.par_dims = #acc<par_dims[block_x]>
+ acc.parallel num_gangs({%c10_i32 : i32}) dataOperands(%dev : memref<1xi32>) {
+ acc.loop gang({static=%c32_i32 : i32}) control(%arg0 : i32) = (%c1_i32 : i32) to (%c100_i32 : i32) step (%c1_i32 : i32) {
+ memref.store %arg0, %dev[%c0] : memref<1xi32>
+ acc.yield
+ } independent
+ acc.yield
+ }
+ acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
+ return
+}
+
+// -----
+
+// gang(static:*) is encoded as chunk size -1.
+// CHECK-LABEL: func.func @parallel_loop_gang_static_star
+func.func @parallel_loop_gang_static_star(%buf: memref<1xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1_i32 = arith.constant 1 : i32
+ %c_m1 = arith.constant -1 : index
+ %c10_i32 = arith.constant 10 : i32
+ %c100_i32 = arith.constant 100 : i32
+
+ %dev = acc.copyin varPtr(%buf : memref<1xi32>) -> memref<1xi32>
+ // CHECK: scf.parallel
+ // CHECK: acc.chunk_size = #acc.chunk_size<-1>
+ // CHECK-SAME: acc.par_dims = #acc<par_dims[block_x]>
+ acc.parallel num_gangs({%c10_i32 : i32}) dataOperands(%dev : memref<1xi32>) {
+ acc.loop gang({static=%c_m1 : index}) control(%arg0 : i32) = (%c1_i32 : i32) to (%c100_i32 : i32) step (%c1_i32 : i32) {
+ memref.store %arg0, %dev[%c0] : memref<1xi32>
+ acc.yield
+ } independent
+ acc.yield
+ }
+ acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
+ return
+}
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 097e34f6332ee..9de3158accee1 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -278,6 +278,22 @@ TEST_F(OpenACCUtilsCGTest, activeParDimsOperationAttributes) {
EXPECT_EQ(getActiveParDimsAttr(op), threadAttr);
}
+TEST_F(OpenACCUtilsCGTest, chunkSizeAttribute) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ Operation *op = module->getOperation();
+
+ EXPECT_FALSE(hasChunkSizeAttr(op));
+ EXPECT_FALSE(getChunkSize(op).has_value());
+
+ setChunkSizeAttr(op, 32);
+ EXPECT_TRUE(hasChunkSizeAttr(op));
+ EXPECT_EQ(getChunkSize(op), 32);
+ EXPECT_EQ(getChunkSizeAttr(op).getChunkSize(), 32);
+
+ setChunkSizeAttr(op, ChunkSizeAttr::get(&context, 64));
+ EXPECT_EQ(getChunkSize(op), 64);
+}
+
//===----------------------------------------------------------------------===//
// buildComputeRegion Tests
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list