Lower IO input with vector subscripts

Define a VectorSubscriptBox class that allow representing and working
with a lowered Designator containing vector subscripts while ensuring
all the subscripts expression are only lowered once.

It's a bit of a super ExtendedValue (but is is not added as such because
it is heavy and its use case is only restricted to IO input so far).

The key point of this class is that it has members functions
`loopOverElements` and `loopOverElementsWhile` that allow creating loops
over each element of the designator, and call a provided callback with
the address of the element.

This class is used in input IO to create an IO runtime call for each
element of the designator, since it is not possible to build a
descriptor for it. The `loopOverElementsWhile` version is required when
error recovery is enabled in IO.

A hidden VectorSubscriptBoxBuilder is in charge of making a custom
Designator<T> visit and create the VectorSubscriptBox. Once this is
done, the VectorSubscriptBox does not equire any front-end data
structures to work with.

The motivation for creating such tool is that the current lowering
array lowering infrastructure is centered around assignment (to a
variable or temporary), and could not be used here to created the loops
over the IO runtime calls without some non trivial modification to it.
Adding complexity to the array expression lowering framework to cover
a corner case did not appear a good idea.

The option of creating a temp, passing it to the runtime, and copying it
back was explored, but it was not possible to guarantee that the
subscript would be evaluated only once given there was no way to "keep
the lowered representation" of the designator between the temp creation
and the copy back (the temp creation requires evaluating the susbcripts
to compute the shape in general). However, note that with the added
VectorSubscriptBox, it would actually now also be possible to deploy a
temp + copy back mechanism. This can be done later if it appeared
beneficial in real world program.

The only cases left TODOs are the cases when one of the Components
is a parent type reference (not yet handled properly in the general case),
the coarray case, and the PDT case (not sure exactly how type params
will be threaded in fir.field_index). All other cases are covered, and
I tried to add exhaustive tests since I do not expect real world program
to be very harsh on this utility (most will just do READ(*, *) x(y)).

Author: jeanPerier

flang/include/flang/Lower/Support/Utils.h

@@ -45,4 +45,25 @@ toEvExpr(const A &x) {
   return Fortran::evaluate::AsGenericExpr(Fortran::common::Clone(x));
 }
 
+template <Fortran::common::TypeCategory FROM>
+static Fortran::evaluate::Expr<Fortran::evaluate::SomeType> ignoreEvConvert(
+    const Fortran::evaluate::Convert<
+        Fortran::evaluate::Type<Fortran::common::TypeCategory::Integer, 8>,
+        FROM> &x) {
+  return toEvExpr(x.left());
+}
+template <typename A>
+static Fortran::evaluate::Expr<Fortran::evaluate::SomeType>
+ignoreEvConvert(const A &x) {
+  return toEvExpr(x);
+}
+/// A vector subscript expression may be wrapped with a cast to INTEGER*8.
+/// Get rid of it here so the vector can be loaded. Add it back when
+/// generating the elemental evaluation (inside the loop nest).
+inline Fortran::evaluate::Expr<Fortran::evaluate::SomeType>
+ignoreEvConvert(const Fortran::evaluate::Expr<Fortran::evaluate::Type<
+                    Fortran::common::TypeCategory::Integer, 8>> &x) {
+  return std::visit([](const auto &v) { return ignoreEvConvert(v); }, x.u);
+}
+
 #endif // FORTRAN_LOWER_SUPPORT_UTILS_H

flang/include/flang/Lower/VectorSubscripts.h

@@ -0,0 +1,152 @@
+//===-- VectorSubscripts.h -- vector subscripts tools -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// \brief Defines a compiler internal representation for lowered designators
+///  containing vector subscripts. This representation allows working on such
+///  designators in custom ways while ensuring the designator subscripts are
+///  only evaluated once. It is mainly intended for cases that do not fit in
+///  the array expression lowering framework like input IO in presence of
+///  vector subscripts.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_LOWER_VECTORSUBSCRIPTS_H
+#define FORTRAN_LOWER_VECTORSUBSCRIPTS_H
+
+#include "flang/Optimizer/Builder/BoxValue.h"
+
+namespace fir {
+class FirOpBuilder;
+}
+
+namespace Fortran {
+
+namespace evaluate {
+template <typename>
+class Expr;
+struct SomeType;
+} // namespace evaluate
+
+namespace lower {
+
+class AbstractConverter;
+class StatementContext;
+
+/// VectorSubscriptBox is a lowered representation for any Designator<T> that
+/// contain at least one vector subscript.
+///
+/// A designator `x%a(i,j)%b(1:foo():1, vector, k)%c%d(m)%e1
+/// Is lowered into:
+///   - an ExtendedValue for ranked base (x%a(i,j)%b)
+///   - mlir:Values and ExtendedValues for the triplet, vector subscript and
+///     scalar subscripts of the ranked array reference (1:foo():1, vector, k)
+///   - a list of fir.field_index and scalar integers mlir::Value for the
+///   component
+///     path at the right of the ranked array ref (%c%d(m)%e).
+///
+/// This representation allows later creating loops over the designator elements
+/// and fir.array_coor to get the element addresses without re-evaluating any
+/// sub-expressions.
+class VectorSubscriptBox {
+public:
+  /// Type of the callbacks that can be passed to work with the element
+  /// addresses.
+  using ElementalGenerator = std::function<void(const fir::ExtendedValue &)>;
+  using ElementalGeneratorWithBoolReturn =
+      std::function<mlir::Value(const fir::ExtendedValue &)>;
+  struct LoweredVectorSubscript {
+    fir::ExtendedValue vector;
+    // Vector size, guaranteed to be of indexType.
+    mlir::Value size;
+  };
+  struct LoweredTriplet {
+    // Triplets value, guaranteed to be of indexType.
+    mlir::Value lb;
+    mlir::Value ub;
+    mlir::Value stride;
+  };
+  using LoweredSubscript =
+      std::variant<mlir::Value, LoweredTriplet, LoweredVectorSubscript>;
+  using MaybeSubstring = llvm::SmallVector<mlir::Value, 2>;
+  VectorSubscriptBox(
+      fir::ExtendedValue &&loweredBase,
+      llvm::SmallVector<LoweredSubscript, 16> &&loweredSubscripts,
+      llvm::SmallVector<mlir::Value> &&componentPath,
+      MaybeSubstring substringBounds, mlir::Type elementType)
+      : loweredBase{std::move(loweredBase)}, loweredSubscripts{std::move(
+                                                 loweredSubscripts)},
+        componentPath{std::move(componentPath)},
+        substringBounds{substringBounds}, elementType{elementType} {};
+
+  /// Loop over the elements described by the VectorSubscriptBox, and call
+  /// \p elementalGenerator inside the loops with the element addresses.
+  void loopOverElements(fir::FirOpBuilder &builder, mlir::Location loc,
+                        const ElementalGenerator &elementalGenerator);
+
+  /// Loop over the elements described by the VectorSubscriptBox while a
+  /// condition is true, and call \p elementalGenerator inside the loops with
+  /// the element addresses. The initial condition value is \p initialCondition,
+  /// and then it is the result of \p elementalGenerator. The value of the
+  /// condition after the loops is returned.
+  mlir::Value loopOverElementsWhile(
+      fir::FirOpBuilder &builder, mlir::Location loc,
+      const ElementalGeneratorWithBoolReturn &elementalGenerator,
+      mlir::Value initialCondition);
+
+  /// Return the type of the elements of the array section.
+  mlir::Type getElementType() { return elementType; }
+
+private:
+  /// Common implementation for DoLoop and IterWhile loop creations.
+  template <typename LoopType, typename Generator>
+  mlir::Value loopOverElementsBase(fir::FirOpBuilder &builder,
+                                   mlir::Location loc,
+                                   const Generator &elementalGenerator,
+                                   mlir::Value initialCondition);
+  /// Create sliceOp for the designator.
+  mlir::Value createSlice(fir::FirOpBuilder &builder, mlir::Location loc);
+
+  /// Create ExtendedValue the element inside the loop.
+  fir::ExtendedValue getElementAt(fir::FirOpBuilder &builder,
+                                  mlir::Location loc, mlir::Value shape,
+                                  mlir::Value slice,
+                                  mlir::ValueRange inductionVariables);
+
+  /// Generate the [lb, ub, step] to loop over the section (in loop order, not
+  /// Fortran dimension order).
+  llvm::SmallVector<std::tuple<mlir::Value, mlir::Value, mlir::Value>>
+  genLoopBounds(fir::FirOpBuilder &builder, mlir::Location loc);
+
+  /// Lowered base of the ranked array ref.
+  fir::ExtendedValue loweredBase;
+  /// Subscripts values of the rank arrayRef part.
+  llvm::SmallVector<LoweredSubscript, 16> loweredSubscripts;
+  /// Scalar subscripts and components at the right of the ranked
+  /// array ref part of any.
+  llvm::SmallVector<mlir::Value> componentPath;
+  /// List of substring bounds if this is a substring (only the lower bound if
+  /// the upper is implicit).
+  MaybeSubstring substringBounds;
+  /// Type of the elements described by this array section.
+  mlir::Type elementType;
+};
+
+/// Lower \p expr, that must be an designator containing vector subscripts, to a
+/// VectorSubscriptBox representation. This causes evaluation of all the
+/// subscripts. Any required clean-ups from subscript expression are added to \p
+/// stmtCtx.
+VectorSubscriptBox genVectorSubscriptBox(
+    mlir::Location loc, Fortran::lower::AbstractConverter &converter,
+    Fortran::lower::StatementContext &stmtCtx,
+    const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &expr);
+
+} // namespace lower
+} // namespace Fortran
+
+#endif // FORTRAN_LOWER_VECTORSUBSCRIPTS_H

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -336,7 +336,7 @@ class FirOpBuilder : public mlir::OpBuilder {
 
   /// Generate code testing \p addr is not a null address.
   mlir::Value genIsNotNull(mlir::Location loc, mlir::Value addr);
-   
+
   /// Generate code testing \p addr is a null address.
   mlir::Value genIsNull(mlir::Location loc, mlir::Value addr);
 
@@ -428,12 +428,37 @@ fir::ExtendedValue componentToExtendedValue(fir::FirOpBuilder &builder,
                                             mlir::Location loc,
                                             mlir::Value component);
 
+/// Given the address of an array element and the ExtendedValue describing the
+/// array, returns the ExtendedValue describing the array element. The purpose
+/// is to propagate the length parameters of the array to the element.
+/// This can be used for elements of `array` or `array(i:j:k)`. If \p element
+/// belongs to an array section `array%x` whose base is \p array,
+/// arraySectionElementToExtendedValue must be used instead.
+fir::ExtendedValue arrayElementToExtendedValue(fir::FirOpBuilder &builder,
+                                               mlir::Location loc,
+                                               const fir::ExtendedValue &array,
+                                               mlir::Value element);
+
+/// Build the ExtendedValue for \p element that is an element of an array or
+/// array section with \p array base (`array` or `array(i:j:k)%x%y`).
+/// If it is an array section, \p slice must be provided and be a fir::SliceOp
+/// that describes the section.
+fir::ExtendedValue arraySectionElementToExtendedValue(
+    fir::FirOpBuilder &builder, mlir::Location loc,
+    const fir::ExtendedValue &array, mlir::Value element, mlir::Value slice);
+
 /// Assign \p rhs to \p lhs. Both \p rhs and \p lhs must be scalar derived
 /// types. The assignment follows Fortran intrinsic assignment semantic for
 /// derived types (10.2.1.3 point 13).
 void genRecordAssignment(fir::FirOpBuilder &builder, mlir::Location loc,
                          const fir::ExtendedValue &lhs,
                          const fir::ExtendedValue &rhs);
+
+/// Compute the extent of (lb:ub:step) as max((ub-lb+step)/step, 0). See Fortran
+/// 2018 9.5.3.3.2 section for more details.
+mlir::Value computeTripletExtent(fir::FirOpBuilder &builder, mlir::Location loc,
+                                 mlir::Value lb, mlir::Value ub,
+                                 mlir::Value step, mlir::Type type);
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

flang/lib/Lower/CMakeLists.txt

@@ -19,6 +19,7 @@ add_flang_library(FortranLower
   PFTBuilder.cpp
   Runtime.cpp
   SymbolMap.cpp
+  VectorSubscripts.cpp
 
   DEPENDS
   FIRBuilder

flang/lib/Lower/ConvertExpr.cpp

@@ -207,54 +207,6 @@ static bool isAllocatableOrPointer(const Fortran::lower::SomeExpr &expr) {
   return sym && Fortran::semantics::IsAllocatableOrPointer(*sym);
 }
 
-/// Given the address of an array element and the ExtendedValue describing the
-/// array, returns the ExtendedValue describing the array element. The purpose
-/// is to propagate the length parameters of the array to the element.
-/// This can be used for elements of `array` or `array(i:j:k)`. If \p element
-/// belongs to an array section `array%x` whose base is \p array,
-/// arraySectionElementToExtendedValue must be used instead.
-static fir::ExtendedValue
-arrayElementToExtendedValue(fir::FirOpBuilder &builder, mlir::Location loc,
-                            const fir::ExtendedValue &array,
-                            mlir::Value element) {
-  return array.match(
-      [&](const fir::CharBoxValue &cb) -> fir::ExtendedValue {
-        return cb.clone(element);
-      },
-      [&](const fir::CharArrayBoxValue &bv) -> fir::ExtendedValue {
-        return bv.cloneElement(element);
-      },
-      [&](const fir::BoxValue &box) -> fir::ExtendedValue {
-        if (box.isCharacter()) {
-          auto len = fir::factory::readCharLen(builder, loc, box);
-          return fir::CharBoxValue{element, len};
-        }
-        if (box.isDerivedWithLengthParameters())
-          TODO(loc, "get length parameters from derived type BoxValue");
-        return element;
-      },
-      [&](const auto &) -> fir::ExtendedValue { return element; });
-}
-
-/// Build the ExtendedValue for \p element that is an element of an array or
-/// array section with \p array base (`array` or `array(i:j:k)%x%y`).
-/// If it is an array section, \p slice must be provided and be a fir::SliceOp
-/// that describes the section.
-static fir::ExtendedValue arraySectionElementToExtendedValue(
-    fir::FirOpBuilder &builder, mlir::Location loc,
-    const fir::ExtendedValue &array, mlir::Value element, mlir::Value slice) {
-  if (!slice)
-    return arrayElementToExtendedValue(builder, loc, array, element);
-  auto sliceOp = mlir::dyn_cast_or_null<fir::SliceOp>(slice.getDefiningOp());
-  assert(sliceOp && "slice must be a sliceOp");
-  if (sliceOp.fields().empty())
-    return arrayElementToExtendedValue(builder, loc, array, element);
-  // For F95, using componentToExtendedValue will work, but when PDTs are
-  // lowered. It will be required to go down the slice to propagate the length
-  // parameters.
-  return fir::factory::componentToExtendedValue(builder, loc, element);
-}
-
 /// Convert the array_load, `load`, to an extended value. If `path` is not
 /// empty, then traverse through the components designated. The base value is
 /// `newBase`. This does not accept an array_load with a slice operand.
@@ -1442,7 +1394,7 @@ class ScalarExprLowering {
     assert(args.size() == seqTy.getDimension());
     auto ty = builder.getRefType(seqTy.getEleTy());
     auto addr = builder.create<fir::CoordinateOp>(loc, ty, base, args);
-    return arrayElementToExtendedValue(builder, loc, array, addr);
+    return fir::factory::arrayElementToExtendedValue(builder, loc, array, addr);
   }
 
   /// Lower an ArrayRef to a fir.coordinate_of using an element offset instead
@@ -1552,7 +1504,8 @@ class ScalarExprLowering {
     auto elementAddr = builder.create<fir::ArrayCoorOp>(
         loc, refTy, addr, shape, /*slice=*/mlir::Value{}, arrayCoorArgs,
         fir::getTypeParams(exv));
-    return arrayElementToExtendedValue(builder, loc, exv, elementAddr);
+    return fir::factory::arrayElementToExtendedValue(builder, loc, exv,
+                                                     elementAddr);
   }
 
   /// Return the coordinate of the array reference.
@@ -4842,8 +4795,8 @@ class ArrayExprLowering {
         mlir::Value coor = builder.create<fir::ArrayCoorOp>(
             loc, refEleTy, memref, shape, slice, indices,
             fir::getTypeParams(extMemref));
-        return arraySectionElementToExtendedValue(builder, loc, extMemref, coor,
-                                                  slice);
+        return fir::factory::arraySectionElementToExtendedValue(
+            builder, loc, extMemref, coor, slice);
       };
     }
     auto arrLoad = builder.create<fir::ArrayLoadOp>(
@@ -4893,8 +4846,8 @@ class ArrayExprLowering {
             llvm::ArrayRef<mlir::NamedAttribute>{
                 Fortran::lower::getAdaptToByRefAttr(builder)});
         builder.create<fir::StoreOp>(loc, base, temp);
-        return arraySectionElementToExtendedValue(builder, loc, extMemref, temp,
-                                                  slice);
+        return fir::factory::arraySectionElementToExtendedValue(
+            builder, loc, extMemref, temp, slice);
       };
     }
     // In the default case, the array reference forwards an `array_fetch` Op
@@ -4903,8 +4856,8 @@ class ArrayExprLowering {
       auto arrFetch = builder.create<fir::ArrayFetchOp>(
           loc, adjustedArraySubtype(arrTy, iters.iterVec()), arrLd,
           iters.iterVec(), arrLdTypeParams);
-      return arraySectionElementToExtendedValue(builder, loc, extMemref,
-                                                arrFetch, slice);
+      return fir::factory::arraySectionElementToExtendedValue(
+          builder, loc, extMemref, arrFetch, slice);
     };
   }