[Mlir-commits] [flang] [mlir] [Flang][mlir][OpenMP] Add affinity clause to omp.task and Flang lowering (PR #179003)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jan 31 10:46:26 PST 2026
https://github.com/chichunchen updated https://github.com/llvm/llvm-project/pull/179003
>From 741cb9b185f887f911274491e02c277acaa5d31f Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Fri, 30 Jan 2026 15:34:45 -0600
Subject: [PATCH 1/2] [Flang][mlir][OpenMP] Add affinity clause to omp.task and
Flang lowering
- Add MLIR OpenMP affinity clause
- Lower flang task affinity to mlir
- Emit TODO for iterator modifier and update negative test
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 17 +++
flang/lib/Lower/OpenMP/ClauseProcessor.h | 1 +
flang/lib/Lower/OpenMP/OpenMP.cpp | 3 +-
.../Lower/OpenMP/Todo/affinity-clause.f90 | 4 +-
flang/test/Lower/OpenMP/task-affinity.f90 | 111 ++++++++++++++++++
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 23 ++++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 6 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 33 +++++-
mlir/test/Dialect/OpenMP/ops.mlir | 30 +++++
9 files changed, 218 insertions(+), 10 deletions(-)
create mode 100644 flang/test/Lower/OpenMP/task-affinity.f90
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 8094c6264b492..b1973a3b8bf06 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -710,6 +710,23 @@ static llvm::StringMap<bool> getTargetFeatures(mlir::ModuleOp module) {
return featuresMap;
}
+bool ClauseProcessor::processAffinity(
+ mlir::omp::AffinityClauseOps &result) const {
+ return findRepeatableClause<omp::clause::Affinity>(
+ [&](const omp::clause::Affinity &clause, const parser::CharBlock &) {
+ if (std::get<std::optional<omp::clause::Iterator>>(clause.t)) {
+ TODO(converter.getCurrentLocation(),
+ "Support for iterator modifiers is not implemented yet");
+ }
+
+ const auto &objects = std::get<omp::ObjectList>(clause.t);
+ if (!objects.empty())
+ genObjectList(objects, converter, result.affinityVars);
+
+ return true;
+ });
+}
+
static void
addAlignedClause(lower::AbstractConverter &converter,
const omp::clause::Aligned &clause,
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.h b/flang/lib/Lower/OpenMP/ClauseProcessor.h
index ba1764ce46821..ca9b28dfdd061 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.h
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.h
@@ -119,6 +119,7 @@ class ClauseProcessor {
bool processDetach(mlir::omp::DetachClauseOps &result) const;
// 'Repeatable' clauses: They can appear multiple times in the clause list.
+ bool processAffinity(mlir::omp::AffinityClauseOps &result) const;
bool processAligned(mlir::omp::AlignedClauseOps &result) const;
bool processAllocate(mlir::omp::AllocateClauseOps &result) const;
bool processCopyin() const;
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 43c6128d46647..3eeaab48d1447 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -1801,6 +1801,7 @@ static void genTaskClauses(
mlir::omp::TaskOperands &clauseOps,
llvm::SmallVectorImpl<const semantics::Symbol *> &inReductionSyms) {
ClauseProcessor cp(converter, semaCtx, clauses);
+ cp.processAffinity(clauseOps);
cp.processAllocate(clauseOps);
cp.processDepend(symTable, stmtCtx, clauseOps);
cp.processFinal(stmtCtx, clauseOps);
@@ -1810,8 +1811,6 @@ static void genTaskClauses(
cp.processPriority(stmtCtx, clauseOps);
cp.processUntied(clauseOps);
cp.processDetach(clauseOps);
-
- cp.processTODO<clause::Affinity>(loc, llvm::omp::Directive::OMPD_task);
}
static void genTaskgroupClauses(
diff --git a/flang/test/Lower/OpenMP/Todo/affinity-clause.f90 b/flang/test/Lower/OpenMP/Todo/affinity-clause.f90
index 3459dd219e425..6be477229286a 100644
--- a/flang/test/Lower/OpenMP/Todo/affinity-clause.f90
+++ b/flang/test/Lower/OpenMP/Todo/affinity-clause.f90
@@ -1,10 +1,10 @@
!RUN: %not_todo_cmd bbc -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
!RUN: %not_todo_cmd %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=50 -o - %s 2>&1 | FileCheck %s
-!CHECK: not yet implemented: Unhandled clause AFFINITY in TASK construct
+!CHECK: Support for iterator modifiers is not implemented yet
subroutine f00(x)
integer :: x(10)
-!$omp task affinity(x)
+!$omp task affinity(iterator(i = 1:10) : x(i))
x = x + 1
!$omp end task
end
diff --git a/flang/test/Lower/OpenMP/task-affinity.f90 b/flang/test/Lower/OpenMP/task-affinity.f90
new file mode 100644
index 0000000000000..66254e48e9b8e
--- /dev/null
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -0,0 +1,111 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=52 -o - %s | FileCheck %s
+
+! scalar element locator
+subroutine omp_task_affinity_elem()
+ implicit none
+ integer, parameter :: n = 100
+ integer :: a(n)
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(a(1))
+ a(1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine omp_task_affinity_elem
+
+! CHECK-LABEL: func.func @_QPomp_task_affinity_elem()
+! CHECK: %[[A1:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_elemEa"}
+! CHECK: omp.parallel {
+! CHECK: omp.single {
+! CHECK: omp.task affinity(%[[A1]]#0 : !fir.ref<!fir.array<100xi32>>) {
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: return
+
+! array section locator
+subroutine omp_task_affinity_array_section()
+ implicit none
+ integer, parameter :: n = 100
+ integer :: a(n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(a(2:50)) private(i)
+ do i = 2, 50
+ a(i) = i
+ end do
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine omp_task_affinity_array_section
+
+! CHECK-LABEL: func.func @_QPomp_task_affinity_array_section()
+! CHECK: %[[A2:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_array_sectionEa"}
+! CHECK: %[[I2:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_array_sectionEi"}
+! CHECK: omp.parallel {
+! CHECK: omp.single {
+! CHECK: omp.task affinity(%[[A2]]#0 : !fir.ref<!fir.array<100xi32>>) private({{.*}} %[[I2]]#0 -> %{{.*}} : !fir.ref<i32>) {
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: return
+
+! scalar variable locator
+subroutine omp_task_affinity_scalar()
+ implicit none
+ integer :: s
+ s = 7
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(s)
+ s = s + 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine omp_task_affinity_scalar
+
+! CHECK-LABEL: func.func @_QPomp_task_affinity_scalar()
+! CHECK: %[[S3:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_scalarEs"}
+! CHECK: omp.parallel {
+! CHECK: omp.single {
+! CHECK: omp.task affinity(%[[S3]]#0 : !fir.ref<i32>) {
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: return
+
+! multiple locators
+subroutine omp_task_affinity_multi()
+ implicit none
+ integer, parameter :: n = 100
+ integer :: a(n), b(n)
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(a(1), b(1))
+ a(2) = 2
+ b(2) = 2
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine omp_task_affinity_multi
+
+! CHECK-LABEL: func.func @_QPomp_task_affinity_multi()
+! CHECK: %[[A4:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_multiEa"}
+! CHECK: %[[B4:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_multiEb"}
+! CHECK: omp.parallel {
+! CHECK: omp.single {
+! CHECK: omp.task affinity(%[[A4]]#0 : !fir.ref<!fir.array<100xi32>>, %[[B4]]#0 : !fir.ref<!fir.array<100xi32>>) {
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: return
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 9114d9c1f0ac1..286b07941a03c 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -23,6 +23,29 @@ include "mlir/Dialect/OpenMP/OpenMPOpBase.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/IR/BuiltinAttributes.td"
+//===----------------------------------------------------------------------===//
+// V5.2: [12.5.1] `affinity` clause
+//===----------------------------------------------------------------------===//
+
+class OpenMP_AffinityClauseSkip<
+ bit traits = false, bit arguments = false, bit assemblyFormat = false,
+ bit description = false, bit extraClassDeclaration = false>
+ : OpenMP_Clause<traits, arguments, assemblyFormat, description,
+ extraClassDeclaration> {
+ let arguments = (ins Variadic<OpenMP_PointerLikeType>:$affinity_vars);
+
+ let optAssemblyFormat = [{
+ `affinity` `(` custom<AffinityClause>($affinity_vars, type($affinity_vars)) `)`
+ }];
+
+ let description = [{
+ The `affinity` clause specifies a locator list used as a hint for task
+ placement / scheduling affinity.
+ }];
+}
+
+def OpenMP_AffinityClause : OpenMP_AffinityClauseSkip<>;
+
//===----------------------------------------------------------------------===//
// V5.2: [6.3] `align` clause
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index dfec6609e1161..49a724fd5446e 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -905,9 +905,9 @@ def TaskOp
traits = [AttrSizedOperandSegments, AutomaticAllocationScope,
OutlineableOpenMPOpInterface],
clauses = [
- // TODO: Complete clause list (affinity, detach).
- OpenMP_AllocateClause, OpenMP_DependClause,
- OpenMP_FinalClause, OpenMP_IfClause,
+ // TODO: Complete clause list (detach).
+ OpenMP_AffinityClause, OpenMP_AllocateClause,
+ OpenMP_DependClause, OpenMP_FinalClause, OpenMP_IfClause,
OpenMP_InReductionClause, OpenMP_MergeableClause,
OpenMP_PriorityClause, OpenMP_PrivateClause,
OpenMP_UntiedClause, OpenMP_DetachClause],
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index c3916219d1c93..e8eebc2a1a4c6 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -3137,9 +3137,10 @@ LogicalResult DeclareReductionOp::verifyRegions() {
void TaskOp::build(OpBuilder &builder, OperationState &state,
const TaskOperands &clauses) {
MLIRContext *ctx = builder.getContext();
- TaskOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
- makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
- clauses.final, clauses.ifExpr, clauses.inReductionVars,
+ TaskOp::build(builder, state, clauses.affinityVars, clauses.allocateVars,
+ clauses.allocatorVars, makeArrayAttr(ctx, clauses.dependKinds),
+ clauses.dependVars, clauses.final, clauses.ifExpr,
+ clauses.inReductionVars,
makeDenseBoolArrayAttr(ctx, clauses.inReductionByref),
makeArrayAttr(ctx, clauses.inReductionSyms), clauses.mergeable,
clauses.priority, /*private_vars=*/clauses.privateVars,
@@ -4534,6 +4535,32 @@ static void printUniformClause(OpAsmPrinter &p, Operation *op,
}
}
+//===----------------------------------------------------------------------===//
+// Parser and printer for Affinity Clause
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseAffinityClause(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &affinityVars,
+ SmallVectorImpl<Type> &affinityTypes) {
+ return parser.parseCommaSeparatedList([&]() -> ParseResult {
+ if (parser.parseOperand(affinityVars.emplace_back()) ||
+ parser.parseColonType(affinityTypes.emplace_back()))
+ return failure();
+ return success();
+ });
+}
+
+static void printAffinityClause(OpAsmPrinter &p, Operation *op,
+ ValueRange affinityVars,
+ TypeRange affinityTypes) {
+ for (unsigned i = 0; i < affinityVars.size(); ++i) {
+ if (i)
+ p << ", ";
+ p << affinityVars[i] << " : " << affinityTypes[i];
+ }
+}
+
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 1e8c9bdeb33ae..5c2849cc9b5ea 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3520,3 +3520,33 @@ func.func @omp_declare_simd_all_clauses(%a: f64, %b: f64,
inbranch
return
}
+
+// CHECK-LABEL: func.func @task_affinity_single
+func.func @task_affinity_single() {
+ // CHECK: %[[A:.*]] = memref.alloca() : memref<100xi32>
+ // CHECK: omp.task affinity(%[[A]] : memref<100xi32>) {
+ // CHECK: omp.terminator
+ // CHECK: }
+ // CHECK: return
+ %a = memref.alloca() : memref<100xi32>
+ omp.task affinity(%a : memref<100xi32>) {
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @task_affinity_multi
+func.func @task_affinity_multi() {
+ // CHECK: %[[A:.*]] = memref.alloca() : memref<64xi32>
+ // CHECK: %[[B:.*]] = memref.alloca() : memref<8xf64>
+ // CHECK: omp.task affinity(%[[A]] : memref<64xi32>, %[[B]] : memref<8xf64>) {
+ // CHECK: omp.terminator
+ // CHECK: }
+ // CHECK: return
+ %a = memref.alloca() : memref<64xi32>
+ %b = memref.alloca() : memref<8xf64>
+ omp.task affinity(%a : memref<64xi32>, %b : memref<8xf64>) {
+ omp.terminator
+ }
+ return
+}
>From 762121b6bd2bf7190152a8906ebb07ae6839838d Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Sat, 31 Jan 2026 12:30:44 -0600
Subject: [PATCH 2/2] Emit error when lowering from mlir to llvmir for affinity
in task
---
.../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 5 +++++
mlir/test/Target/LLVMIR/openmp-todo.mlir | 10 ++++++++++
2 files changed, 15 insertions(+)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 672e87790456d..4343ba9c61a48 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -321,6 +321,10 @@ static LogicalResult checkImplementationStatus(Operation &op) {
<< " operation";
};
+ auto checkAffinity = [&todo](auto op, LogicalResult &result) {
+ if (!op.getAffinityVars().empty())
+ result = todo("affinity");
+ };
auto checkAllocate = [&todo](auto op, LogicalResult &result) {
if (!op.getAllocateVars().empty() || !op.getAllocatorVars().empty())
result = todo("allocate");
@@ -417,6 +421,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
checkThreadLimit(op, result);
})
.Case([&](omp::TaskOp op) {
+ checkAffinity(op, result);
checkAllocate(op, result);
checkInReduction(op, result);
})
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 36338e5cb1bed..70e4edf0705f2 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -487,3 +487,13 @@ llvm.func @wsloop_order(%lb : i32, %ub : i32, %step : i32) {
}
llvm.return
}
+
+// -----
+llvm.func @task_affinity(%x : !llvm.ptr) {
+ // expected-error at below {{not yet implemented: Unhandled clause affinity in omp.task operation}}
+ // expected-error at below {{LLVM Translation failed for operation: omp.task}}
+ omp.task affinity(%x : !llvm.ptr) {
+ omp.terminator
+ }
+ llvm.return
+}
More information about the Mlir-commits
mailing list