[flang-commits] [flang] e4160cb - [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (#182222)
via flang-commits
flang-commits at lists.llvm.org
Sun Mar 8 16:32:22 PDT 2026
Author: Chi-Chun, Chen
Date: 2026-03-08T18:32:17-05:00
New Revision: e4160cb52d48f63abb520dcba8d212ea79f617cf
URL: https://github.com/llvm/llvm-project/commit/e4160cb52d48f63abb520dcba8d212ea79f617cf
DIFF: https://github.com/llvm/llvm-project/commit/e4160cb52d48f63abb520dcba8d212ea79f617cf.diff
LOG: [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (#182222)
This patch translates the Flang AST to the OpenMP dialect for the
affinity clause, including support for the iterator modifier.
2/3 in stack for implementing affinity clause with iterator modifier
1/3 #182218
2/3 #182222
3/3 #182223
Added:
Modified:
flang/lib/Lower/OpenMP/ClauseProcessor.cpp
flang/lib/Lower/OpenMP/Utils.cpp
flang/lib/Lower/OpenMP/Utils.h
flang/test/Lower/OpenMP/task-affinity.f90
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/test/Dialect/OpenMP/ops.mlir
mlir/test/Target/LLVMIR/openmp-todo.mlir
Removed:
flang/test/Lower/OpenMP/Todo/affinity-clause.f90
################################################################################
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 9696bc21b5fbb..4eed4231646ae 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -202,6 +202,125 @@ getIfClauseOperand(lower::AbstractConverter &converter,
ifVal);
}
+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();
+
+ 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();
+ 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 &stExpr = 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 fir::ConvertOp::create(builder, loc, builder.getIndexType(), v);
+ };
+
+ r.lb = toIndex(builder, loc, lbVal);
+ r.ub = toIndex(builder, loc, ubVal);
+
+ if (stExpr) {
+ mlir::Value stVal =
+ fir::getBase(converter.genExprValue(toEvExpr(*stExpr), stmtCtx));
+ r.step = toIndex(builder, loc, stVal);
+ } else {
+ r.step = mlir::arith::ConstantIndexOp::create(builder, loc, 1);
+ }
+
+ return r;
+}
+
+template <typename BodyFn>
+static mlir::Value buildIteratorOp(Fortran::lower::AbstractConverter &converter,
+ mlir::Location loc, mlir::Type iterTy,
+ llvm::ArrayRef<IteratorRange> ranges,
+ BodyFn &&bodyGen) {
+
+ 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::IteratorOp::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) {
+ mlir::Value ivVal = ivs[i];
+ mlir::Type ivTy = converter.genType(*ranges[i].ivSym);
+ if (ivVal.getType() != ivTy)
+ ivVal = fir::ConvertOp::create(builder, loc, ivTy, ivVal);
+ symMap.addSymbol(*ranges[i].ivSym, ivVal, /*force=*/true);
+ }
+
+ mlir::omp::YieldOp::create(builder, loc, bodyGen(builder, loc, ivs));
+
+ return itOp.getResult();
+}
+
+template <typename ClauseTuple>
+static void collectIteratorIVs(
+ const ClauseTuple &clause, Fortran::lower::AbstractConverter &converter,
+ Fortran::lower::StatementContext &stmtCtx,
+ llvm::SmallVectorImpl<IteratorRange> &iteratorRanges,
+ llvm::SmallPtrSetImpl<const Fortran::semantics::Symbol *> &ivSyms) {
+ auto &iteratorModifier =
+ std::get<std::optional<omp::clause::Iterator>>(clause.t);
+ if (!iteratorModifier.has_value())
+ return;
+
+ mlir::Location clauseLocation = converter.getCurrentLocation();
+ const auto &iteratorModifierSpecs = *iteratorModifier;
+ iteratorRanges.reserve(iteratorModifierSpecs.size());
+ for (const auto &itSpec : iteratorModifierSpecs)
+ iteratorRanges.push_back(lowerIteratorRange<Fortran::evaluate::SomeType>(
+ converter, itSpec, stmtCtx, clauseLocation));
+
+ for (const IteratorRange &r : iteratorRanges)
+ ivSyms.insert(&r.ivSym->GetUltimate());
+}
+
//===----------------------------------------------------------------------===//
// ClauseProcessor unique clauses
//===----------------------------------------------------------------------===//
@@ -761,14 +880,82 @@ bool ClauseProcessor::processAffinity(
mlir::omp::AffinityClauseOps &result) const {
return findRepeatableClause<omp::clause::Affinity>(
[&](const omp::clause::Affinity &clause, const parser::CharBlock &) {
- if (std::get<std::optional<omp::clause::Iterator>>(clause.t)) {
- TODO(converter.getCurrentLocation(),
- "Support for iterator modifiers is not implemented yet");
- }
-
const auto &objects = std::get<omp::ObjectList>(clause.t);
- if (!objects.empty())
- genObjectList(objects, converter, result.affinityVars);
+ 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());
+ mlir::Type iterTy =
+ mlir::omp::IteratedType::get(&converter.getMLIRContext(), entryTy);
+
+ auto makeAffinityEntry = [&](fir::FirOpBuilder &b, mlir::Location l,
+ mlir::Type entryTy, mlir::Value addr,
+ mlir::Value len) -> mlir::Value {
+ mlir::Value addrI8 = fir::ConvertOp::create(b, l, refI8Ty, addr);
+ return mlir::omp::AffinityEntryOp::create(b, l, entryTy, addrI8, len)
+ .getResult();
+ };
+
+ llvm::SmallVector<IteratorRange> iteratorRanges;
+ llvm::SmallPtrSet<const Fortran::semantics::Symbol *, 4> ivSyms;
+
+ auto &iteratorModifier =
+ std::get<std::optional<omp::clause::Iterator>>(clause.t);
+ collectIteratorIVs(clause, converter, stmtCtx, iteratorRanges, ivSyms);
+
+ for (const omp::Object &object : objects) {
+ llvm::SmallVector<mlir::Value> bounds;
+ std::stringstream asFortran;
+ if (iteratorModifier.has_value() &&
+ hasIteratorIVReference(object, ivSyms)) {
+ mlir::Value iterHandle = buildIteratorOp(
+ converter, clauseLocation, iterTy, iteratorRanges,
+ [&](fir::FirOpBuilder &builder, mlir::Location loc,
+ llvm::ArrayRef<mlir::Value> /*ivs*/) -> mlir::Value {
+ lower::StatementContext iterStmtCtx;
+
+ if (std::optional<llvm::SmallVector<mlir::Value>>
+ loweredIndices = getIteratorElementIndices(
+ converter, object, iterStmtCtx, loc)) {
+ 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);
+ hlfir::Entity entity{info.addr};
+ mlir::Value iteratedAddr = genIteratorCoordinate(
+ converter, entity, *loweredIndices, loc);
+ mlir::Value len = genElementSizeInBytes(
+ builder, loc, builder.getDataLayout(), entity);
+ return makeAffinityEntry(builder, loc, entryTy,
+ iteratedAddr, len);
+ }
+
+ TODO(loc, "object type not supported by iterator modifier");
+ });
+ 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(),
+ hlfir::Entity{info.addr}, bounds);
+ 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 394c7a485525f..a0fd386e8fecb 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -913,6 +913,308 @@ void collectLoopRelatedInfo(
convertLoopBounds(converter, currentLocation, result, loopVarTypeSize);
}
+// Lower an affinity object to the raw storage address.
+// The lowering paths feeding this helper are mixed: some produce HLFIR
+// entities such as hlfir.designate/hlfir.declare, while others already
+// produce raw FIR addresses such as fir.box_addr. Normalize entity-like values
+// to a raw address, and leave already-raw addresses unchanged.
+mlir::Value genAffinityAddr(Fortran::lower::AbstractConverter &converter,
+ const omp::Object &object,
+ Fortran::lower::StatementContext &stmtCtx,
+ mlir::Location loc) {
+ fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+
+ auto genRawAddress = [&](mlir::Value v) -> mlir::Value {
+ // Examples seen here include hlfir.designate for a(i), hlfir.declare for
+ // whole objects like dummy/character arrays, fir.load of a pointer box,
+ // and already-raw fir.box_addr results. Only the entity-like cases can be
+ // wrapped as hlfir::Entity; the raw address cases must be returned as-is.
+ if (!hlfir::isFortranEntity(v))
+ return v;
+
+ hlfir::Entity entity{v};
+ // Pointer/allocatable entities need to be dereferenced first so affinity
+ // uses the pointee storage rather than the box address.
+ entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
+ return hlfir::genVariableRawAddress(loc, builder, entity);
+ };
+
+ // Designators such as affinity(a(3)) or affinity(a(1:10)) lower through
+ // genExprAddr. The base may still be an HLFIR entity, or may already be a
+ // raw FIR address after earlier lowering.
+ if (auto expr = object.ref()) {
+ fir::ExtendedValue exv =
+ converter.genExprAddr(toEvExpr(*expr), stmtCtx, &loc);
+ mlir::Value baseAddr = fir::getBase(exv);
+ return genRawAddress(baseAddr);
+ }
+
+ // Whole objects such as affinity(a) come from the symbol address directly.
+ const Fortran::semantics::Symbol *sym = object.sym();
+ assert(sym && "expected symbol in affinity object");
+ mlir::Value symAddr = converter.getSymbolAddress(*sym);
+ return genRawAddress(symAddr);
+}
+
+// Compute the size in bytes of a single element described by an HLFIR entity.
+// This returns the per-element byte size only; callers handle any array extent
+// or section span separately.
+mlir::Value genElementSizeInBytes(fir::FirOpBuilder &builder,
+ mlir::Location loc,
+ const mlir::DataLayout &dl,
+ hlfir::Entity entity) {
+ // Boxed entities carry the runtime element size in the descriptor.
+ if (entity.isBoxAddressOrValue())
+ return fir::ConvertOp::create(
+ builder, loc, builder.getI64Type(),
+ fir::BoxEleSizeOp::create(builder, loc, builder.getIndexType(),
+ entity));
+
+ mlir::Type elemTy = entity.getFortranElementType();
+
+ if (auto charTy = mlir::dyn_cast<fir::CharacterType>(elemTy)) {
+ // Non-box character entities expose length separately; multiply it by the
+ // character kind byte width.
+ mlir::Value charLen = hlfir::genCharLength(loc, builder, entity);
+ mlir::Value charBytes = builder.createIntegerConstant(
+ loc, builder.getI64Type(), charTy.getFKind());
+ return mlir::arith::MulIOp::create(
+ builder, loc,
+ fir::ConvertOp::create(builder, loc, builder.getI64Type(), charLen),
+ charBytes);
+ }
+
+ // PDTs with length parameters and assumed-rank entities do not currently
+ // have a precise byte size here, so keep the existing conservative 0.
+ if (fir::isRecordWithTypeParameters(elemTy) || entity.isAssumedRank())
+ return builder.createIntegerConstant(loc, builder.getI64Type(), 0);
+
+ // Trivial non-box entities have a fixed element size in the data layout.
+ return builder.createIntegerConstant(
+ loc, builder.getI64Type(), static_cast<int64_t>(dl.getTypeSize(elemTy)));
+}
+
+// Compute the total number of elements in a whole affinity object.
+static mlir::Value getTotalElements(fir::FirOpBuilder &builder,
+ mlir::Location loc, hlfir::Entity entity) {
+ if (entity.isAssumedRank())
+ return builder.createIntegerConstant(loc, builder.getI64Type(), 0);
+
+ assert(!entity.isScalar() &&
+ "expected non-scalar entity to compute total elements");
+
+ mlir::Value total =
+ builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+ for (mlir::Value extent : hlfir::genExtentsVector(loc, builder, entity))
+ total = mlir::arith::MulIOp::create(builder, loc, total, extent);
+ return fir::ConvertOp::create(builder, loc, builder.getI64Type(), total);
+}
+
+// Compute the contiguous element span covered by an array section.
+// This is not the number of selected elements. Instead, it is the inclusive
+// distance from the lowest addressed element in the section to the highest
+// addressed element, using Fortran column-major layout. genAffinityLen later
+// multiplies this span by the element size to get the byte length.
+//
+// For each dimension d:
+// delta_d = upper_d - lower_d
+// distance_d = product(fullExtents[0..d-1])
+// with distance_0 = 1.
+//
+// Example:
+// integer :: a(5, 7)
+// !$omp task affinity(a(2:4, 3:5))
+// The section selects 9 elements, but its contiguous span runs from a(2,3) to
+// a(4,5). In linearized column-major indices, those are 11 and 23, so the
+// span is 23 - 11 + 1 = 13 elements.
+//
+// Strides in the section bounds do not change this computation: the span still
+// covers the full contiguous address range between the first and last element.
+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>();
+ assert(mb && "expected omp.map_bounds for affinity section span");
+ 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);
+}
+
+// Compute the byte length covered by an affinity object.
+// For a scalar or single element, this is the element size. For a section, it
+// is the span of the section in elements multiplied by the element size. For a
+// whole array object, it is the total number of elements multiplied by the
+// element size.
+mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
+ const mlir::DataLayout &dl, hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> bounds) {
+ mlir::Value elemBytes = genElementSizeInBytes(builder, loc, dl, entity);
+
+ // Scalar entities and single designated elements contribute exactly one
+ // element to the affinity object.
+ if (entity.isScalar())
+ return elemBytes;
+
+ if (!bounds.empty()) {
+ // Array sections carry explicit bounds describing the covered span.
+ mlir::Value spanElems = computeBoundsSpan(builder, loc, bounds, entity);
+ return mlir::arith::MulIOp::create(builder, loc, spanElems, elemBytes);
+ }
+
+ // Whole-array objects have no explicit bounds here, so use the extents of
+ // the entity itself.
+ return mlir::arith::MulIOp::create(
+ builder, loc, getTotalElements(builder, loc, entity), elemBytes);
+}
+
+bool hasIteratorIVReference(
+ 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;
+}
+
+// Build the array coordinate for an object that uses iterator variables.
+// If the object is a section, use the first element of that section
+// as the coordinate. Currently only support top-level ArrayRef designators.
+//
+// Examples:
+// a(i, j) -> coordinates for a(i, j)
+// a(i:i+1, j+2) -> coordinates for a(i, j+2)
+std::optional<llvm::SmallVector<mlir::Value>> getIteratorElementIndices(
+ Fortran::lower::AbstractConverter &converter, const omp::Object &object,
+ Fortran::lower::StatementContext &stmtCtx, mlir::Location loc) {
+ const std::optional<ExprTy> &ref = object.ref();
+ assert(ref && "expected iterator-dependent object to have a reference");
+
+ std::optional<Fortran::evaluate::DataRef> dataRef =
+ Fortran::evaluate::ExtractDataRef(*ref);
+ if (!dataRef)
+ return std::nullopt;
+ const auto *arrayRef = std::get_if<Fortran::evaluate::ArrayRef>(&dataRef->u);
+ if (!arrayRef || arrayRef->subscript().empty())
+ return std::nullopt;
+
+ auto &builder = converter.getFirOpBuilder();
+ const Fortran::semantics::Symbol *sym = object.sym();
+ assert(sym && "expected symbol for iterator-dependent object");
+ fir::ExtendedValue dataExv = converter.getSymbolExtendedValue(*sym);
+ mlir::Value one =
+ builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+ llvm::SmallVector<mlir::Value> indices;
+ indices.reserve(arrayRef->subscript().size());
+
+ for (const auto &[dim, subscript] : llvm::enumerate(arrayRef->subscript())) {
+ mlir::Value idx;
+ if (const auto *triplet =
+ std::get_if<Fortran::evaluate::Triplet>(&subscript.u)) {
+ // Sections use the first element of the section as the base address, so
+ // the coordinate for this dimension comes from the triplet lower bound.
+ std::optional<
+ Fortran::evaluate::Expr<Fortran::evaluate::SubscriptInteger>>
+ lowerBound = triplet->lower();
+ if (!lowerBound) {
+ // Get lower bound if not provided by user.
+ // For example: !$omp task affinity(iterator(i = 1:n, j = 1:m) : a(:i+1,
+ // j+2))
+ idx = fir::factory::readLowerBound(builder, loc, dataExv, dim, one);
+ } else {
+ idx = fir::getBase(
+ createSomeExtendedExpression(loc, converter, toEvExpr(*lowerBound),
+ converter.getSymbolMap(), stmtCtx));
+ }
+ } else {
+ // Not handling vector subscripts for now.
+ if (subscript.Rank() > 0)
+ return std::nullopt;
+
+ const auto *indirect =
+ std::get_if<Fortran::evaluate::IndirectSubscriptIntegerExpr>(
+ &subscript.u);
+ assert(indirect && "expected non-triplet subscript");
+
+ // Scalar subscripts, including reordered indices and expressions like
+ // i+1 or j+2, lower directly through expression lowering.
+ idx = fir::getBase(createSomeExtendedExpression(
+ loc, converter, toEvExpr(indirect->value()), converter.getSymbolMap(),
+ stmtCtx));
+ }
+ indices.push_back(idx);
+ }
+
+ return indices;
+}
+
+// Build the element address for an iterator-dependent affinity object from a
+// base entity and lowered indices.
+mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
+ hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> ivs,
+ mlir::Location loc) {
+ auto &builder = converter.getFirOpBuilder();
+ mlir::Value base = entity.getBase();
+
+ // 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);
+ }
+
+ // Build shape from the entity extents
+ mlir::Value shape;
+ auto extents = hlfir::genExtentsVector(loc, builder, entity);
+ assert(extents.size() == ivs.size() &&
+ "expected the number of extents and iteration variables to match for "
+ "iterator");
+ if (entity.mayHaveNonDefaultLowerBounds()) {
+ llvm::SmallVector<mlir::Value> lowerBounds;
+ lowerBounds.reserve(ivs.size());
+ for (unsigned dim = 0; dim < ivs.size(); ++dim)
+ lowerBounds.push_back(hlfir::genLBound(loc, builder, entity, dim));
+ shape = builder.genShape(loc, lowerBounds, extents);
+ } else {
+ 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
} // namespace lower
} // namespace Fortran
diff --git a/flang/lib/Lower/OpenMP/Utils.h b/flang/lib/Lower/OpenMP/Utils.h
index f707557197847..44381da853ecb 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -10,11 +10,13 @@
#define FORTRAN_LOWER_OPENMPUTILS_H
#include "flang/Lower/OpenMP/Clauses.h"
+#include "flang/Optimizer/Builder/HLFIRTools.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Value.h"
#include "llvm/Support/CommandLine.h"
#include <cstdint>
+#include <optional>
extern llvm::cl::opt<bool> treatIndexAsSection;
@@ -189,6 +191,40 @@ void collectTileSizesFromOpenMPConstruct(
llvm::SmallVectorImpl<int64_t> &tileSizes,
Fortran::semantics::SemanticsContext &semaCtx);
+mlir::Value genElementSizeInBytes(fir::FirOpBuilder &builder,
+ mlir::Location loc,
+ const mlir::DataLayout &dl,
+ hlfir::Entity entity);
+
+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, hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> bounds);
+
+struct IteratorRange {
+ mlir::Value lb;
+ mlir::Value ub;
+ mlir::Value step;
+ Fortran::semantics::Symbol *ivSym = nullptr;
+};
+
+bool hasIteratorIVReference(
+ const omp::Object &object,
+ const llvm::SmallPtrSetImpl<const Fortran::semantics::Symbol *> &ivSyms);
+
+mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
+ hlfir::Entity entity,
+ llvm::ArrayRef<mlir::Value> ivs,
+ mlir::Location loc);
+
+std::optional<llvm::SmallVector<mlir::Value>> getIteratorElementIndices(
+ Fortran::lower::AbstractConverter &converter, const omp::Object &object,
+ Fortran::lower::StatementContext &stmtCtx, mlir::Location loc);
+
} // namespace omp
} // namespace lower
} // namespace Fortran
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 66254e48e9b8e..e8a7b4bc44f9c 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -1,32 +1,36 @@
! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=52 -o - %s | FileCheck %s
-! scalar element locator
+! Non-iterator tests
+
subroutine omp_task_affinity_elem()
implicit none
integer, parameter :: n = 100
integer :: a(n)
!$omp parallel
- !$omp single
!$omp task affinity(a(1))
a(1) = 1
!$omp end task
- !$omp end single
!$omp end parallel
end subroutine omp_task_affinity_elem
! CHECK-LABEL: func.func @_QPomp_task_affinity_elem()
-! CHECK: %[[A1:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_elemEa"}
+! CHECK: %[[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: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
-
-! array section locator
+! 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: %[[C4:.*]] = arith.constant 4 : i64
+! 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: %[[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
integer, parameter :: n = 100
@@ -34,78 +38,524 @@ 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
! 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: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
-
-! scalar variable locator
+! 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: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[SPAN_I64:.*]] = fir.convert {{.*}} : (index) -> 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
integer :: s
s = 7
!$omp parallel
- !$omp single
!$omp task affinity(s)
s = s + 1
!$omp end task
- !$omp end single
!$omp end parallel
end subroutine omp_task_affinity_scalar
! CHECK-LABEL: func.func @_QPomp_task_affinity_scalar()
-! CHECK: %[[S3:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_scalarEs"}
+! CHECK: %[[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: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
-
-! multiple locators
+! 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>) {
+
subroutine omp_task_affinity_multi()
implicit none
integer, parameter :: n = 100
integer :: a(n), b(n)
!$omp parallel
- !$omp single
!$omp task affinity(a(1), b(1))
a(2) = 2
b(2) = 2
!$omp end task
- !$omp end single
!$omp end parallel
end subroutine omp_task_affinity_multi
! CHECK-LABEL: func.func @_QPomp_task_affinity_multi()
-! CHECK: %[[A4:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_multiEa"}
-! CHECK: %[[B4:.*]]:2 = hlfir.declare {{.*}} {uniq_name = "_QFomp_task_affinity_multiEb"}
+! CHECK: omp.parallel {
+! CHECK: %[[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>) {
+
+subroutine whole_array_affinity()
+ implicit none
+ integer :: a(10)
+
+ !$omp task affinity(a)
+ a(1) = 1
+ !$omp end task
+end subroutine whole_array_affinity
+
+! CHECK-LABEL: func.func @_QPwhole_array_affinity()
+! CHECK: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFwhole_array_affinityEa"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
+! CHECK: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[C4]] : i64
+! CHECK: %[[ADDRI8:.*]] = fir.convert %[[A]]#0 : (!fir.ref<!fir.array<10xi32>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDRI8]], %[[LEN]] : (!fir.ref<i8>, i64) -> !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: omp.task affinity(%[[A4]]#0 : !fir.ref<!fir.array<100xi32>>, %[[B4]]#0 : !fir.ref<!fir.array<100xi32>>) {
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: omp.terminator
-! CHECK: }
-! CHECK: return
+! 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: %[[C4:.*]] = arith.constant 4 : i64
+! CHECK: %[[SPANI64:.*]] = fir.convert {{.*}} : (index) -> 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>){{.*}} {
+
+subroutine assumed_shape_affinity(a)
+ integer, intent(inout) :: a(:)
+
+ !$omp task affinity(a)
+ a(1) = 1
+ !$omp end task
+end subroutine
+
+! CHECK-LABEL: func.func @_QPassumed_shape_affinity(
+! CHECK: %[[A:.*]]:2 = hlfir.declare %arg0 dummy_scope %{{.*}} arg 1 {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QFassumed_shape_affinityEa"}
+! CHECK: %[[ELEM:.*]] = fir.box_elesize %[[A]]#0 : (!fir.box<!fir.array<?xi32>>) -> index
+! CHECK: %[[ELEM_I64:.*]] = fir.convert %[[ELEM]] : (index) -> i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[ELEM_I64]] : i64
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine allocatable_affinity(n)
+ integer, intent(in) :: n
+ integer, allocatable :: a(:)
+
+ allocate(a(n))
+ !$omp task affinity(a)
+ a(1) = 1
+ !$omp end task
+end subroutine
+
+! CHECK-LABEL: func.func @_QPallocatable_affinity(
+! CHECK: %[[A:.*]]:2 = hlfir.declare %{{.*}} {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFallocatable_affinityEa"}
+! CHECK: %[[ABOX:.*]] = fir.load %[[A]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+! CHECK: %[[ELEM:.*]] = fir.box_elesize %{{.*}} : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> index
+! CHECK: %[[ELEM_I64:.*]] = fir.convert %[[ELEM]] : (index) -> i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[ELEM_I64]] : i64
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine pointer_affinity(n)
+ integer, intent(in) :: n
+ integer, target :: tgt(n)
+ integer, pointer :: p(:)
+
+ p => tgt
+ !$omp task affinity(p)
+ p(1) = 1
+ !$omp end task
+end subroutine
+
+! CHECK-LABEL: func.func @_QPpointer_affinity(
+! CHECK: %[[P:.*]]:2 = hlfir.declare %{{.*}} {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFpointer_affinityEp"}
+! CHECK: %[[PBOX:.*]] = fir.load %[[P]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
+! CHECK: %[[ELEM:.*]] = fir.box_elesize %{{.*}} : (!fir.box<!fir.ptr<!fir.array<?xi32>>>) -> index
+! CHECK: %[[ELEM_I64:.*]] = fir.convert %[[ELEM]] : (index) -> i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[ELEM_I64]] : i64
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine char_const_len_affinity()
+ character(len=7) :: a(8)
+
+ !$omp task affinity(a)
+ a(1) = "abcdefg"
+ !$omp end task
+end subroutine
+
+! CHECK-LABEL: func.func @_QPchar_const_len_affinity()
+! CHECK: %[[DECLARE:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) typeparams %c7 {uniq_name = "_QFchar_const_len_affinityEa"}
+! CHECK: %[[CHARLEN_I64:.*]] = fir.convert %c7 : (index) -> i64
+! CHECK: %[[ELEMSIZE:.*]] = arith.muli %[[CHARLEN_I64]], %{{.*}} : i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[ELEMSIZE]] : i64
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine char_runtime_len_affinity(n, l)
+ integer, intent(in) :: n, l
+ character(len=l) :: a(n)
+
+ !$omp task affinity(a)
+ a(1) = repeat("x", l)
+ !$omp end task
+end subroutine
+
+! CHECK-LABEL: func.func @_QPchar_runtime_len_affinity(
+! CHECK: %[[DECLARE:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) typeparams %{{.*}} {uniq_name = "_QFchar_runtime_len_affinityEa"}
+! CHECK: %[[ELEM:.*]] = fir.box_elesize %[[DECLARE]]#0 : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+! CHECK: %[[ELEM_I64:.*]] = fir.convert %[[ELEM]] : (index) -> i64
+! CHECK: %[[LEN:.*]] = arith.muli %{{.*}}, %[[ELEM_I64]] : i64
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %{{.*}}, %[[LEN]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+module task_affinity_polymorphic_mod
+ type :: t
+ integer :: x
+ end type
+contains
+ subroutine task_affinity_poly()
+ class(t), allocatable :: a
+
+ allocate(a)
+ !$omp task affinity(a) shared(a)
+ select type (a)
+ type is (t)
+ a%x = 1
+ end select
+ !$omp end task
+ end subroutine
+end module
+
+! CHECK-LABEL: func.func @_QMtask_affinity_polymorphic_modPtask_affinity_poly()
+! CHECK: %[[DECLARE:.*]]:2 = hlfir.declare %{{.*}} {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMtask_affinity_polymorphic_modFtask_affinity_polyEa"}
+! CHECK: %[[ALOAD:.*]] = fir.load %[[DECLARE]]#0 : !fir.ref<!fir.class<!fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>>>
+! CHECK: %[[ADDR:.*]] = fir.box_addr %[[ALOAD]] : (!fir.class<!fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>>) -> !fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>
+! CHECK: %[[SIZELOAD:.*]] = fir.load %[[DECLARE]]#0 : !fir.ref<!fir.class<!fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>>>
+! CHECK: %[[SIZE:.*]] = fir.box_elesize %[[SIZELOAD]] : (!fir.class<!fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>>) -> index
+! CHECK: %[[SIZE_I64:.*]] = fir.convert %[[SIZE]] : (index) -> i64
+! CHECK: %[[ADDR_I8:.*]] = fir.convert %[[ADDR]] : (!fir.heap<!fir.type<_QMtask_affinity_polymorphic_modTt{x:i32}>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDR_I8]], %[[SIZE_I64]] : (!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>) {
+
+! Iterator tests
+
+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.iterator(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IV_I32:.*]] = fir.convert %[[IV]] : (index) -> i32
+! CHECK: %[[IV_I64:.*]] = fir.convert %[[IV_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE:.*]] = fir.shape %c16 : (index) -> !fir.shape<1>
+! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV_I64]] : (!fir.ref<!fir.array<16xi32>>, !fir.shape<1>, i64) -> !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_nondefault_lb()
+ implicit none
+ integer, parameter :: n = 8
+ integer :: a(0:n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 0:n) : a(i))
+ a(0) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_nondefault_lb()
+! CHECK: %[[ITERATED_NDLB:.*]] = omp.iterator(%[[IV_NDLB:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IV_NDLB_I32:.*]] = fir.convert %[[IV_NDLB]] : (index) -> i32
+! CHECK: %[[IV_NDLB_I64:.*]] = fir.convert %[[IV_NDLB_I32]] : (i32) -> i64
+! CHECK: %[[SHIFT_NDLB:.*]] = fir.shape_shift %c0, %c9 : (index, index) -> !fir.shapeshift<1>
+! CHECK: %[[COOR_NDLB:.*]] = fir.array_coor {{.*}}(%[[SHIFT_NDLB]]) %[[IV_NDLB_I64]] : (!fir.box<!fir.array<9xi32>>, !fir.shapeshift<1>, i64) -> !fir.ref<i32>
+! CHECK: %[[ELEM_NDLB:.*]] = fir.box_elesize %{{.*}} : (!fir.box<!fir.array<9xi32>>) -> index
+! CHECK: %[[ELEM_NDLB_I64:.*]] = fir.convert %[[ELEM_NDLB]] : (index) -> i64
+! CHECK: %[[ADDRI8_NDLB:.*]] = fir.convert %[[COOR_NDLB]] : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY_NDLB:.*]] = omp.affinity_entry %[[ADDRI8_NDLB]], %[[ELEM_NDLB_I64]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.yield(%[[ENTRY_NDLB]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>)
+! CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>
+! CHECK: omp.task affinity(%[[ITERATED_NDLB]] : !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>) {
+
+subroutine task_affinity_iterator_nondefault_lb_2d()
+ implicit none
+ integer, parameter :: n = 4, m = 6
+ integer :: a(0:n, -1:m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 0:n, j = -1:m) : a(i, j))
+ a(0, -1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_nondefault_lb_2d()
+! CHECK: %[[ITERATED_NDLB2:.*]] = omp.iterator(%[[IV0_NDLB2:.*]]: index, %[[IV1_NDLB2:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IV0_NDLB2_I32:.*]] = fir.convert %[[IV0_NDLB2]] : (index) -> i32
+! CHECK: %[[IV1_NDLB2_I32:.*]] = fir.convert %[[IV1_NDLB2]] : (index) -> i32
+! CHECK: %[[IV0_NDLB2_I64:.*]] = fir.convert %[[IV0_NDLB2_I32]] : (i32) -> i64
+! CHECK: %[[IV1_NDLB2_I64:.*]] = fir.convert %[[IV1_NDLB2_I32]] : (i32) -> i64
+! CHECK: %[[SHIFT_NDLB2:.*]] = fir.shape_shift %c0, %c5, %c-1, %c8 : (index, index, index, index) -> !fir.shapeshift<2>
+! CHECK: %[[COOR_NDLB2:.*]] = fir.array_coor {{.*}}(%[[SHIFT_NDLB2]]) %[[IV0_NDLB2_I64]], %[[IV1_NDLB2_I64]] : (!fir.box<!fir.array<5x8xi32>>, !fir.shapeshift<2>, i64, i64) -> !fir.ref<i32>
+! CHECK: %[[ELEM_NDLB2:.*]] = fir.box_elesize %{{.*}} : (!fir.box<!fir.array<5x8xi32>>) -> index
+! CHECK: %[[ELEM_NDLB2_I64:.*]] = fir.convert %[[ELEM_NDLB2]] : (index) -> i64
+! CHECK: %[[ADDRI8_NDLB2:.*]] = fir.convert %[[COOR_NDLB2]] : (!fir.ref<i32>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY_NDLB2:.*]] = omp.affinity_entry %[[ADDRI8_NDLB2]], %[[ELEM_NDLB2_I64]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+! CHECK: omp.yield(%[[ENTRY_NDLB2]] : !omp.affinity_entry_ty<!fir.ref<i8>, i64>)
+! CHECK: } -> !omp.iterated<!omp.affinity_entry_ty<!fir.ref<i8>, i64>>
+! CHECK: omp.task affinity(%[[ITERATED_NDLB2]] : !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.iterator(%[[IV0:.*]]: index, %[[IV1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IV0_I32:.*]] = fir.convert %[[IV0]] : (index) -> i32
+! CHECK: %[[IV1_I32:.*]] = fir.convert %[[IV1]] : (index) -> i32
+! CHECK: %[[IV0_I64:.*]] = fir.convert %[[IV0_I32]] : (i32) -> i64
+! CHECK: %[[IV1_I64:.*]] = fir.convert %[[IV1_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE:.*]] = fir.shape %c4, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[IV0_I64]], %[[IV1_I64]] : (!fir.ref<!fir.array<4x6xi32>>, !fir.shape<2>, i64, i64) -> !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_iterator_reordered()
+ integer, parameter :: n = 4, m = 6
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n, j = 1:m) : a(j, i)) shared(a)
+ a(1, 1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_reordered()
+! CHECK: %[[ITER:.*]] = omp.iterator(%[[IV0:.*]]: index, %[[IV1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[RO_IV0_I32:.*]] = fir.convert %[[IV0]] : (index) -> i32
+! CHECK: %[[RO_IV1_I32:.*]] = fir.convert %[[IV1]] : (index) -> i32
+! CHECK: %[[RO_IV1_I64:.*]] = fir.convert %[[RO_IV1_I32]] : (i32) -> i64
+! CHECK: %[[RO_IV0_I64:.*]] = fir.convert %[[RO_IV0_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE:.*]] = fir.shape %c4, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR:.*]] = fir.array_coor {{.*}}(%[[SHAPE]]) %[[RO_IV1_I64]], %[[RO_IV0_I64]] : (!fir.ref<!fir.array<4x6xi32>>, !fir.shape<2>, i64, i64) -> !fir.ref<i32>
+
+subroutine task_affinity_iterator_expr_subscript()
+ integer, parameter :: n = 5, m = 6
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n-1, j = 1:m) : a(i+1, j)) shared(a)
+ a(1, 1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_expr_subscript()
+! CHECK: %[[ITER2:.*]] = omp.iterator(%[[IVA:.*]]: index, %[[IVB:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IVA_I32:.*]] = fir.convert %[[IVA]] : (index) -> i32
+! CHECK: %[[IVB_I32:.*]] = fir.convert %[[IVB]] : (index) -> i32
+! CHECK: %[[C1_I32:.*]] = arith.constant 1 : i32
+! CHECK: %[[IP1_I32:.*]] = arith.addi %[[IVA_I32]], %[[C1_I32]] : i32
+! CHECK: %[[IP1_I64:.*]] = fir.convert %[[IP1_I32]] : (i32) -> i64
+! CHECK: %[[IVB_I64:.*]] = fir.convert %[[IVB_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE2:.*]] = fir.shape %c5, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR2:.*]] = fir.array_coor {{.*}}(%[[SHAPE2]]) %[[IP1_I64]], %[[IVB_I64]] : (!fir.ref<!fir.array<5x6xi32>>, !fir.shape<2>, i64, i64) -> !fir.ref<i32>
+
+subroutine task_affinity_iterator_section_subscript()
+ integer, parameter :: n = 5, m = 6
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n, j = 1:m) : a(i:i+1, j+2)) shared(a)
+ a(1, 1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_section_subscript()
+! CHECK: %[[ITER3:.*]] = omp.iterator(%[[IVS0:.*]]: index, %[[IVS1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IVS0_I32:.*]] = fir.convert %[[IVS0]] : (index) -> i32
+! CHECK: %[[IVS1_I32:.*]] = fir.convert %[[IVS1]] : (index) -> i32
+! CHECK: %[[IVS0_I64:.*]] = fir.convert %[[IVS0_I32]] : (i32) -> i64
+! CHECK: %[[C2_I32:.*]] = arith.constant 2 : i32
+! CHECK: %[[JP2_I32:.*]] = arith.addi %[[IVS1_I32]], %[[C2_I32]] : i32
+! CHECK: %[[JP2_I64:.*]] = fir.convert %[[JP2_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE3:.*]] = fir.shape %c5, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR3:.*]] = fir.array_coor {{.*}}(%[[SHAPE3]]) %[[IVS0_I64]], %[[JP2_I64]] : (!fir.ref<!fir.array<5x6xi32>>, !fir.shape<2>, i64, i64) -> !fir.ref<i32>
+
+subroutine task_affinity_iterator_section_implicit_lower()
+ integer, parameter :: n = 5, m = 6
+ integer :: a(n, m)
+ integer :: i, j
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n, j = 1:m) : a(:i+1, j+2)) shared(a)
+ a(1, 1) = 1
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_section_implicit_lower()
+! CHECK: %[[ITER4:.*]] = omp.iterator(%[[IVT0:.*]]: index, %[[IVT1:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IVT1_I32:.*]] = fir.convert %[[IVT1]] : (index) -> i32
+! CHECK: %[[C1_IDX:.*]] = arith.constant 1 : index
+! CHECK: %[[C2_I32_2:.*]] = arith.constant 2 : i32
+! CHECK: %[[JP2_I32_2:.*]] = arith.addi %[[IVT1_I32]], %[[C2_I32_2]] : i32
+! CHECK: %[[JP2_I64_2:.*]] = fir.convert %[[JP2_I32_2]] : (i32) -> i64
+! CHECK: %[[SHAPE4:.*]] = fir.shape %c5, %c6 : (index, index) -> !fir.shape<2>
+! CHECK: %[[COOR4:.*]] = fir.array_coor {{.*}}(%[[SHAPE4]]) %[[C1_IDX]], %[[JP2_I64_2]] : (!fir.ref<!fir.array<5x6xi32>>, !fir.shape<2>, index, i64) -> !fir.ref<i32>
+
+subroutine task_affinity_iterator_char_simple()
+ integer, parameter :: n = 8
+ character(len=7) :: a(n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n) : a(i))
+ a(1) = "abcdefg"
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_char_simple()
+! CHECK: %[[ITER5:.*]] = omp.iterator(%[[IVC:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IVC_I32:.*]] = fir.convert %[[IVC]] : (index) -> i32
+! CHECK: %[[IVC_I64:.*]] = fir.convert %[[IVC_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE5:.*]] = fir.shape {{.*}} : (index) -> !fir.shape<1>
+! CHECK: %[[COOR5:.*]] = fir.array_coor {{.*}}(%[[SHAPE5]]) %[[IVC_I64]] : ({{.*}}, !fir.shape<1>, i64) -> !fir.ref<!fir.char<1,7>>
+! CHECK: %[[C1_I64:.*]] = arith.constant 1 : i64
+! CHECK: %[[C7_I64:.*]] = fir.convert %c7 : (index) -> i64
+! CHECK: %[[ELEM5:.*]] = arith.muli %[[C7_I64]], %[[C1_I64]] : i64
+! CHECK: %[[ADDR5:.*]] = fir.convert %[[COOR5]] : (!fir.ref<!fir.char<1,7>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY5:.*]] = omp.affinity_entry %[[ADDR5]], %[[ELEM5]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine task_affinity_iterator_char_expr_subscript()
+ integer, parameter :: n = 8
+ character(len=7) :: a(n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n-1) : a(i+1))
+ a(1) = "abcdefg"
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_char_expr_subscript()
+! CHECK: %[[ITER6:.*]] = omp.iterator(%[[IVC2:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[IVC2_I32:.*]] = fir.convert %[[IVC2]] : (index) -> i32
+! CHECK: %[[C1_I32_6:.*]] = arith.constant 1 : i32
+! CHECK: %[[IP1C_I32:.*]] = arith.addi %[[IVC2_I32]], %[[C1_I32_6]] : i32
+! CHECK: %[[IP1C_I64:.*]] = fir.convert %[[IP1C_I32]] : (i32) -> i64
+! CHECK: %[[SHAPE6:.*]] = fir.shape {{.*}} : (index) -> !fir.shape<1>
+! CHECK: %[[COOR6:.*]] = fir.array_coor {{.*}}(%[[SHAPE6]]) %[[IP1C_I64]] : ({{.*}}, !fir.shape<1>, i64) -> !fir.ref<!fir.char<1,7>>
+! CHECK: %[[C1_I64_2:.*]] = arith.constant 1 : i64
+! CHECK: %[[C7_I64_2:.*]] = fir.convert %c7 : (index) -> i64
+! CHECK: %[[ELEM6:.*]] = arith.muli %[[C7_I64_2]], %[[C1_I64_2]] : i64
+! CHECK: %[[ADDR6:.*]] = fir.convert %[[COOR6]] : (!fir.ref<!fir.char<1,7>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY6:.*]] = omp.affinity_entry %[[ADDR6]], %[[ELEM6]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
+
+subroutine task_affinity_iterator_char_runtime(n, l)
+ integer, intent(in) :: n, l
+ character(len=l) :: a(n)
+ integer :: i
+
+ !$omp parallel
+ !$omp single
+ !$omp task affinity(iterator(i = 1:n) : a(i))
+ a(1) = repeat("x", l)
+ !$omp end task
+ !$omp end single
+ !$omp end parallel
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtask_affinity_iterator_char_runtime(
+! CHECK: %[[A:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) typeparams %{{.*}} {uniq_name = "_QFtask_affinity_iterator_char_runtimeEa"}
+! CHECK: %[[ITER:.*]] = omp.iterator(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+! CHECK: %[[COOR:.*]] = fir.array_coor %[[A]]#0({{.*}}) {{.*}} : (!fir.box<!fir.array<?x!fir.char<1,?>>>, !fir.shape<1>, i64) -> !fir.ref<!fir.char<1,?>>
+! CHECK: %[[ELEM:.*]] = fir.box_elesize %[[A]]#0 : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
+! CHECK: %[[ELEM_I64:.*]] = fir.convert %[[ELEM]] : (index) -> i64
+! CHECK: %[[ADDR:.*]] = fir.convert %[[COOR]] : (!fir.ref<!fir.char<1,?>>) -> !fir.ref<i8>
+! CHECK: %[[ENTRY:.*]] = omp.affinity_entry %[[ADDR]], %[[ELEM_I64]] : (!fir.ref<i8>, i64) -> !omp.affinity_entry_ty<!fir.ref<i8>, i64>
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 20f14c0124969..23c2fbdfd7368 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_PointerLikeType>:$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/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/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 4b85d2bbf6db5..f388b7ccb591b 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1434,6 +1434,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,
@@ -3140,10 +3190,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,
@@ -4607,24 +4657,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);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index c02c134b008a0..b908874c2010b 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>) {
+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() {
- // 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_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
@@ -3588,3 +3600,57 @@ 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.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.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>)
+ } -> !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
+}
+
+// CHECK-LABEL: func.func @omp_task_affinity_iterator_2d
+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.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.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
+
+ %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
+}
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 1a1286cb30251..2500d546fcf4d 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -464,10 +464,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 flang-commits
mailing list