[llvm-branch-commits] [flang] [mlir] [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (PR #182222)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Feb 24 15:54:17 PST 2026
https://github.com/chichunchen updated https://github.com/llvm/llvm-project/pull/182222
>From 1eaded7e0dac63055a6d9635bbf47e31155cb929 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/6] [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 3cce660b43a33..88c8ab4f6f949 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
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Iterator Op
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 7d73063321299..1a312b9f2e65f 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 937cb4df60eb3ea7afeefc0a89b412050fcfa42f 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/6] 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 f2d039d7bf5ba..4d2cb8d8a5877 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 104d60d77ea5c920c37d2139a3771166d4555d73 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/6] 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 1a312b9f2e65f..97a22f3e91674 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3600,3 +3600,56 @@ func.func @omp_iterator_2d(%s2 : !llvm.struct<(ptr, i64)>) -> () {
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 7a84f3cfb3336109af538cb43411c0fea955432f 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/6] 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
>From 351236e1bb617101c29e00ce6af3ff8116b36b78 Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Tue, 24 Feb 2026 14:16:17 -0600
Subject: [PATCH 5/6] Rewrite functions in affinity utility functions with
hlfir apis
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 70 ++----
flang/lib/Lower/OpenMP/Utils.cpp | 251 ++++++++++-----------
flang/lib/Lower/OpenMP/Utils.h | 8 +-
flang/test/Lower/OpenMP/task-affinity.f90 | 115 ++++++----
4 files changed, 222 insertions(+), 222 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index f2246377b5ffd..cc08cffc55b03 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -202,14 +202,21 @@ getIfClauseOperand(lower::AbstractConverter &converter,
ifVal);
}
-template <typename IteratorSpecT>
+template <typename SomeType, 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);
+ using IdTy =
+ Fortran::lower::omp::IdTyTemplate<Fortran::evaluate::Expr<SomeType>>;
+ using ExprTy = Fortran::evaluate::Expr<SomeType>;
+
+ using ObjTy = tomp::type::ObjectT<IdTy, ExprTy>;
+ using RangeTy = tomp::type::RangeT<ExprTy>;
+
+ const ObjTy &ivObj = std::get<1>(itSpec.t);
+ const RangeTy &range = std::get<2>(itSpec.t);
IteratorRange r;
r.ivSym = ivObj.sym();
@@ -228,8 +235,7 @@ static IteratorRange lowerIteratorRange(
mlir::Value v) -> mlir::Value {
if (v.getType().isIndex())
return v;
- return mlir::arith::IndexCastOp::create(builder, loc,
- builder.getIndexType(), v);
+ return fir::ConvertOp::create(builder, loc, builder.getIndexType(), v);
};
r.lb = toIndex(builder, loc, lbVal);
@@ -853,38 +859,10 @@ bool ClauseProcessor::processAffinity(
&context, refI8Ty, builder.getI64Type());
mlir::Type iterTy = mlir::omp::IteratedType::get(&context, entryTy);
- 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);
+ mlir::Value addrI8 = fir::ConvertOp::create(b, l, refI8Ty, addr);
return mlir::omp::AffinityEntryOp::create(b, l, entryTy, addrI8, len)
.getResult();
};
@@ -900,7 +878,8 @@ bool ClauseProcessor::processAffinity(
iteratorRanges.reserve(iteratorModifierSpecs.size());
for (const auto &itSpec : iteratorModifierSpecs)
iteratorRanges.push_back(
- lowerIteratorRange(converter, itSpec, stmtCtx, clauseLocation));
+ lowerIteratorRange<Fortran::evaluate::SomeType>(
+ converter, itSpec, stmtCtx, clauseLocation));
for (const IteratorRange &r : iteratorRanges)
ivSyms.insert(&r.ivSym->GetUltimate());
@@ -920,31 +899,28 @@ bool ClauseProcessor::processAffinity(
Fortran::lower::getDataOperandBaseAddr(
converter, builder, *sym, loc,
/*unwrapFirBox=*/false);
- // TODO check correctness of genIteratorCoordinate
+ hlfir::Entity entity{info.addr};
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()));
+ genIteratorCoordinate(converter, entity, ivs, loc);
+ mlir::Value len = builder.createIntegerConstant(
+ loc, builder.getIntegerType(64),
+ getElementBytesOrZero(entity, builder.getDataLayout()));
return makeAffinityEntry(builder, loc, entryTy, addr, len);
});
- iterHandle.dump();
result.iterated.push_back(iterHandle);
} else {
mlir::Value addr =
genAffinityAddr(converter, object, stmtCtx, clauseLocation);
+ // get hlfir.declare for length calculation
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
+ mlir::Value len =
+ genAffinityLen(builder, clauseLocation, builder.getDataLayout(),
+ hlfir::Entity{info.addr}, bounds);
result.affinityVars.push_back(
makeAffinityEntry(builder, clauseLocation, entryTy, addr, len));
}
diff --git a/flang/lib/Lower/OpenMP/Utils.cpp b/flang/lib/Lower/OpenMP/Utils.cpp
index 83d1a977cb3e3..185cab55cd777 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -921,102 +921,136 @@ 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))
+ fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+
+ auto genRawAddress = [&](mlir::Value v) -> mlir::Value {
+ // Only wrap with hlfir::Entity if it is a Fortran entity (avoids asserts).
+ if (!hlfir::isFortranEntity(v))
+ return v;
+
+ hlfir::Entity entity{v};
+ entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
+ return hlfir::genVariableRawAddress(loc, builder, entity);
+ };
+
+ // 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);
+ mlir::Value baseAddr = fir::getBase(exv);
+ return genRawAddress(baseAddr);
}
- // Fallback to base symbol address: affinity(a)
+ // 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;
+ mlir::Value symAddr = converter.getSymbolAddress(*sym);
+ return genRawAddress(symAddr);
}
-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();
+int64_t getElementBytesOrZero(hlfir::Entity entity,
+ const mlir::DataLayout &dl) {
+ if (entity.isBoxAddressOrValue() || entity.isAssumedRank())
+ return 0;
- // ((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::Type elemTy = entity.getFortranElementType();
-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))
+ if (auto charTy = mlir::dyn_cast<fir::CharacterType>(elemTy)) {
+ auto len = hlfir::getCharLengthIfConst(entity);
+ if (!len)
return 0;
- mlir::Type eleTy = fir::unwrapPassByRefType(baseTy);
- eleTy = fir::unwrapSequenceType(eleTy);
- return static_cast<int64_t>(dl.getTypeSize(eleTy));
- };
+ const int64_t charBytes = charTy.getFKind();
+ return charBytes * (*len);
+ }
- auto getWholeObjectBytesIfStaticOrZero = [&](mlir::Type addrTy) -> int64_t {
- if (isDescriptorLike(addrTy))
- return 0;
+ if (fir::isRecordWithTypeParameters(elemTy))
+ return 0;
- mlir::Type eleTy = fir::unwrapPassByRefType(addrTy);
+ return static_cast<int64_t>(dl.getTypeSize(elemTy));
+}
- // Scalar
- if (!mlir::isa<fir::SequenceType>(eleTy))
- return static_cast<int64_t>(dl.getTypeSize(eleTy));
+// return 0 if the total number of elements cannot be determined at compile time
+static int getTotalElement(hlfir::Entity entity, mlir::Location loc,
+ fir::FirOpBuilder &builder) {
+ if (entity.isBoxAddressOrValue() || entity.isAssumedRank())
+ return 0;
+ assert(!entity.isScalar() &&
+ "expected non-scalar entity to compute total elements");
+ const int rank = entity.getRank();
+ assert(rank >= 0 && "expected non-negative rank for non-assumed-rank entity");
+
+ int totalElems = 1;
+ for (unsigned d = 0; d < static_cast<unsigned>(rank); ++d) {
+ mlir::Value extent = hlfir::genExtent(loc, builder, entity, d);
+ auto cst = fir::getIntIfConstant(extent);
+ if (!cst)
+ return 0;
+ totalElems *= *cst;
+ }
+ return totalElems;
+}
- // 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;
- }
+// Compute the byte span of a Fortran array section described by bounds.
+// The span is (highest_address - lowest_address + 1) * element_size
+// where addresses are computed according to Fortran column-major layout.
+// For each dimension d, the highest address is:
+// highest_address = (ub_d - lb_d) * distance_d
+// distance_d = product of fullExtents[0..d-1]
+// distance_0 = 1
+// For example:
+// integer a[5][7], a(2:4, 3:5), lowest = a(2,3), highest = a(4,5)
+// lowest = (2-1) * (3-1) * 5 = 11
+// highest = (4-1) * (5-1) * 5 = 23
+// total span = 23 - 11 + 1 = 13 elements
+static mlir::Value computeBoundsSpan(fir::FirOpBuilder &builder,
+ mlir::Location loc,
+ llvm::ArrayRef<mlir::Value> bounds,
+ hlfir::Entity entity) {
+ assert(!bounds.empty() && "expected non-empty bounds to compute span");
+ auto fullExtents = hlfir::genExtentsVector(loc, builder, entity);
+ assert(fullExtents.size() == bounds.size() &&
+ "expected bounds and full extents to have the same size");
+ mlir::Value one =
+ builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+ mlir::Value span = one; // inclusive: +1
+ mlir::Value distance = one; // column-major linearization factor
+ for (auto [b, extent] : llvm::zip(bounds, fullExtents)) {
+ auto mb = b.getDefiningOp<mlir::omp::MapBoundsOp>();
+ mlir::Value delta = mlir::arith::SubIOp::create(
+ builder, loc, mb.getUpperBound(), mb.getLowerBound());
+
+ span = mlir::arith::AddIOp::create(
+ builder, loc, span,
+ mlir::arith::MulIOp::create(builder, loc, delta, distance));
+
+ distance = mlir::arith::MulIOp::create(builder, loc, distance, extent);
+ }
+ // Convert from index to i64 (bounds are in index type)
+ return fir::ConvertOp::create(builder, loc, builder.getI64Type(), span);
+}
- int64_t elemBytes = static_cast<int64_t>(dl.getTypeSize(seqTy.getEleTy()));
- return elems * elemBytes;
- };
+mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
+ const mlir::DataLayout &dl, hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> bounds) {
+ int64_t elemBytes = getElementBytesOrZero(entity, dl);
- // 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);
-
+ mlir::Value spanElems = computeBoundsSpan(builder, loc, bounds, entity);
return mlir::arith::MulIOp::create(
- builder, loc, numElems,
- builder.createIntegerConstant(loc, builder.getI64Type(), elemBytes));
+ builder, loc, spanElems,
+ builder.createIntegerConstant(loc, builder.getIntegerType(64),
+ elemBytes));
}
// explicit ref => element size (a(3), a(i))
- if (hasRef) {
- int64_t elemBytes = getElementBytesOrZero(addr.getType());
+ if (entity.isScalar()) {
+ elemBytes = getElementBytesOrZero(entity, dl);
return builder.createIntegerConstant(loc, builder.getI64Type(), elemBytes);
}
// whole object => whole size if static, else 0
- int64_t wholeBytes = getWholeObjectBytesIfStaticOrZero(addr.getType());
+ int64_t wholeBytes =
+ getTotalElement(entity, loc, builder) * getElementBytesOrZero(entity, dl);
return builder.createIntegerConstant(loc, builder.getI64Type(), wholeBytes);
}
@@ -1038,69 +1072,34 @@ bool hasIVReference(
}
mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
- mlir::Value base,
+ hlfir::Entity entity,
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{});
- }
+ mlir::Value base = entity.getBase();
- // 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));
+ // If base is a reference-to-box, load it so array_coor sees the box value
+ if (auto refTy = mlir::dyn_cast<fir::ReferenceType>(base.getType())) {
+ if (mlir::isa<fir::BoxType>(refTy.getEleTy()))
+ base = fir::LoadOp::create(builder, loc, base);
}
- return fir::CoordinateOp::create(builder, loc, eleRefTy, base, subsI32);
+ // Build shape from the entity extents
+ auto extents = hlfir::genExtentsVector(loc, builder, entity);
+ assert(extents.size() == ivs.size() &&
+ "expected the number of extents and iteration variables to match for "
+ "iterator");
+ mlir::Value shape = fir::ShapeOp::create(builder, loc, extents);
+
+ mlir::Type elementToRefTy =
+ fir::ReferenceType::get(entity.getFortranElementType());
+
+ return fir::ArrayCoorOp::create(builder, loc, elementToRefTy,
+ /*memref=*/base,
+ /*shape=*/shape,
+ /*slice=*/mlir::Value{},
+ /*indices=*/ivs,
+ /*typeparams=*/mlir::ValueRange{});
}
} // namespace omp
diff --git a/flang/lib/Lower/OpenMP/Utils.h b/flang/lib/Lower/OpenMP/Utils.h
index 10552a68908d8..b6fe8482add8f 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -189,14 +189,16 @@ void collectTileSizesFromOpenMPConstruct(
llvm::SmallVectorImpl<int64_t> &tileSizes,
Fortran::semantics::SemanticsContext &semaCtx);
+int64_t getElementBytesOrZero(hlfir::Entity entity, const mlir::DataLayout &dl);
+
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);
+ const mlir::DataLayout &dl, hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> bounds);
struct IteratorRange {
mlir::Value lb;
@@ -210,7 +212,7 @@ bool hasIVReference(
const llvm::SmallPtrSetImpl<const Fortran::semantics::Symbol *> &ivSyms);
mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
- mlir::Value base,
+ hlfir::Entity entity,
llvm::ArrayRef<mlir::Value> ivs,
mlir::Location loc);
diff --git a/flang/test/Lower/OpenMP/task-affinity.f90 b/flang/test/Lower/OpenMP/task-affinity.f90
index 3e3f8acad7c62..615ccd74624a2 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -16,20 +16,19 @@ 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: %[[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: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[ELEM:.*]] = hlfir.designate %[[A]]#0 (%[[C1]]) : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+! CHECK: %[[C0:.*]] = arith.constant 0 : index
+! CHECK: %[[ONE:.*]] = arith.constant 1 : index
+! CHECK: %[[SUB:.*]] = arith.subi %[[C0]], %[[C0]] : index
+! CHECK: %[[MUL:.*]] = arith.muli %[[SUB]], %[[ONE]] : index
+! CHECK: %[[ADD:.*]] = arith.addi %[[ONE]], %[[MUL]] : index
+! CHECK: %[[CAST:.*]] = fir.convert %[[ADD]] : (index) -> 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>) {
subroutine omp_task_affinity_array_section()
implicit none
@@ -50,25 +49,17 @@ 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: %[[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: %[[C2:.*]] = arith.constant 2 : index
+! CHECK: %[[C50:.*]] = arith.constant 50 : index
+! CHECK: %[[C1:.*]] = arith.constant 1 : index
+! CHECK: %[[SHAPE:.*]] = fir.shape %{{.*}} : (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: %[[SPAN_I64:.*]] = fir.convert {{.*}} : (index) -> i64
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[LEN:.*]] = arith.muli %[[SPAN_I64]], %[[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 -> %{{.*}} : !fir.ref<i32>) {
subroutine omp_task_affinity_scalar()
implicit none
@@ -128,14 +119,15 @@ subroutine task_affinity_iterator_simple()
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>>) {
+! CHECK: %[[ITERATED:.*]] = omp.iterators(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[SHAPE:.*]] = fir.shape %c16 : (index) -> !fir.shape<1>
+! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV]] : (!fir.ref<!fir.array<16xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[COOR]] : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[C4]] : (!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
@@ -157,11 +149,42 @@ subroutine task_affinity_iterator_multi_dimension()
! 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: %[[SHAPE:.*]] = fir.shape %c4, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV0]], %[[IV1]] : (!fir.ref<!fir.array<4x6xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[COOR]] : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[C4]] : (!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>>)
+
+subroutine task_affinity_slice_2d
+ integer, parameter :: n = 5
+ integer, parameter :: m = 7
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(a(2:4, 3:5))
+ do i = 1, n
+ do j = 1, m
+ a(i, j) = i + j
+ end do
+ end do
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_slice_2d()
+! CHECK: omp.parallel {
+! CHECK: omp.single {
+! CHECK: %[[BOX:.*]] = hlfir.designate {{.*}} : (!fir.ref<!fir.array<5x7xi32>>, index, index, index, index, index, index, !fir.shape<2>) -> !fir.box<!fir.array<3x3xi32>>
+! CHECK: %[[BASE:.*]] = fir.box_addr %[[BOX]] : (!fir.box<!fir.array<3x3xi32>>) -> !fir.ref<!fir.array<3x3xi32>>
+! CHECK: %[[SPANI64:.*]] = fir.convert {{.*}} : (index) -> i64
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[LEN:.*]] = arith.muli %[[SPANI64]], %[[C4]] : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[BASE]] : (!fir.ref<!fir.array<3x3xi32>>) -> !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>){{.*}} {
>From cb7fc4a5c1e0da12225e38afc7d51c97550681dc Mon Sep 17 00:00:00 2001
From: cchen <chichun.chen at hpe.com>
Date: Tue, 24 Feb 2026 17:16:17 -0600
Subject: [PATCH 6/6] Fix build and tests
- Update due to rebase of omp.iterators -> omp.iterator
- Fix build by moving HLFIROps.h from Utils.h to Utils.cpp
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 2 +-
flang/lib/Lower/OpenMP/Utils.cpp | 1 -
flang/lib/Lower/OpenMP/Utils.h | 1 +
flang/test/Lower/OpenMP/task-affinity.f90 | 4 ++--
mlir/test/Dialect/OpenMP/ops.mlir | 8 ++++----
5 files changed, 8 insertions(+), 8 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index cc08cffc55b03..eaa13efe8cafb 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -270,7 +270,7 @@ static mlir::Value buildIteratorOp(Fortran::lower::AbstractConverter &converter,
steps.push_back(r.step);
}
- auto itOp = mlir::omp::IteratorsOp::create(
+ auto itOp = mlir::omp::IteratorOp::create(
builder, loc, iterTy, mlir::ValueRange{lbs}, mlir::ValueRange{ubs},
mlir::ValueRange{steps});
diff --git a/flang/lib/Lower/OpenMP/Utils.cpp b/flang/lib/Lower/OpenMP/Utils.cpp
index 185cab55cd777..6e07a9e656351 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -24,7 +24,6 @@
#include <flang/Optimizer/Builder/BoxValue.h>
#include <flang/Optimizer/Builder/FIRBuilder.h>
#include <flang/Optimizer/Builder/Todo.h>
-#include <flang/Optimizer/HLFIR/HLFIROps.h>
#include <flang/Parser/openmp-utils.h>
#include <flang/Parser/parse-tree.h>
#include <flang/Parser/tools.h>
diff --git a/flang/lib/Lower/OpenMP/Utils.h b/flang/lib/Lower/OpenMP/Utils.h
index b6fe8482add8f..c811e6e8cea0f 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -10,6 +10,7 @@
#define FORTRAN_LOWER_OPENMPUTILS_H
#include "flang/Lower/OpenMP/Clauses.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Value.h"
diff --git a/flang/test/Lower/OpenMP/task-affinity.f90 b/flang/test/Lower/OpenMP/task-affinity.f90
index 615ccd74624a2..2289e2cf8051a 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -119,7 +119,7 @@ subroutine task_affinity_iterator_simple()
end subroutine
! CHECK-LABEL: func.func @_QPtask_affinity_iterator_simple()
-! CHECK: %[[ITERATED:.*]] = omp.iterators(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[ITERATED:.*]] = omp.iterator(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
! CHECK: %[[SHAPE:.*]] = fir.shape %c16 : (index) -> !fir.shape<1>
! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV]] : (!fir.ref<!fir.array<16xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
! CHECK: %[[C4:.*]] = arith.constant 4 : i64
@@ -148,7 +148,7 @@ subroutine task_affinity_iterator_multi_dimension()
end subroutine
! CHECK-LABEL: func.func @_QPtask_affinity_iterator_multi_dimension()
-! CHECK: %[[ITER:.*]] = omp.iterators(%[[IV0:.*]]: index, %[[IV1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[ITER:.*]] = omp.iterator(%[[IV0:.*]]: index, %[[IV1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
! CHECK: %[[SHAPE:.*]] = fir.shape %c4, %c6 : (index, index) -> !fir.shape<2>
! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV0]], %[[IV1]] : (!fir.ref<!fir.array<4x6xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
! CHECK: %[[C4:.*]] = arith.constant 4 : i64
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 97a22f3e91674..6d549ff890f24 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3604,13 +3604,13 @@ func.func @omp_iterator_2d(%s2 : !llvm.struct<(ptr, i64)>) -> () {
// 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: %[[IT:.*]] = omp.iterator(%[[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) {
+ %it = omp.iterator(%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>)
@@ -3627,14 +3627,14 @@ func.func @omp_task_affinity_iterator_2d(%lb0 : index, %ub0 : index, %st0 : inde
%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: %[[IT:.*]] = omp.iterator(%[[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) {
+ %it = omp.iterator(%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
More information about the llvm-branch-commits
mailing list