[llvm-branch-commits] [flang] [mlir] [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (PR #182219)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Wed Feb 18 21:13:04 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-openmp
@llvm/pr-subscribers-mlir
Author: Chi-Chun, Chen (chichunchen)
<details>
<summary>Changes</summary>
This patch translates flang ast to OpenMP dialect for affinity clause
including the iterator modifier.
---
Patch is 43.90 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/182219.diff
11 Files Affected:
- (modified) flang/lib/Lower/OpenMP/ClauseProcessor.cpp (+191-6)
- (modified) flang/lib/Lower/OpenMP/Utils.cpp (+186)
- (modified) flang/lib/Lower/OpenMP/Utils.h (+25)
- (removed) flang/test/Lower/OpenMP/Todo/affinity-clause.f90 (-10)
- (modified) flang/test/Lower/OpenMP/task-affinity.f90 (+104-48)
- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td (+4-2)
- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td (+8)
- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+17)
- (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+70-18)
- (modified) mlir/test/Dialect/OpenMP/ops.mlir (+84-19)
- (modified) mlir/test/Target/LLVMIR/openmp-todo.mlir (+6-4)
``````````diff
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index e62395676a696..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,14 +842,113 @@ 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();
+ 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(&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);
+ return mlir::omp::AffinityEntryOp::create(b, l, entryTy, addrI8, len)
+ .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());
}
- const auto &objects = std::get<omp::ObjectList>(clause.t);
- if (!objects.empty())
- genObjectList(objects, converter, result.affinityVars);
+ for (const omp::Object &object : objects) {
+ llvm::SmallVector<mlir::Value> bounds;
+ std::stringstream asFortran;
+ 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 e9ba5f386803a..83d1a977cb3e3 100644
--- a/flang/lib/Lower/OpenMP/Utils.cpp
+++ b/flang/lib/Lower/OpenMP/Utils.cpp
@@ -917,6 +917,192 @@ 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);
+}
+
+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 f707557197847..10552a68908d8 100644
--- a/flang/lib/Lower/OpenMP/Utils.h
+++ b/flang/lib/Lower/OpenMP/Utils.h
@@ -189,6 +189,31 @@ 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);
+
+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/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..3e3f8acad7c62 100644
--- a/flang/test/Lower/OpenMP/task-affinity.f90
+++ b/flang/test/Lower/OpenMP/task-affinity.f90
@@ -7,26 +7,30 @@ 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: %[[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.ar...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/182219
More information about the llvm-branch-commits
mailing list