[Mlir-commits] [mlir] [mlir][acc] ACCComputeLowering needs to account for device_type par (PR #201267)
Razvan Lupusoru
llvmlistbot at llvm.org
Tue Jun 2 21:28:04 PDT 2026
https://github.com/razvanlupusoru created https://github.com/llvm/llvm-project/pull/201267
When assigning parallelism for compute constructs or loops, device_type parallelism must be first considered as a group for all available (gang, worker, vector) - if any of these have device_type setting, then those are the only ones that should be considered. Only if the loop has no device_type specific parallelism then default parallelism should be assigned.
>From a9b27a67ff770c9bb8e9952f0c79314138e6f07e Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Tue, 2 Jun 2026 14:41:04 -0700
Subject: [PATCH] [mlir][acc] ACCComputeLowering needs to account for
device_type par
When assigning parallelism for compute constructs or loops, device_type
parallelism must be first considered as a group for all available (gang,
worker, vector) - if any of these have device_type setting, then those
are the only ones that should be considered. Only if the loop has no
device_type specific parallelism then default parallelism should be
assigned.
---
.../mlir/Dialect/OpenACC/OpenACCOps.td | 12 +++++
mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp | 44 +++++++++++++++++--
.../OpenACC/Transforms/ACCComputeLowering.cpp | 27 +++++++++---
...-compute-lowering-compute-device-type.mlir | 25 +++++++++++
...acc-compute-lowering-loop-device-type.mlir | 23 ++++++++++
5 files changed, 122 insertions(+), 9 deletions(-)
create mode 100644 mlir/test/Dialect/OpenACC/acc-compute-lowering-compute-device-type.mlir
create mode 100644 mlir/test/Dialect/OpenACC/acc-compute-lowering-loop-device-type.mlir
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 32ecaa6bc2d42..26a9344fa3a15 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -1829,6 +1829,10 @@ def OpenACC_ParallelOp
mlir::Operation::operand_range
getNumGangsValues(mlir::acc::DeviceType deviceType);
+ /// Return true if the op has any num_gangs, num_workers, or vector_length
+ /// clause for the given device_type.
+ bool hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType);
+
/// Return true if the op has the wait attribute for the
/// mlir::acc::DeviceType::None device_type.
bool hasWaitOnly();
@@ -2156,6 +2160,10 @@ def OpenACC_KernelsOp
mlir::Operation::operand_range
getNumGangsValues(mlir::acc::DeviceType deviceType);
+ /// Return true if the op has any num_gangs, num_workers, or vector_length
+ /// clause for the given device_type.
+ bool hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType);
+
/// Return true if the op has the wait attribute for the
/// mlir::acc::DeviceType::None device_type.
bool hasWaitOnly();
@@ -2808,6 +2816,10 @@ def OpenACC_LoopOp
// 'default'/None device-type.
bool hasDefaultGangWorkerVector();
+ // Return whether this LoopOp has a gang, worker, or vector for the given
+ // device-type.
+ bool hasAnyGangWorkerVector(DeviceType deviceType);
+
// Used to obtain the parallelism mode for the requested device type.
// This first checks if the mode is set for the device_type requested.
// And if not, it returns the non-device_type mode.
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 449a9b588910f..8462c346fd7c6 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -2173,6 +2173,31 @@ ParallelOp::getNumGangsValues(mlir::acc::DeviceType deviceType) {
getNumGangsSegments(), deviceType);
}
+static bool hasAnyGangWorkerVectorForDeviceType(
+ std::optional<mlir::ArrayAttr> numGangsDeviceType,
+ mlir::Operation::operand_range numGangs,
+ std::optional<llvm::ArrayRef<int32_t>> numGangsSegments,
+ std::optional<mlir::ArrayAttr> numWorkersDeviceType,
+ mlir::Operation::operand_range numWorkers,
+ std::optional<mlir::ArrayAttr> vectorLengthDeviceType,
+ mlir::Operation::operand_range vectorLength,
+ mlir::acc::DeviceType deviceType) {
+ return !getValuesFromSegments(numGangsDeviceType, numGangs, numGangsSegments,
+ deviceType)
+ .empty() ||
+ getValueInDeviceTypeSegment(numWorkersDeviceType, numWorkers,
+ deviceType) ||
+ getValueInDeviceTypeSegment(vectorLengthDeviceType, vectorLength,
+ deviceType);
+}
+
+bool acc::ParallelOp::hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType) {
+ return hasAnyGangWorkerVectorForDeviceType(
+ getNumGangsDeviceType(), getNumGangs(), getNumGangsSegments(),
+ getNumWorkersDeviceType(), getNumWorkers(), getVectorLengthDeviceType(),
+ getVectorLength(), deviceType);
+}
+
bool acc::ParallelOp::hasWaitOnly() {
return hasWaitOnly(mlir::acc::DeviceType::None);
}
@@ -3039,6 +3064,13 @@ KernelsOp::getNumGangsValues(mlir::acc::DeviceType deviceType) {
getNumGangsSegments(), deviceType);
}
+bool acc::KernelsOp::hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType) {
+ return hasAnyGangWorkerVectorForDeviceType(
+ getNumGangsDeviceType(), getNumGangs(), getNumGangsSegments(),
+ getNumWorkersDeviceType(), getNumWorkers(), getVectorLengthDeviceType(),
+ getVectorLength(), deviceType);
+}
+
bool acc::KernelsOp::hasWaitOnly() {
return hasWaitOnly(mlir::acc::DeviceType::None);
}
@@ -3968,9 +4000,15 @@ bool acc::LoopOp::hasParallelismFlag(DeviceType dt) {
}
bool acc::LoopOp::hasDefaultGangWorkerVector() {
- return hasVector() || getVectorValue() || hasWorker() || getWorkerValue() ||
- hasGang() || getGangValue(GangArgType::Num) ||
- getGangValue(GangArgType::Dim) || getGangValue(GangArgType::Static);
+ return hasAnyGangWorkerVector(DeviceType::None);
+}
+
+bool acc::LoopOp::hasAnyGangWorkerVector(DeviceType deviceType) {
+ return hasVector(deviceType) || getVectorValue(deviceType) ||
+ hasWorker(deviceType) || getWorkerValue(deviceType) ||
+ hasGang(deviceType) || getGangValue(GangArgType::Num, deviceType) ||
+ getGangValue(GangArgType::Dim, deviceType) ||
+ getGangValue(GangArgType::Static, deviceType);
}
acc::LoopParMode
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
index 80b4570587a17..7e1bf74cc3728 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
@@ -185,11 +185,30 @@ static void insertParDim(SmallVectorImpl<GPUParallelDimAttr> &parDims,
parDims.insert(lb, parDim);
}
+/// Return the device type from which gang/worker/vector clauses should be read.
+/// If the requested device type has any such clauses, use that exclusively;
+/// otherwise fall back to the default (DeviceType::None).
+static DeviceType getGangWorkerVectorDeviceType(LoopOp loopOp,
+ DeviceType deviceType) {
+ if (deviceType != DeviceType::None && loopOp.hasAnyGangWorkerVector(deviceType))
+ return deviceType;
+ return DeviceType::None;
+}
+
+template <typename ComputeConstructT>
+static DeviceType getParDimsDeviceType(ComputeConstructT computeOp,
+ DeviceType deviceType) {
+ if (deviceType != DeviceType::None && computeOp.hasAnyGangWorkerVector(deviceType))
+ return deviceType;
+ return DeviceType::None;
+}
+
/// Map loop parallelism clauses (gang/worker/vector) to GPU parallel
/// dimensions using the given mapping policy.
static SmallVector<GPUParallelDimAttr>
getParallelDimensions(LoopOp loopOp, const ACCToGPUMappingPolicy &policy,
DeviceType deviceType) {
+ deviceType = getGangWorkerVectorDeviceType(loopOp, deviceType);
SmallVector<GPUParallelDimAttr> parDims;
auto *ctx = loopOp->getContext();
@@ -229,12 +248,12 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
if (isEffectivelySerial(computeOp))
return {ParWidthOp::create(rewriter, loc, Value(), policy.seqDim(ctx))};
+ deviceType = getParDimsDeviceType(computeOp, deviceType);
+
SmallVector<Value> values;
auto indexTy = rewriter.getIndexType();
auto numGangs = computeOp.getNumGangsValues(deviceType);
- if (numGangs.empty())
- numGangs = computeOp.getNumGangsValues();
for (auto [gangDimIdx, gangSize] : llvm::enumerate(numGangs)) {
auto gangLevel = getGangParLevel(gangDimIdx + 1);
values.push_back(ParWidthOp::create(
@@ -245,8 +264,6 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
}
Value numWorkers = computeOp.getNumWorkersValue(deviceType);
- if (!numWorkers)
- numWorkers = computeOp.getNumWorkersValue();
if (numWorkers) {
values.push_back(ParWidthOp::create(
rewriter, loc,
@@ -256,8 +273,6 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
}
Value vectorLength = computeOp.getVectorLengthValue(deviceType);
- if (!vectorLength)
- vectorLength = computeOp.getVectorLengthValue();
if (vectorLength) {
values.push_back(ParWidthOp::create(
rewriter, loc,
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute-device-type.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute-device-type.mlir
new file mode 100644
index 0000000000000..572235b89106c
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute-device-type.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s -acc-compute-lowering=device-type=nvidia | FileCheck %s
+
+// Default num_gangs, nvidia vector_length: with device-type=nvidia only vector applies.
+// CHECK-LABEL: func.func @parallel_default_gangs_nvidia_vector_length
+func.func @parallel_default_gangs_nvidia_vector_length(%buf: memref<4xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c4_i32 = arith.constant 4 : i32
+ %c32_i32 = arith.constant 32 : i32
+
+ %dev = acc.copyin varPtr(%buf : memref<4xi32>) -> memref<4xi32>
+ // CHECK-NOT: acc.par_width {{.*}} {par_dim = #acc.par_dim<block_x>}
+ // CHECK: acc.par_width {{.*}} {par_dim = #acc.par_dim<thread_x>}
+ acc.parallel num_gangs({%c4_i32 : i32}) vector_length(%c32_i32 : i32 [#acc.device_type<nvidia>]) dataOperands(%dev : memref<4xi32>) {
+ acc.loop control(%i : index) = (%c0 : index) to (%c4 : index) step (%c1 : index) {
+ %vi = arith.index_cast %i : index to i32
+ memref.store %vi, %dev[%i] : memref<4xi32>
+ acc.yield
+ } attributes {independent = [#acc.device_type<none>]}
+ acc.yield
+ }
+ acc.copyout accPtr(%dev : memref<4xi32>) to varPtr(%buf : memref<4xi32>)
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop-device-type.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop-device-type.mlir
new file mode 100644
index 0000000000000..fb7c9b35e0909
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop-device-type.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s -acc-compute-lowering=device-type=nvidia | FileCheck %s
+
+// Gang on default, vector on nvidia: with device-type=nvidia only vector applies.
+// CHECK-LABEL: func.func @parallel_loop_gang_default_vector_nvidia
+func.func @parallel_loop_gang_default_vector_nvidia(%buf: memref<1xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1_i32 = arith.constant 1 : 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.par_dims = #acc<par_dims[block_x]>
+ // CHECK: acc.par_dims = #acc<par_dims[thread_x]>
+ acc.parallel num_gangs({%c10_i32 : i32}) dataOperands(%dev : memref<1xi32>) {
+ acc.loop gang control(%arg0 : i32) = (%c1_i32 : i32) to (%c100_i32 : i32) step (%c1_i32 : i32) {
+ memref.store %arg0, %dev[%c0] : memref<1xi32>
+ acc.yield
+ } attributes {auto_ = [#acc.device_type<none>], gang = [#acc.device_type<none>], vector = [#acc.device_type<nvidia>]}
+ acc.yield
+ }
+ acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
+ return
+}
More information about the Mlir-commits
mailing list