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

[Jason Van Beusekom 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;
----------------
Jason-Van-Beusekom wrote:

added comment in https://github.com/llvm/llvm-project/pull/212646/commits/7d7669f1398a9d2eb4ae5afa7ae784895fc2bbde

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