Merge pull request #1057 from flang-compiler/jpr-vector-io-2

Lower IO input with vector subscripts

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);
     };
   }