https://github.com/chichunchen updated https://github.com/llvm/llvm-project/pull/182222
>From 1eaded7e0dac63055a6d9635bbf47e31155cb929 Mon Sep 17 00:00:00 2001 From: cchen <[email protected]> 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 <[email protected]> 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 <[email protected]> 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 <[email protected]> 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@below {{not yet implemented: Unhandled clause affinity in omp.task operation}} - // expected-error@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@below {{not yet implemented: omp.affinity_entry}} + // expected-error@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 <[email protected]> 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 <[email protected]> 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 _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
