[Mlir-commits] [flang] [mlir] [mlir][flang][OpenMP] Select iterators per locator and allow multi-block regions (PR #227454)

[llvmlistbot at llvm.org]

================
@@ -1046,62 +1044,46 @@ bool ClauseProcessor::processAffinity(
               .getResult();
         };
 
-        llvm::SmallVector<IteratorRange> iteratorRanges;
-        llvm::SmallPtrSet<const Fortran::semantics::Symbol *, 4> ivSyms;
-
-        auto &iteratorModifier =
-            std::get<std::optional<omp::clause::Iterator>>(clause.t);
-        collectIteratorIVs(clause, converter, stmtCtx, iteratorRanges, ivSyms);
+        auto iteratorRanges = lowerIteratorRanges(clause, converter, stmtCtx);
 
         TodoLocators(clauseLocation, objects);
 
-        for (const omp::Object &object : objects) {
+        auto genEntry = [&](const omp::Object &object,
+                            lower::StatementContext &localStmtCtx) {
+          mlir::Value addr =
+              genAffinityAddr(converter, object, localStmtCtx, clauseLocation);
           llvm::SmallVector<mlir::Value> bounds;
           std::stringstream asFortran;
-          if (iteratorModifier.has_value() &&
-              hasIteratorIVReference(object, ivSyms)) {
+          fir::factory::AddrAndBoundsInfo info =
+              lower::gatherDataOperandAddrAndBounds<mlir::omp::MapBoundsOp,
+                                                    mlir::omp::MapBoundsType>(
+                  converter, builder, semaCtx, localStmtCtx, *object.sym(),
+                  object.ref(), clauseLocation, asFortran, bounds,
+                  treatIndexAsSection);
+          hlfir::Entity entity{info.addr};
+          mlir::Value len =
+              object.ref() && object.ref()->Rank() == 0
+                  ? genElementSizeInBytes(builder, clauseLocation,
+                                          builder.getDataLayout(), entity)
+                  : genAffinityLen(builder, clauseLocation,
+                                   builder.getDataLayout(), entity, bounds);
+          return makeAffinityEntry(builder, clauseLocation, entryTy, addr, len);
+        };
+
+        for (const omp::Object &object : objects) {
+          llvm::SmallVector<IteratorRange> objectRanges =
+              getIteratorRangesForObject(object, iteratorRanges);
+          if (!objectRanges.empty()) {
             mlir::Value iterHandle = buildIteratorOp(
-                converter, clauseLocation, iterTy, iteratorRanges,
+                converter, clauseLocation, iterTy, objectRanges,
                 [&](fir::FirOpBuilder &builder, mlir::Location loc,
                     llvm::ArrayRef<mlir::Value> /*ivs*/) -> mlir::Value {
                   lower::StatementContext iterStmtCtx;
-
-                  if (std::optional<llvm::SmallVector<mlir::Value>>
-                          loweredIndices = getIteratorElementIndices(
-                              converter, object, iterStmtCtx, loc)) {
-                    const Fortran::semantics::Symbol *sym = object.sym();
-                    assert(sym && "expected symbol for iterator object");
-                    fir::factory::AddrAndBoundsInfo info =
-                        Fortran::lower::getDataOperandBaseAddr(
-                            converter, builder, *sym, loc,
-                            /*unwrapFirBox=*/false);
-                    hlfir::Entity entity{info.addr};
-                    mlir::Value iteratedAddr = genIteratorCoordinate(
-                        converter, entity, *loweredIndices, loc);
-                    mlir::Value len = genElementSizeInBytes(
-                        builder, loc, builder.getDataLayout(), entity);
-                    return makeAffinityEntry(builder, loc, entryTy,
-                                             iteratedAddr, len);
-                  }
-
-                  TODO(loc, "object type not supported by iterator modifier");
+                  return genEntry(object, iterStmtCtx);
----------------
MattPD wrote:

Could the vector-subscript check cover earlier part references in the designator? In `a([i])%field(1)`, the final subscript is scalar, but `a([i])` has a vector subscript:

```fortran
subroutine s(a, m)
  type t
    integer :: field(2)
  end type
  type(t) :: a(8)
  integer :: m
  !$omp task affinity(iterator(i=1:m): a([i])%field(1))
  !$omp end task
end
```

The source passes semantic checking at this revision and at its [parent commit](https://github.com/llvm/llvm-project/commit/5655cabbd7b780c9b3bcfbbb07a619139990f06c). With `flang -c -O1 -fopenmp -fopenmp-version=52 repro.f90`, this revision aborts at the scalar-rank assertion in `genScalarAssignment` instead of reporting the vector subscript as not yet implemented. The parent commit fails earlier with `symbol is not mapped to any IR value`.

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