[flang-commits] [flang] [mlir] [Flang][OpenMP] Support iterator modifier in map and motion clauses (PR #197757)

[via flang-commits]

================
@@ -1283,6 +1283,189 @@ mlir::Value genIteratorCoordinate(Fortran::lower::AbstractConverter &converter,
                                   /*typeparams=*/mlir::ValueRange{});
 }
 
+// Lower the stable base entity for an iterator map/motion locator.
+//
+// MapInfoOp stores this base in var_ptr and represents the iterator-selected
+// elements or sections with map bounds. For component array locators such as
+// v%a(i) or v%a(i:i+1), return v%a so each iteration maps offsets from the
+// same base.
+static std::optional<hlfir::Entity>
+getIteratorMapEntity(Fortran::lower::AbstractConverter &converter,
+                     fir::FirOpBuilder &builder,
+                     Fortran::semantics::SemanticsContext &semaCtx,
+                     Fortran::lower::StatementContext &stmtCtx,
+                     const omp::Object &object, mlir::Location loc) {
+  const semantics::Symbol *sym = object.sym();
+  assert(sym && "expected symbol for iterator object");
+
+  if (!sym->owner().IsDerivedType()) {
+    // Keep the same base address as ordinary map lowering for non-component
+    // objects.
+    fir::factory::AddrAndBoundsInfo info =
+        Fortran::lower::getDataOperandBaseAddr(converter, builder, *sym, loc,
+                                               /*unwrapFirBox=*/false);
+    return hlfir::Entity{info.addr};
+  }
+
+  const std::optional<ExprTy> &ref = object.ref();
+  if (!ref)
+    return std::nullopt;
+
+  auto arrayRef = Fortran::lower::detail::getRef<evaluate::ArrayRef>(*ref);
+  if (!arrayRef)
+    return std::nullopt;
+
+  evaluate::ExpressionAnalyzer ea{semaCtx};
+  std::optional<ExprTy> arrayBase;
+  const evaluate::NamedEntity &base = arrayRef->base();
+  // Component array references carry the array base separately from the
+  // subscript list. Lower that base, e.g. v%a in v%a(i), and leave the
+  // subscript-dependent part to map bounds.
+  if (const semantics::SymbolRef *symRef = base.UnwrapSymbolRef())
+    arrayBase = ea.Designate(evaluate::DataRef{*symRef});
+  else if (const evaluate::Component *component = base.UnwrapComponent())
+    arrayBase = ea.Designate(evaluate::DataRef{*component});
+  else
+    llvm_unreachable("unexpected NamedEntity");
+
+  assert(arrayBase);
+  // Preserve mutable-box lowering for allocatable and pointer bases; MapInfoOp
+  // still records the address returned by the lowered base expression.
+  fir::ExtendedValue dataExv;
+  if (semantics::IsAllocatableOrPointer(base.GetLastSymbol()))
+    dataExv = converter.genExprMutableBox(loc, *arrayBase);
+  else
+    dataExv = converter.genExprAddr(loc, *arrayBase, stmtCtx);
+  return hlfir::Entity{fir::getBase(dataExv)};
+}
+
+// Build normalized map bounds for an iterator-dependent map/motion object.
+// MapInfoOp keeps the array base address as var_ptr; these bounds describe the
+// selected element or contiguous array section for each dimension.
+//
+// Examples:
+//   a(i)     -> lower_bound == upper_bound == i - base lower bound
+//   a(i:i+1) -> lower_bound == i - base lower bound,
+//               upper_bound == i + 1 - base lower bound
+static std::optional<llvm::SmallVector<mlir::Value>>
+genIteratorMapBounds(Fortran::lower::AbstractConverter &converter,
+                     hlfir::Entity entity, const omp::Object &object,
+                     Fortran::lower::StatementContext &stmtCtx,
+                     mlir::Location loc) {
+  const std::optional<ExprTy> &ref = object.ref();
+  assert(ref && "expected iterator-dependent object to have a reference");
+
+  std::optional<Fortran::evaluate::DataRef> dataRef =
+      Fortran::evaluate::ExtractDataRef(*ref);
+  if (!dataRef)
+    return std::nullopt;
+  const auto *arrayRef = std::get_if<Fortran::evaluate::ArrayRef>(&dataRef->u);
+  if (!arrayRef || arrayRef->subscript().empty())
+    return std::nullopt;
+
+  auto &builder = converter.getFirOpBuilder();
+  using SubscriptExpr =
+      Fortran::evaluate::Expr<Fortran::evaluate::SubscriptInteger>;
+  mlir::Type idxTy = builder.getIndexType();
+  mlir::Type boundTy = builder.getType<mlir::omp::MapBoundsType>();
+  mlir::Value zero = builder.createIntegerConstant(loc, idxTy, 0);
+  mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+
+  auto convertToIndex = [&](mlir::Value value) -> mlir::Value {
+    if (value.getType().isIndex())
+      return value;
+    return fir::ConvertOp::create(builder, loc, idxTy, value);
+  };
+
+  auto lowerSubscriptToIndex = [&](const SubscriptExpr &expr) -> mlir::Value {
+    mlir::Value value = fir::getBase(createSomeExtendedExpression(
+        loc, converter, toEvExpr(expr), converter.getSymbolMap(), stmtCtx));
+    return convertToIndex(value);
+  };
+
+  llvm::SmallVector<mlir::Value> bounds;
+  bounds.reserve(arrayRef->subscript().size());
+  // Translate each Fortran subscript into an OpenMP map bound for the
+  // corresponding array dimension.
+  for (const auto &[dim, subscript] : llvm::enumerate(arrayRef->subscript())) {
----------------
chichunchen wrote:

Done.

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