[flang-commits] [flang] [mlir] [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (PR #182222)

[Tom Eccles via flang-commits]

================
@@ -756,14 +848,83 @@ 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 makeAffinityEntry = [&](fir::FirOpBuilder &b, mlir::Location l,
+                                     mlir::Type entryTy, mlir::Value addr,
+                                     mlir::Value len) -> mlir::Value {
+          mlir::Value addrI8 = fir::ConvertOp::create(b, l, refI8Ty, addr);
+          return mlir::omp::AffinityEntryOp::create(b, l, entryTy, addrI8, len)
+              .getResult();
+        };
+
+        llvm::SmallVector<IteratorRange> iteratorRanges;
+        llvm::SmallPtrSet<const Fortran::semantics::Symbol *, 4> ivSyms;
+
+        // 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<Fortran::evaluate::SomeType>(
+                    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);
+                  hlfir::Entity entity{info.addr};
+                  mlir::Value addr =
+                      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);
+                });
+            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(),
----------------
tblah wrote:

Is it an error for genAffinityLen to return 0 here? My understanding is that this works so far because we will take the other path where bounds are available

https://github.com/llvm/llvm-project/pull/182222