[flang-commits] [flang] [Flang] add support for collapsed io calls for implicit do loops (PR #212646)

[via flang-commits]

================
@@ -701,6 +701,244 @@ static mlir::func::FuncOp getOutputFunc(mlir::Location loc,
                                                                    builder);
 }
 
+/// Return the evaluate expression of a non-implied-do io item, or nullptr if
+/// the item is itself an io-implied-do.
+static const Fortran::lower::SomeExpr *
+getIoItemLeafExpr(const Fortran::parser::OutputItem &item) {
+  if (const auto *parserExpr = std::get_if<Fortran::parser::Expr>(&item.u))
+    return Fortran::semantics::GetExpr(*parserExpr);
+  return nullptr;
+}
+static const Fortran::lower::SomeExpr *
+getIoItemLeafExpr(const Fortran::parser::InputItem &item) {
+  if (const auto *var = std::get_if<Fortran::parser::Variable>(&item.u))
+    return Fortran::semantics::GetExpr(*var);
+  return nullptr;
+}
+
+template <typename A>
+static bool ioExprReferencesSymbol(const A &expr,
+                                   const Fortran::semantics::Symbol &sym) {
+  const Fortran::semantics::Symbol &ultimate = sym.GetUltimate();
+  for (const Fortran::semantics::SymbolRef &ref :
+       Fortran::evaluate::CollectSymbols(expr)) {
+    if (&ref->GetUltimate() == &ultimate)
+      return true;
+  }
+  return false;
+}
+
+template <typename A>
+static bool ioExprReferencesVolatile(const A &expr) {
+  for (const Fortran::semantics::SymbolRef &ref :
+       Fortran::evaluate::CollectSymbols(expr))
+    if (ref->GetUltimate().attrs().test(Fortran::semantics::Attr::VOLATILE))
+      return true;
+  return false;
+}
+
+namespace {
+struct CollapsedImpliedDo {
+  Fortran::lower::SomeExpr section;
+  const Fortran::semantics::Symbol *loopSym;
+  Fortran::lower::SomeExpr loopLowerValue;
+  Fortran::lower::SomeExpr loopStepValue;
+};
+} // namespace
+
+/// Try to recognize a (single-level) io-implied-do as an equivalent contiguous
+/// array section. Returns the section expression on success, std::nullopt
+/// otherwise
+template <typename ImpliedDo>
+static std::optional<CollapsedImpliedDo>
+matchContiguousImpliedDo(Fortran::lower::AbstractConverter &converter,
+                         const ImpliedDo &impliedDo, bool isInput) {
+
+  // Only collapse a single body item
+  const auto &items = std::get<0>(impliedDo.t);
+  if (items.size() != 1)
+    return std::nullopt;
+
+  const Fortran::lower::SomeExpr *leaf = getIoItemLeafExpr(items.front());
+  if (!leaf)
+    return std::nullopt;
+
+  const Fortran::parser::IoImpliedDoControl &control = std::get<1>(impliedDo.t);
+  const Fortran::semantics::Symbol *loopSym =
+      Fortran::parser::UnwrapRef<Fortran::parser::Name>(control.Name()).symbol;
+  if (!loopSym)
+    return std::nullopt;
+
+  std::optional<Fortran::evaluate::DataRef> dataRef =
+      Fortran::evaluate::ExtractDataRef(*leaf);
+  if (!dataRef)
+    return std::nullopt;
+
+  const auto *arrayRef = std::get_if<Fortran::evaluate::ArrayRef>(&dataRef->u);
+  if (!arrayRef || !arrayRef->base().IsSymbol())
+    return std::nullopt;
+
+  const Fortran::lower::SomeExpr *lowerExpr =
+      Fortran::semantics::GetExpr(control.Lower());
+  const Fortran::lower::SomeExpr *upperExpr =
+      Fortran::semantics::GetExpr(control.Upper());
+  const Fortran::lower::SomeExpr *stepExpr =
+      control.Step() ? Fortran::semantics::GetExpr(*control.Step()) : nullptr;
+  if (!lowerExpr || !upperExpr || (control.Step() && !stepExpr))
+    return std::nullopt;
+
+  if (ioExprReferencesSymbol(*lowerExpr, *loopSym) ||
+      ioExprReferencesSymbol(*upperExpr, *loopSym) ||
+      (stepExpr && ioExprReferencesSymbol(*stepExpr, *loopSym)))
+    return std::nullopt;
+
+  Fortran::evaluate::FoldingContext &foldingContext =
+      converter.getFoldingContext();
+  if (Fortran::evaluate::FindImpureCall(foldingContext, *lowerExpr) ||
+      Fortran::evaluate::FindImpureCall(foldingContext, *upperExpr) ||
+      (stepExpr &&
+       Fortran::evaluate::FindImpureCall(foldingContext, *stepExpr)))
+    return std::nullopt;
+
+  if (ioExprReferencesVolatile(*lowerExpr) ||
+      ioExprReferencesVolatile(*upperExpr) ||
+      (stepExpr && ioExprReferencesVolatile(*stepExpr)))
+    return std::nullopt;
----------------
jeanPerier wrote:

This is an implementation limitation because the code here is re-evaluating the lower and step to compute the final value, right?
Not asking you to lift it, it is OK to have implementation limitations in the code, please just explain the rational for the restriction in comment.

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