[Flang][mlir][OpenMP] Add affinity clause to omp.task and Flang lowering#179003
Merged
Conversation
Contributor
chichunchen
commented
Jan 31, 2026
chichunchen
commented
- Add MLIR OpenMP affinity clause
- Lower flang task affinity to mlir
- Emit TODO for iterator modifier and update negative test
- Add MLIR OpenMP affinity clause - Lower flang task affinity to mlir - Emit TODO for iterator modifier and update negative test
llvmbot
added
mlir
flang
Member
|
@llvm/pr-subscribers-mlir-llvm @llvm/pr-subscribers-mlir-openmp Author: Chi-Chun, Chen (chichunchen) Changes
Full diff: https://github.com/llvm/llvm-project/pull/179003.diff 9 Files Affected:
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
+}
|
Member
|
@llvm/pr-subscribers-flang-fir-hlfir Author: Chi-Chun, Chen (chichunchen) Changes
Full diff: https://github.com/llvm/llvm-project/pull/179003.diff 9 Files Affected:
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
+}
|
Contributor
Author
|
ping |
rishabhmadan19
pushed a commit
to rishabhmadan19/llvm-project
that referenced
this pull request
Feb 9, 2026
…ing (llvm#179003) - Add MLIR OpenMP affinity clause - Lower flang task affinity to mlir - Emit TODO for iterator modifier and update negative test
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.