[llvm-branch-commits] [flang] [mlir] [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (PR #182219)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Wed Feb 18 21:12:35 PST 2026
https://github.com/chichunchen created https://github.com/llvm/llvm-project/pull/182219
This patch translates flang ast to OpenMP dialect for affinity clause
including the iterator modifier.
>From b6401192571b996c843e731ad074534866cdfeb9 Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Wed, 18 Feb 2026 14:46:36 -0600
Subject: [PATCH 1/4] [Flang][mlir][OpenMP] Support affinity clause codegen in
Flang
This patch translate flang ast to OpenMP dialect for affinity clause
including the iterator modifier.
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 66 ++++++++++-
flang/lib/Lower/OpenMP/Utils.cpp | 103 ++++++++++++++++++
flang/lib/Lower/OpenMP/Utils.h | 9 ++
flang/test/Lower/OpenMP/task-affinity.f90 | 66 ++++++++---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 2 +-
.../mlir/Dialect/OpenMP/OpenMPOpBase.td | 8 ++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 17 +++
mlir/test/Dialect/OpenMP/ops.mlir | 60 ++++++----
8 files changed, 290 insertions(+), 41 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index e62395676a696..90956701093ea 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -762,8 +762,70 @@ bool ClauseProcessor::processAffinity(
}
const auto &objects = std::get<omp::ObjectList>(clause.t);
- if (!objects.empty())
- genObjectList(objects, converter, result.affinityVars);
+ lower::StatementContext stmtCtx;
+ auto &builder = converter.getFirOpBuilder();
+ auto &context = converter.getMLIRContext();
+ mlir::Location clauseLocation = converter.getCurrentLocation();
+
+ mlir::Type refI8Ty = fir::ReferenceType::get(builder.getIntegerType(8));
+ mlir::Type entryTy = mlir::omp::AffinityEntryType::get(
+ &context, refI8Ty, builder.getI64Type());
+
+ auto normalizeAddr = [](fir::FirOpBuilder &b, mlir::Location l,
+ mlir::Type addrTy,
+ mlir::Value v) -> mlir::Value {
+ mlir::Value addr = v;
+
+ // ref-to-box -> load box -> box_addr
+ if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(addr.getType())) {
+ if (auto innerBoxTy =
+ mlir::dyn_cast<fir::BoxType>(refTy.getEleTy())) {
+ mlir::Value boxVal = fir::LoadOp::create(b, l, innerBoxTy, addr);
+ mlir::Type boxedEleTy = innerBoxTy.getEleTy();
+ addr = fir::BoxAddrOp::create(
+ b, l, fir::ReferenceType::get(boxedEleTy), boxVal);
+ }
+ }
+
+ // box value -> box_addr
+ if (auto boxTy = mlir::dyn_cast<fir::BoxType>(addr.getType())) {
+ mlir::Type boxedEleTy = boxTy.getEleTy();
+ addr = fir::BoxAddrOp::create(
+ b, l, fir::ReferenceType::get(boxedEleTy), addr);
+ }
+
+ assert(mlir::isa<fir::ReferenceType>(addr.getType()) &&
+ "expect fir.ref after normalization");
+ return fir::ConvertOp::create(b, l, addrTy, addr);
+ };
+
+ auto makeAffinityEntry = [&](fir::FirOpBuilder &b, mlir::Location l,
+ mlir::Type entryTy, mlir::Value addr,
+ mlir::Value len) -> mlir::Value {
+ mlir::Value addrI8 = normalizeAddr(b, l, refI8Ty, addr);
+ return mlir::omp::AffinityEntryOp::create(b, l, entryTy, addrI8, len)
+ .getResult();
+ };
+
+ for (const omp::Object &object : objects) {
+ llvm::SmallVector<mlir::Value> bounds;
+ std::stringstream asFortran;
+ mlir::Value addr =
+ genAffinityAddr(converter, object, stmtCtx, clauseLocation);
+ fir::factory::AddrAndBoundsInfo info =
+ lower::gatherDataOperandAddrAndBounds<mlir::omp::MapBoundsOp,
+ mlir::omp::MapBoundsType>(
+ converter, builder, semaCtx, stmtCtx, *object.sym(),
+ object.ref(), clauseLocation, asFortran, bounds,
+ treatIndexAsSection);
+ mlir::Value len = genAffinityLen(
+ builder, clauseLocation, builder.getDataLayout(), info.addr,
+ bounds, static_cast<bool>(object.ref()));
+ // info.addr is not the base address so use the result from
+ // genAffinityAddr instead
+ result.affinityVars.push_back(
+ makeAffinityEntry(builder, clauseLocation, entryTy, addr, len));
+ }
return true;
});
diff --git a/flang/lib/Lower/OpenMP/Utils.cpp b/flang/lib/Lower/OpenMP/Utils.cpp
index e9ba5f386803a..875099b846995 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -917,6 +917,109 @@ void collectLoopRelatedInfo(
convertLoopBounds(converter, currentLocation, result, loopVarTypeSize);
}
+mlir::Value genAffinityAddr(Fortran::lower::AbstractConverter &converter,
+ const omp::Object &object,
+ Fortran::lower::StatementContext &stmtCtx,
+ mlir::Location loc) {
+ // Get address from expression if it exists: affinity(a(3)), affinity(a(1:10))
+ if (auto expr = object.ref()) {
+ fir::ExtendedValue exv =
+ converter.genExprAddr(toEvExpr(*expr), stmtCtx, &loc);
+ return fir::getBase(exv);
+ }
+
+ // Fallback to base symbol address: affinity(a)
+ const Fortran::semantics::Symbol *sym = object.sym();
+ assert(sym && "expected symbol in affinity object");
+ mlir::Value addr = converter.getSymbolAddress(*sym);
+
+ if (mlir::isa<fir::BoxType>(addr.getType())) {
+ addr = fir::BoxAddrOp::create(converter.getFirOpBuilder(), loc, addr);
+ }
+ return addr;
+}
+
+static mlir::Value buildNumElemsFromMapBound(fir::FirOpBuilder &builder,
+ mlir::Location loc,
+ mlir::omp::MapBoundsOp mb) {
+ mlir::Value lb = mb.getLowerBound();
+ mlir::Value ub = mb.getUpperBound();
+ mlir::Value stride = mb.getStride();
+
+ // ((ub - lb) / stride) + 1
+ mlir::Value diff = mlir::arith::SubIOp::create(builder, loc, ub, lb);
+ mlir::Value div = mlir::arith::DivUIOp::create(builder, loc, diff, stride);
+ mlir::Value one =
+ builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+ mlir::Value result = mlir::arith::AddIOp::create(builder, loc, div, one);
+
+ return mlir::arith::IndexCastOp::create(builder, loc, builder.getI64Type(),
+ result);
+}
+
+mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
+ const mlir::DataLayout &dl, mlir::Value addr,
+ llvm::ArrayRef<mlir::Value> bounds, bool hasRef) {
+ auto isDescriptorLike = [](mlir::Type t) -> bool {
+ t = fir::unwrapPassByRefType(t);
+ return mlir::isa<fir::BoxType, fir::ClassType>(t);
+ };
+
+ auto getElementBytesOrZero = [&](mlir::Type baseTy) -> int64_t {
+ if (isDescriptorLike(baseTy))
+ return 0;
+ mlir::Type eleTy = fir::unwrapPassByRefType(baseTy);
+ eleTy = fir::unwrapSequenceType(eleTy);
+ return static_cast<int64_t>(dl.getTypeSize(eleTy));
+ };
+
+ auto getWholeObjectBytesIfStaticOrZero = [&](mlir::Type addrTy) -> int64_t {
+ if (isDescriptorLike(addrTy))
+ return 0;
+
+ mlir::Type eleTy = fir::unwrapPassByRefType(addrTy);
+
+ // Scalar
+ if (!mlir::isa<fir::SequenceType>(eleTy))
+ return static_cast<int64_t>(dl.getTypeSize(eleTy));
+
+ // Array with static extents
+ auto seqTy = mlir::cast<fir::SequenceType>(eleTy);
+ int64_t elems = 1;
+ for (int64_t d : seqTy.getShape()) {
+ if (d < 0)
+ return 0; // dynamic extent => unknown here
+ elems *= d;
+ }
+
+ int64_t elemBytes = static_cast<int64_t>(dl.getTypeSize(seqTy.getEleTy()));
+ return elems * elemBytes;
+ };
+
+ // Return the length of the first dimension if bounds are available
+ if (!bounds.empty()) {
+ auto mb = bounds.front().getDefiningOp<mlir::omp::MapBoundsOp>();
+ mlir::Value numElems = buildNumElemsFromMapBound(builder, loc, mb);
+ int64_t elemBytes = getElementBytesOrZero(addr.getType());
+ if (elemBytes == 0)
+ return builder.createIntegerConstant(loc, builder.getI64Type(), 0);
+
+ return mlir::arith::MulIOp::create(
+ builder, loc, numElems,
+ builder.createIntegerConstant(loc, builder.getI64Type(), elemBytes));
+ }
+
+ // explicit ref => element size (a(3), a(i))
+ if (hasRef) {
+ int64_t elemBytes = getElementBytesOrZero(addr.getType());
+ return builder.createIntegerConstant(loc, builder.getI64Type(), elemBytes);
+ }
+
+ // whole object => whole size if static, else 0
+ int64_t wholeBytes = getWholeObjectBytesIfStaticOrZero(addr.getType());
+ return builder.createIntegerConstant(loc, builder.getI64Type(), wholeBytes);
+}
+
} // namespace omp
} // namespace lower
} // namespace Fortran
diff --git a/flang/lib/Lower/OpenMP/Utils.h b/flang/lib/Lower/OpenMP/Utils.h
index f707557197847..d6d3887e41fff 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -189,6 +189,15 @@ void collectTileSizesFromOpenMPConstruct(
llvm::SmallVectorImpl<int64_t> &tileSizes,
Fortran::semantics::SemanticsContext &semaCtx);
+mlir::Value genAffinityAddr(Fortran::lower::AbstractConverter &converter,
+ const omp::Object &object,
+ Fortran::lower::StatementContext &stmtCtx,
+ mlir::Location loc);
+
+mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
+ const mlir::DataLayout &dl, mlir::Value addr,
+ llvm::ArrayRef<mlir::Value> bounds, bool hasRef);
+
} // namespace omp
} // namespace lower
} // namespace Fortran
diff --git a/flang/test/Lower/OpenMP/task-affinity.f90 b/flang/test/Lower/OpenMP/task-affinity.f90
index 66254e48e9b8e..d75f1712df503 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -16,17 +16,29 @@ subroutine omp_task_affinity_elem()
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: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFomp_task_affinity_elemEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xi32>>, !fir.ref<!fir.array<100xi32>>)
! CHECK: omp.parallel {
! CHECK: omp.single {
-! CHECK: omp.task affinity(%[[A1]]#0 : !fir.ref<!fir.array<100xi32>>) {
+! CHECK: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[ELEM:.*]] = hlfir.designate %[[A]]#0 (%[[C1]]) : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+! CHECK: %[[C1_0:.*]] = arith.constant 1 : index
+! CHECK: %[[C0:.*]] = arith.constant 0 : index
+! CHECK: %[[SUB:.*]] = arith.subi %[[C0]], %[[C0]] : index
+! CHECK: %[[DIV:.*]] = arith.divui %[[SUB]], %[[C1_0]] : index
+! CHECK: %[[C1_1:.*]] = arith.constant 1 : index
+! CHECK: %[[ADD:.*]] = arith.addi %[[DIV]], %[[C1_1]] : index
+! CHECK: %[[CAST:.*]] = arith.index_cast %[[ADD]] : index to i64
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[LEN:.*]] = arith.muli %[[CAST]], %[[C4]] : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[ELEM]] : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
! 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
@@ -45,18 +57,35 @@ subroutine omp_task_affinity_array_section()
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: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFomp_task_affinity_array_sectionEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xi32>>, !fir.ref<!fir.array<100xi32>>)
+! CHECK: %[[I:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFomp_task_affinity_array_sectionEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
! CHECK: omp.parallel {
! CHECK: omp.single {
-! CHECK: omp.task affinity(%[[A2]]#0 : !fir.ref<!fir.array<100xi32>>) private({{.*}} %[[I2]]#0 -> %{{.*}} : !fir.ref<i32>) {
+! CHECK: %[[C2:.*]] = arith.constant 2 : index
+! CHECK: %[[C50:.*]] = arith.constant 50 : index
+! CHECK: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[C49:.*]] = arith.constant 49 : index
+! CHECK: %[[SHAPE:.*]] = fir.shape %[[C49]] : (index) -> !fir.shape<1>
+! CHECK: %[[SLICE:.*]] = hlfir.designate %[[A]]#0 (%[[C2]]:%[[C50]]:%[[C1]]) shape %[[SHAPE]] : (!fir.ref<!fir.array<100xi32>>, index, index, index, !fir.shape<1>) -> !fir.ref<!fir.array<49xi32>>
+! CHECK: %[[C1_0:.*]] = arith.constant 1 : index
+! CHECK: %[[C1_1:.*]] = arith.constant 1 : index
+! CHECK: %[[C49_2:.*]] = arith.constant 49 : index
+! CHECK: %[[SUB:.*]] = arith.subi %[[C49_2]], %[[C1_1]] : index
+! CHECK: %[[DIV:.*]] = arith.divui %[[SUB]], %[[C1_0]] : index
+! CHECK: %[[C1_3:.*]] = arith.constant 1 : index
+! CHECK: %[[ADD:.*]] = arith.addi %[[DIV]], %[[C1_3]] : index
+! CHECK: %[[CAST:.*]] = arith.index_cast %[[ADD]] : index to i64
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[LEN:.*]] = arith.muli %[[CAST]], %[[C4]] : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[SLICE]] : (!fir.ref<!fir.array<49xi32>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) private(@_QFomp_task_affinity_array_sectionEi_private_i32 %[[I]]#0 -> %[[P:.*]] : !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
@@ -72,17 +101,23 @@ subroutine omp_task_affinity_scalar()
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: %[[S:.*]] = fir.alloca i32 {bindc_name = "s", uniq_name = "_QFomp_task_affinity_scalarEs"}
+! CHECK: %[[SDECL:.*]]:2 = hlfir.declare %[[S]] {uniq_name = "_QFomp_task_affinity_scalarEs"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: hlfir.assign %{{.*}} to %[[SDECL]]#0 : i32, !fir.ref<i32>
! CHECK: omp.parallel {
! CHECK: omp.single {
-! CHECK: omp.task affinity(%[[S3]]#0 : !fir.ref<i32>) {
+! CHECK: %[[LEN:.*]] = arith.constant 4 : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[SDECL]]#0 : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
! CHECK: omp.terminator
! CHECK: }
+! CHECK: omp.terminator
+! CHECK: }
! CHECK: omp.terminator
! CHECK: }
! CHECK: return
-! multiple locators
subroutine omp_task_affinity_multi()
implicit none
integer, parameter :: n = 100
@@ -99,13 +134,16 @@ subroutine omp_task_affinity_multi()
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: %[[AADDR:.*]] = fir.convert %{{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[AENT:.*]] = omp.affinity_entry %[[AADDR]], %{{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: %[[BADDR:.*]] = fir.convert %{{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[BENT:.*]] = omp.affinity_entry %[[BADDR]], %{{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.task affinity(%[[AENT]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>, %[[BENT]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
! CHECK: omp.terminator
! CHECK: }
+! 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 ba52e52ebf58d..bf59f9d108501 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -32,7 +32,7 @@ class OpenMP_AffinityClauseSkip<
bit description = false, bit extraClassDeclaration = false>
: OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
- let arguments = (ins Variadic<OpenMP_PointerLikeType>:$affinity_vars);
+ let arguments = (ins Variadic<OpenMP_AffinityEntryType>:$affinity_vars);
let optAssemblyFormat = [{
`affinity` `(` custom<AffinityClause>($affinity_vars, type($affinity_vars)) `)`
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
index 4dd8e91585a66..c1017826ab0c9 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
@@ -38,6 +38,14 @@ def OpenMP_MapBoundsType : OpenMP_Type<"MapBounds", "map_bounds_ty"> {
let summary = "Type for representing omp map clause bounds information";
}
+def OpenMP_AffinityEntryType
+ : OpenMP_Type<"AffinityEntry", "affinity_entry_ty"> {
+ let summary = "Type for representing omp affinity clause locator information";
+
+ let parameters = (ins "Type":$addrType, "Type":$lenType);
+ let assemblyFormat = "`<` $addrType `,` $lenType `>`";
+}
+
def OpenMP_IteratedType : OpenMP_Type<"Iterated", "iterated"> {
let summary = "OpenMP iterator-produced list handle";
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 9705185656907..87f61a27c4447 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -2300,6 +2300,23 @@ def DeclareSimdOp
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// Affinity Entry Op
+//===----------------------------------------------------------------------===//
+
+def AffinityEntryOp : OpenMP_Op<"affinity_entry", [Pure]> {
+ let summary = "OpenMP affinity clause entry value";
+
+ let arguments = (ins OpenMP_PointerLikeType:$addr, IntLikeType:$len);
+
+ let results = (outs OpenMP_AffinityEntryType:$entry);
+
+ let assemblyFormat = [{
+ $addr `,` $len `:` `(` type($addr) `,` type($len) `)` `->`
+ qualified(type($entry)) attr-dict
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Iterators Op
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 6186ff35c8df7..fa4d8ec490f32 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3522,30 +3522,42 @@ func.func @omp_declare_simd_all_clauses(%a: f64, %b: f64,
}
// 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>) {
+func.func @task_affinity_single(%ptr: !llvm.ptr) {
+ // CHECK: %[[LEN:.*]] = llvm.mlir.constant(400 : i64) : i64
+ // CHECK: %[[AE:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: omp.task affinity(%[[AE]] : !omp.affinity_entry_ty<!llvm.ptr, i64>) {
+ // CHECK: omp.terminator
+ // CHECK: }
+ // CHECK: return
+ %len = llvm.mlir.constant(400 : i64) : i64
+ %ae = omp.affinity_entry %ptr, %len
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ omp.task affinity(%ae : !omp.affinity_entry_ty<!llvm.ptr, i64>) {
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @task_affinity_multi(
+func.func @task_affinity_multi(%ptr1: !llvm.ptr, %ptr2: !llvm.ptr) {
+ // CHECK: %[[LEN1:.*]] = llvm.mlir.constant(400 : i64) : i64
+ // CHECK: %[[AE1:.*]] = omp.affinity_entry %{{.*}}, %[[LEN1]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: %[[LEN2:.*]] = llvm.mlir.constant(800 : i64) : i64
+ // CHECK: %[[AE2:.*]] = omp.affinity_entry %{{.*}}, %[[LEN2]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: omp.task affinity(%[[AE1]] : !omp.affinity_entry_ty<!llvm.ptr, i64>, %[[AE2]] : !omp.affinity_entry_ty<!llvm.ptr, i64>) {
+ // CHECK: omp.terminator
+ // CHECK: }
+ // CHECK: return
+ %len1 = llvm.mlir.constant(400 : i64) : i64
+ %ae1 = omp.affinity_entry %ptr1, %len1
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+
+ %len2 = llvm.mlir.constant(800 : i64) : i64
+ %ae2 = omp.affinity_entry %ptr2, %len2
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+
+ omp.task affinity(%ae1 : !omp.affinity_entry_ty<!llvm.ptr, i64>,
+ %ae2 : !omp.affinity_entry_ty<!llvm.ptr, i64>) {
omp.terminator
}
return
>From 9605493ba9c5eda9a23fe400c2d9009aa934bc8f Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Wed, 18 Feb 2026 13:25:17 -0600
Subject: [PATCH 2/4] Emit omp.iterator in affinity clause from Flang
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 163 +++++++++++++++---
flang/lib/Lower/OpenMP/Utils.cpp | 83 +++++++++
flang/lib/Lower/OpenMP/Utils.h | 16 ++
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 6 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 88 ++++++++--
5 files changed, 316 insertions(+), 40 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 90956701093ea..f2246377b5ffd 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -202,6 +202,92 @@ getIfClauseOperand(lower::AbstractConverter &converter,
ifVal);
}
+template <typename IteratorSpecT>
+static IteratorRange lowerIteratorRange(
+ Fortran::lower::AbstractConverter &converter, const IteratorSpecT &itSpec,
+ Fortran::lower::StatementContext &stmtCtx, mlir::Location loc) {
+ auto &builder = converter.getFirOpBuilder();
+
+ const auto &ivObj = std::get<1>(itSpec.t);
+ const auto &range = std::get<2>(itSpec.t);
+
+ IteratorRange r;
+ r.ivSym = ivObj.sym();
+ assert(r.ivSym && "expected iterator induction symbol");
+
+ const auto &lbExpr = std::get<0>(range.t);
+ const auto &ubExpr = std::get<1>(range.t);
+ const auto &stExprOpt = std::get<2>(range.t);
+
+ mlir::Value lbVal =
+ fir::getBase(converter.genExprValue(toEvExpr(lbExpr), stmtCtx));
+ mlir::Value ubVal =
+ fir::getBase(converter.genExprValue(toEvExpr(ubExpr), stmtCtx));
+
+ auto toIndex = [](fir::FirOpBuilder &builder, mlir::Location loc,
+ mlir::Value v) -> mlir::Value {
+ if (v.getType().isIndex())
+ return v;
+ return mlir::arith::IndexCastOp::create(builder, loc,
+ builder.getIndexType(), v);
+ };
+
+ r.lb = toIndex(builder, loc, lbVal);
+ r.ub = toIndex(builder, loc, ubVal);
+
+ if (stExprOpt) {
+ mlir::Value stVal =
+ fir::getBase(converter.genExprValue(toEvExpr(*stExprOpt), stmtCtx));
+ r.step = toIndex(builder, loc, stVal);
+ } else {
+ r.step = mlir::arith::ConstantIndexOp::create(builder, loc, 1);
+ }
+
+ return r;
+}
+
+template <typename BuildBodyFn>
+static mlir::Value buildIteratorOp(Fortran::lower::AbstractConverter &converter,
+ mlir::Location loc, mlir::Type iterTy,
+ llvm::ArrayRef<IteratorRange> ranges,
+ BuildBodyFn &&buildBody) {
+
+ auto &builder = converter.getFirOpBuilder();
+
+ llvm::SmallVector<mlir::Value> lbs, ubs, steps;
+ lbs.reserve(ranges.size());
+ ubs.reserve(ranges.size());
+ steps.reserve(ranges.size());
+ for (auto &r : ranges) {
+ lbs.push_back(r.lb);
+ ubs.push_back(r.ub);
+ steps.push_back(r.step);
+ }
+
+ auto itOp = mlir::omp::IteratorsOp::create(
+ builder, loc, iterTy, mlir::ValueRange{lbs}, mlir::ValueRange{ubs},
+ mlir::ValueRange{steps});
+
+ mlir::OpBuilder::InsertionGuard guard(builder);
+
+ mlir::Region ® = itOp.getRegion();
+ mlir::Block *body = builder.createBlock(®);
+
+ llvm::SmallVector<mlir::Value> ivs;
+ ivs.reserve(ranges.size());
+ for (size_t i = 0; i < ranges.size(); ++i)
+ ivs.push_back(body->addArgument(builder.getIndexType(), loc));
+
+ Fortran::lower::SymMap &symMap = converter.getSymbolMap();
+ Fortran::lower::SymMapScope scope(symMap);
+ for (size_t i = 0; i < ranges.size(); ++i)
+ symMap.addSymbol(*ranges[i].ivSym, ivs[i], /*force=*/true);
+
+ mlir::omp::YieldOp::create(builder, loc, buildBody(builder, loc, ivs));
+
+ return itOp.getResult();
+}
+
//===----------------------------------------------------------------------===//
// ClauseProcessor unique clauses
//===----------------------------------------------------------------------===//
@@ -756,11 +842,6 @@ 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);
lower::StatementContext stmtCtx;
auto &builder = converter.getFirOpBuilder();
@@ -770,6 +851,7 @@ bool ClauseProcessor::processAffinity(
mlir::Type refI8Ty = fir::ReferenceType::get(builder.getIntegerType(8));
mlir::Type entryTy = mlir::omp::AffinityEntryType::get(
&context, refI8Ty, builder.getI64Type());
+ mlir::Type iterTy = mlir::omp::IteratedType::get(&context, entryTy);
auto normalizeAddr = [](fir::FirOpBuilder &b, mlir::Location l,
mlir::Type addrTy,
@@ -807,24 +889,65 @@ bool ClauseProcessor::processAffinity(
.getResult();
};
+ llvm::SmallVector<IteratorRange> iteratorRanges;
+ llvm::SmallPtrSet<const Fortran::semantics::Symbol *, 4> ivSyms;
+
+ // If iterator modifier exists, collect ranges and IV symbols.
+ auto &iteratorModifier =
+ std::get<std::optional<omp::clause::Iterator>>(clause.t);
+ if (iteratorModifier.has_value()) {
+ const auto &iteratorModifierSpecs = *iteratorModifier;
+ iteratorRanges.reserve(iteratorModifierSpecs.size());
+ for (const auto &itSpec : iteratorModifierSpecs)
+ iteratorRanges.push_back(
+ lowerIteratorRange(converter, itSpec, stmtCtx, clauseLocation));
+
+ for (const IteratorRange &r : iteratorRanges)
+ ivSyms.insert(&r.ivSym->GetUltimate());
+ }
+
for (const omp::Object &object : objects) {
llvm::SmallVector<mlir::Value> bounds;
std::stringstream asFortran;
- mlir::Value addr =
- genAffinityAddr(converter, object, stmtCtx, clauseLocation);
- fir::factory::AddrAndBoundsInfo info =
- lower::gatherDataOperandAddrAndBounds<mlir::omp::MapBoundsOp,
- mlir::omp::MapBoundsType>(
- converter, builder, semaCtx, stmtCtx, *object.sym(),
- object.ref(), clauseLocation, asFortran, bounds,
- treatIndexAsSection);
- mlir::Value len = genAffinityLen(
- builder, clauseLocation, builder.getDataLayout(), info.addr,
- bounds, static_cast<bool>(object.ref()));
- // info.addr is not the base address so use the result from
- // genAffinityAddr instead
- result.affinityVars.push_back(
- makeAffinityEntry(builder, clauseLocation, entryTy, addr, len));
+ if (iteratorModifier.has_value() && hasIVReference(object, ivSyms)) {
+ mlir::Value iterHandle = buildIteratorOp(
+ converter, clauseLocation, iterTy, iteratorRanges,
+ [&](fir::FirOpBuilder &builder, mlir::Location loc,
+ llvm::ArrayRef<mlir::Value> ivs) -> mlir::Value {
+ const Fortran::semantics::Symbol *sym = object.sym();
+ assert(sym && "expected symbol for iterator object");
+ fir::factory::AddrAndBoundsInfo info =
+ Fortran::lower::getDataOperandBaseAddr(
+ converter, builder, *sym, loc,
+ /*unwrapFirBox=*/false);
+ // TODO check correctness of genIteratorCoordinate
+ mlir::Value addr =
+ genIteratorCoordinate(converter, info.addr, ivs, loc);
+ // Length of iterator-based affinity entry set as element size
+ mlir::Value len = genAffinityLen(
+ builder, clauseLocation, builder.getDataLayout(),
+ info.addr, bounds, static_cast<bool>(object.ref()));
+ return makeAffinityEntry(builder, loc, entryTy, addr, len);
+ });
+ iterHandle.dump();
+ result.iterated.push_back(iterHandle);
+ } else {
+ mlir::Value addr =
+ genAffinityAddr(converter, object, stmtCtx, clauseLocation);
+ fir::factory::AddrAndBoundsInfo info =
+ lower::gatherDataOperandAddrAndBounds<mlir::omp::MapBoundsOp,
+ mlir::omp::MapBoundsType>(
+ converter, builder, semaCtx, stmtCtx, *object.sym(),
+ object.ref(), clauseLocation, asFortran, bounds,
+ treatIndexAsSection);
+ mlir::Value len = genAffinityLen(
+ builder, clauseLocation, builder.getDataLayout(), info.addr,
+ bounds, static_cast<bool>(object.ref()));
+ // info.addr is not the base address so use the result from
+ // genAffinityAddr instead
+ result.affinityVars.push_back(
+ makeAffinityEntry(builder, clauseLocation, entryTy, addr, len));
+ }
}
return true;
diff --git a/flang/lib/Lower/OpenMP/Utils.cpp b/flang/lib/Lower/OpenMP/Utils.cpp
index 875099b846995..83d1a977cb3e3 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -1020,6 +1020,89 @@ mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
return builder.createIntegerConstant(loc, builder.getI64Type(), wholeBytes);
}
+bool hasIVReference(
+ const omp::Object &object,
+ const llvm::SmallPtrSetImpl<const Fortran::semantics::Symbol *> &ivSyms) {
+ auto ref = object.ref();
+ if (!ref)
+ return false;
+
+ Fortran::lower::SomeExpr expr = toEvExpr(*ref);
+
+ for (Fortran::evaluate::SymbolRef s : CollectSymbols(expr)) {
+ const Fortran::semantics::Symbol &ult = s->GetUltimate();
+ if (ivSyms.contains(&ult))
+ return true;
+ }
+ return false;
+}
+
+mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
+ mlir::Value base,
+ llvm::ArrayRef<mlir::Value> ivs,
+ mlir::Location loc) {
+ auto &builder = converter.getFirOpBuilder();
+ mlir::Type baseTy = base.getType();
+
+ // If base is a reference-to-box, load it to get the box value.
+ if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(baseTy)) {
+ if (auto innerBoxTy = mlir::dyn_cast<fir::BoxType>(refTy.getEleTy())) {
+ base = fir::LoadOp::create(builder, loc, innerBoxTy, base);
+ baseTy = base.getType();
+ }
+ }
+
+ // descriptor-backed arrays (assumed-shape dummies etc.)
+ if (auto boxTy = mlir::dyn_cast<fir::BoxType>(baseTy)) {
+ // Build !fir.shape<rank> from descriptor extents.
+ const unsigned rank = ivs.size();
+ llvm::SmallVector<mlir::Value> extents;
+ extents.reserve(rank);
+
+ for (unsigned d = 0; d < rank; ++d) {
+ mlir::Value dim = builder.createIntegerConstant(loc, builder.getI32Type(),
+ static_cast<int64_t>(d));
+ auto dims = fir::BoxDimsOp::create(builder, loc,
+ /*lbType=*/builder.getIndexType(),
+ /*extentType=*/builder.getIndexType(),
+ /*strideType=*/builder.getIndexType(),
+ base, dim);
+ extents.push_back(dims.getExtent());
+ }
+
+ mlir::Value shape = fir::ShapeOp::create(builder, loc, extents);
+
+ // Result element reference type.
+ mlir::Type boxedEleTy = boxTy.getEleTy(); // e.g. !fir.array<?x?xi32>
+ if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(boxedEleTy))
+ boxedEleTy = seqTy.getEleTy();
+ mlir::Type eleRefTy = fir::ReferenceType::get(boxedEleTy);
+
+ return fir::ArrayCoorOp::create(builder, loc, eleRefTy,
+ /*memref=*/base,
+ /*shape=*/shape,
+ /*slice=*/mlir::Value{},
+ /*indices=*/mlir::ValueRange{ivs},
+ /*typeparams=*/mlir::ValueRange{});
+ }
+
+ // explicit-shape arrays lowered as !fir.ref<!fir.array<...>>
+ // base must be a reference to a SequenceType.
+ auto baseRefTy = mlir::cast<fir::ReferenceType>(baseTy);
+ auto seqTy = mlir::cast<fir::SequenceType>(baseRefTy.getEleTy());
+ mlir::Type eleRefTy = fir::ReferenceType::get(seqTy.getEleTy());
+
+ // coordinate_of expects i32 subscripts.
+ llvm::SmallVector<mlir::Value> subsI32;
+ subsI32.reserve(ivs.size());
+ for (mlir::Value iv : ivs) {
+ subsI32.push_back(mlir::arith::IndexCastOp::create(
+ builder, loc, builder.getI32Type(), iv));
+ }
+
+ return fir::CoordinateOp::create(builder, loc, eleRefTy, base, subsI32);
+}
+
} // namespace omp
} // namespace lower
} // namespace Fortran
diff --git a/flang/lib/Lower/OpenMP/Utils.h b/flang/lib/Lower/OpenMP/Utils.h
index d6d3887e41fff..10552a68908d8 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -198,6 +198,22 @@ mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
const mlir::DataLayout &dl, mlir::Value addr,
llvm::ArrayRef<mlir::Value> bounds, bool hasRef);
+struct IteratorRange {
+ mlir::Value lb;
+ mlir::Value ub;
+ mlir::Value step;
+ Fortran::semantics::Symbol *ivSym = nullptr;
+};
+
+bool hasIVReference(
+ const omp::Object &object,
+ const llvm::SmallPtrSetImpl<const Fortran::semantics::Symbol *> &ivSyms);
+
+mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
+ mlir::Value base,
+ llvm::ArrayRef<mlir::Value> ivs,
+ mlir::Location loc);
+
} // namespace omp
} // namespace lower
} // namespace Fortran
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index bf59f9d108501..52defbbb04c27 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -32,10 +32,12 @@ class OpenMP_AffinityClauseSkip<
bit description = false, bit extraClassDeclaration = false>
: OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
- let arguments = (ins Variadic<OpenMP_AffinityEntryType>:$affinity_vars);
+ let arguments = (ins Variadic<OpenMP_IteratedType>:$iterated,
+ Variadic<OpenMP_AffinityEntryType>:$affinity_vars);
let optAssemblyFormat = [{
- `affinity` `(` custom<AffinityClause>($affinity_vars, type($affinity_vars)) `)`
+ `affinity` `(` custom<AffinityClause>($iterated, $affinity_vars,
+ type($iterated), type($affinity_vars)) `)`
}];
let description = [{
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index b55e278e08255..b51fd0002862f 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1431,6 +1431,56 @@ static void printUseDeviceAddrUseDevicePtrRegion(OpAsmPrinter &p, Operation *op,
printBlockArgRegion(p, op, region, args);
}
+template <typename ParsePrefixFn>
+static ParseResult parseSplitIteratedList(
+ OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &iteratedVars,
+ SmallVectorImpl<Type> &iteratedTypes,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &plainVars,
+ SmallVectorImpl<Type> &plainTypes, ParsePrefixFn &&parsePrefix) {
+
+ return parser.parseCommaSeparatedList([&]() -> ParseResult {
+ if (failed(parsePrefix()))
+ return failure();
+
+ OpAsmParser::UnresolvedOperand v;
+ Type ty;
+ if (parser.parseOperand(v) || parser.parseColonType(ty))
+ return failure();
+
+ if (llvm::isa<mlir::omp::IteratedType>(ty)) {
+ iteratedVars.push_back(v);
+ iteratedTypes.push_back(ty);
+ } else {
+ plainVars.push_back(v);
+ plainTypes.push_back(ty);
+ }
+ return success();
+ });
+}
+
+template <typename PrintPrefixFn>
+static void printSplitIteratedList(OpAsmPrinter &p, ValueRange iteratedVars,
+ TypeRange iteratedTypes,
+ ValueRange plainVars, TypeRange plainTypes,
+ PrintPrefixFn &&printPrefixForPlain,
+ PrintPrefixFn &&printPrefixForIterated) {
+
+ bool first = true;
+ auto emit = [&](Value v, Type t, auto &&printPrefix) {
+ if (!first)
+ p << ", ";
+ printPrefix(v, t);
+ p << v << " : " << t;
+ first = false;
+ };
+
+ for (unsigned i = 0; i < iteratedVars.size(); ++i)
+ emit(iteratedVars[i], iteratedTypes[i], printPrefixForIterated);
+ for (unsigned i = 0; i < plainVars.size(); ++i)
+ emit(plainVars[i], plainTypes[i], printPrefixForPlain);
+}
+
/// Verifies Reduction Clause
static LogicalResult
verifyReductionVarList(Operation *op, std::optional<ArrayAttr> reductionSyms,
@@ -3137,10 +3187,10 @@ LogicalResult DeclareReductionOp::verifyRegions() {
void TaskOp::build(OpBuilder &builder, OperationState &state,
const TaskOperands &clauses) {
MLIRContext *ctx = builder.getContext();
- TaskOp::build(builder, state, clauses.affinityVars, clauses.allocateVars,
- clauses.allocatorVars, makeArrayAttr(ctx, clauses.dependKinds),
- clauses.dependVars, clauses.final, clauses.ifExpr,
- clauses.inReductionVars,
+ TaskOp::build(builder, state, clauses.iterated, 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,
@@ -4604,24 +4654,26 @@ static void printUniformClause(OpAsmPrinter &p, Operation *op,
static ParseResult parseAffinityClause(
OpAsmParser &parser,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &iterated,
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();
- });
+ SmallVectorImpl<Type> &iteratedTypes,
+ SmallVectorImpl<Type> &affinityVarTypes) {
+ if (failed(parseSplitIteratedList(
+ parser, iterated, iteratedTypes, affinityVars, affinityVarTypes,
+ /*parsePrefix=*/[&]() -> ParseResult { return success(); })))
+ 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];
- }
+ ValueRange iterated, ValueRange affinityVars,
+ TypeRange iteratedTypes,
+ TypeRange affinityVarTypes) {
+ auto nop = [&](Value, Type) {};
+ printSplitIteratedList(p, iterated, iteratedTypes, affinityVars,
+ affinityVarTypes,
+ /*plain prefix*/ nop,
+ /*iterated prefix*/ nop);
}
//===----------------------------------------------------------------------===//
>From 31c91766932b3c8eca4486625865bf90a72bc2af Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Wed, 18 Feb 2026 14:15:02 -0600
Subject: [PATCH 3/4] Add iterator test and remove redundant check lines
---
.../Lower/OpenMP/Todo/affinity-clause.f90 | 10 ---
flang/test/Lower/OpenMP/task-affinity.f90 | 88 +++++++++++--------
mlir/test/Dialect/OpenMP/ops.mlir | 53 +++++++++++
3 files changed, 106 insertions(+), 45 deletions(-)
delete mode 100644 flang/test/Lower/OpenMP/Todo/affinity-clause.f90
diff --git a/flang/test/Lower/OpenMP/Todo/affinity-clause.f90 b/flang/test/Lower/OpenMP/Todo/affinity-clause.f90
deleted file mode 100644
index 6be477229286a..0000000000000
--- a/flang/test/Lower/OpenMP/Todo/affinity-clause.f90
+++ /dev/null
@@ -1,10 +0,0 @@
-!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: Support for iterator modifiers is not implemented yet
-subroutine f00(x)
- integer :: x(10)
-!$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
index d75f1712df503..3e3f8acad7c62 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -7,18 +7,15 @@ subroutine omp_task_affinity_elem()
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: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFomp_task_affinity_elemEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xi32>>, !fir.ref<!fir.array<100xi32>>)
! CHECK: omp.parallel {
-! CHECK: omp.single {
! CHECK: %[[C1:.*]] = arith.constant 1 : index
! CHECK: %[[ELEM:.*]] = hlfir.designate %[[A]]#0 (%[[C1]]) : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
! CHECK: %[[C1_0:.*]] = arith.constant 1 : index
@@ -33,11 +30,6 @@ end subroutine omp_task_affinity_elem
! CHECK: %[[ADDRI8:.*]] = fir.convert %[[ELEM]] : (!fir.ref<i32>) -> !fir.ref<i8>
! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
subroutine omp_task_affinity_array_section()
implicit none
@@ -46,13 +38,11 @@ subroutine omp_task_affinity_array_section()
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
@@ -60,7 +50,6 @@ end subroutine omp_task_affinity_array_section
! CHECK: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFomp_task_affinity_array_sectionEa"} : (!fir.ref<!fir.array<100xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<100xi32>>, !fir.ref<!fir.array<100xi32>>)
! CHECK: %[[I:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFomp_task_affinity_array_sectionEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
! CHECK: omp.parallel {
-! CHECK: omp.single {
! CHECK: %[[C2:.*]] = arith.constant 2 : index
! CHECK: %[[C50:.*]] = arith.constant 50 : index
! CHECK: %[[C1:.*]] = arith.constant 1 : index
@@ -80,11 +69,6 @@ end subroutine omp_task_affinity_array_section
! CHECK: %[[ADDRI8:.*]] = fir.convert %[[SLICE]] : (!fir.ref<!fir.array<49xi32>>) -> !fir.ref<i8>
! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) private(@_QFomp_task_affinity_array_sectionEi_private_i32 %[[I]]#0 -> %[[P:.*]] : !fir.ref<i32>) {
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
subroutine omp_task_affinity_scalar()
implicit none
@@ -92,11 +76,9 @@ subroutine omp_task_affinity_scalar()
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
@@ -105,18 +87,10 @@ end subroutine omp_task_affinity_scalar
! CHECK: %[[SDECL:.*]]:2 = hlfir.declare %[[S]] {uniq_name = "_QFomp_task_affinity_scalarEs"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
! CHECK: hlfir.assign %{{.*}} to %[[SDECL]]#0 : i32, !fir.ref<i32>
! CHECK: omp.parallel {
-! CHECK: omp.single {
! CHECK: %[[LEN:.*]] = arith.constant 4 : i64
! CHECK: %[[ADDRI8:.*]] = fir.convert %[[SDECL]]#0 : (!fir.ref<i32>) -> !fir.ref<i8>
! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
! CHECK: omp.task affinity(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
subroutine omp_task_affinity_multi()
implicit none
@@ -124,26 +98,70 @@ subroutine omp_task_affinity_multi()
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: omp.parallel {
-! CHECK: omp.single {
! CHECK: %[[AADDR:.*]] = fir.convert %{{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
! CHECK: %[[AENT:.*]] = omp.affinity_entry %[[AADDR]], %{{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
! CHECK: %[[BADDR:.*]] = fir.convert %{{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
! CHECK: %[[BENT:.*]] = omp.affinity_entry %[[BADDR]], %{{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
! CHECK: omp.task affinity(%[[AENT]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>, %[[BENT]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>) {
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
+
+subroutine task_affinity_iterator_simple()
+ integer, parameter :: n = 16
+ integer :: a(n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n) : a(i))
+ a(i) = i
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_simple()
+! CHECK: %[[ITERATED:.*]] = omp.iterators(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: {{.*}} = arith.index_cast %[[IV]] : index to i32
+! CHECK: {{.*}} = fir.coordinate_of {{.*}} : (!fir.ref<!fir.array<16xi32>>, i32) -> !fir.ref<i32>
+! CHECK: {{.*}} = fir.convert {{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry {{.*}}, {{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.yield(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>)
+! CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>
+! CHECK: omp.task affinity(%{{.*}} : !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>) {
+
+subroutine task_affinity_iterator_multi_dimension()
+ integer, parameter :: n = 4, m = 6
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 3:n, j = 1:m) : a(i, j)) shared(a)
+ do i = 1, n
+ do j = 1, m
+ a(i, j) = 100*i + j
+ end do
+ end do
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_multi_dimension()
+! CHECK: %[[ITER:.*]] = omp.iterators(%[[IV0:.*]]: index, %[[IV1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: {{.*}} = arith.index_cast %[[IV0]] : index to i32
+! CHECK: {{.*}} = arith.index_cast %[[IV1]] : index to i32
+! CHECK: {{.*}} = fir.coordinate_of {{.*}}, {{.*}}, {{.*}} : (!fir.ref<!fir.array<4x6xi32>>, i32, i32) -> !fir.ref<i32>
+! CHECK: {{.*}} = fir.convert {{.*}} : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry {{.*}}, {{.*}} : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.yield(%[[ENTRY]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>)
+! CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>
+! CHECK: omp.task affinity(%[[ITER]] : !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>)
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index fa4d8ec490f32..efb3be121399e 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3591,3 +3591,56 @@ func.func @omp_iterators_2d(%lb0 : index, %ub0 : index, %st0 : index,
return
}
+
+// CHECK-LABEL: func.func @omp_task_affinity_iterator_1d(
+func.func @omp_task_affinity_iterator_1d(%lb : index, %ub : index, %step : index,
+ %addr : !llvm.ptr, %len : i64) -> () {
+ // CHECK: %[[IT:.*]] = omp.iterators(%[[IV:.*]]: index) = (%[[LB:.*]] to %[[UB:.*]] step %[[ST:.*]]) {
+ // CHECK: %[[E:.*]] = omp.affinity_entry %[[ADDR:.*]], %[[LEN:.*]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: omp.yield(%[[E]] : !omp.affinity_entry_ty<!llvm.ptr, i64>)
+ // CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>
+ // CHECK: omp.task affinity(%[[IT]] : !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>) {
+ // CHECK: }
+ %it = omp.iterators(%iv: index) = (%lb to %ub step %step) {
+ %e = omp.affinity_entry %addr, %len
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ omp.yield(%e : !omp.affinity_entry_ty<!llvm.ptr, i64>)
+ } -> !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>
+
+ omp.task affinity(%it : !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>) {
+ omp.terminator
+ }
+
+ return
+}
+
+func.func @omp_task_affinity_iterator_2d(%lb0 : index, %ub0 : index, %st0 : index,
+ %lb1 : index, %ub1 : index, %st1 : index,
+ %addr0 : !llvm.ptr, %addr1 : !llvm.ptr,
+ %len0 : i64, %len1 : i64) -> () {
+ // CHECK: %[[IT:.*]] = omp.iterators(%[[I:.*]]: index, %[[J:.*]]: index) = (%[[LB0:.*]] to %[[UB0:.*]] step %[[ST0:.*]], %[[LB1:.*]] to %[[UB1:.*]] step %[[ST1:.*]]) {
+ // CHECK: %[[E0:.*]] = omp.affinity_entry %[[A0:.*]], %[[L0:.*]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: %[[E1:.*]] = omp.affinity_entry %[[A1:.*]], %[[L1:.*]] : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ // CHECK: omp.yield(%[[E1]] : !omp.affinity_entry_ty<!llvm.ptr, i64>)
+ // CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>
+ // CHECK: omp.task affinity(%[[IT]] : !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>) {
+ // CHECK: }
+ %it = omp.iterators(%i: index, %j: index) = (%lb0 to %ub0 step %st0, %lb1 to %ub1 step %st1) {
+ %use_i = arith.addi %i, %lb0 : index
+ %use_j = arith.addi %j, %lb1 : index
+ %_ = arith.cmpi ult, %use_i, %use_j : index
+
+ %e0 = omp.affinity_entry %addr0, %len0
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ %e1 = omp.affinity_entry %addr1, %len1
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+
+ omp.yield(%e1 : !omp.affinity_entry_ty<!llvm.ptr, i64>)
+ } -> !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>
+
+ omp.task affinity(%it : !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>) {
+ omp.terminator
+ }
+
+ return
+}
>From 8fdddc751c96e21979b13bb5613b8189bd308ebd Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Wed, 18 Feb 2026 15:04:48 -0600
Subject: [PATCH 4/4] Fix LLVMIR openmp todo test for affinity clause
---
mlir/test/Target/LLVMIR/openmp-todo.mlir | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 9a10ad74baeb6..ae02b5878f763 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -475,10 +475,12 @@ llvm.func @wsloop_order(%lb : i32, %ub : i32, %step : i32) {
}
// -----
-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) {
+llvm.func @task_affinity(%ptr : !llvm.ptr, %len : i64) {
+ // expected-error at below {{not yet implemented: omp.affinity_entry}}
+ // expected-error at below {{LLVM Translation failed for operation: omp.affinity_entry}}
+ %ae = omp.affinity_entry %ptr, %len
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ omp.task affinity(%ae : !omp.affinity_entry_ty<!llvm.ptr, i64>) {
omp.terminator
}
llvm.return
More information about the llvm-branch-commits
mailing list