[flang] Lower IO input with vector subscripts

This patch adds lowering for IO input with vector subscripts.
It defines 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.

This patch is part of the upstreaming effort from fir-dev branch.

Reviewed By: PeteSteinfeld

Differential Revision: https://reviews.llvm.org/D121806

Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/include/flang/Lower/AbstractConverter.h

@@ -159,6 +159,11 @@ class AbstractConverter {
   /// Generate the type from a Variable
   virtual mlir::Type genType(const pft::Variable &) = 0;
 
+  /// Register a runtime derived type information object symbol to ensure its
+  /// object will be generated as a global.
+  virtual void registerRuntimeTypeInfo(mlir::Location loc,
+                                       SymbolRef typeInfoSym) = 0;
+
   //===--------------------------------------------------------------------===//
   // Locations
   //===--------------------------------------------------------------------===//

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

@@ -57,4 +57,25 @@ static Fortran::lower::SomeExpr toEvExpr(const A &x) {
   return Fortran::evaluate::AsGenericExpr(Fortran::common::Clone(x));
 }
 
+template <Fortran::common::TypeCategory FROM>
+static Fortran::lower::SomeExpr ignoreEvConvert(
+    const Fortran::evaluate::Convert<
+        Fortran::evaluate::Type<Fortran::common::TypeCategory::Integer, 8>,
+        FROM> &x) {
+  return toEvExpr(x.left());
+}
+template <typename A>
+static Fortran::lower::SomeExpr 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::lower::SomeExpr
+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,154 @@
+//===-- 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 {
+    LoweredVectorSubscript(fir::ExtendedValue &&vector, mlir::Value size)
+        : vector{std::move(vector)}, size{size} {}
+    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/Transforms/Passes.h

@@ -38,6 +38,7 @@ std::unique_ptr<mlir::Pass> createMemoryAllocationPass();
 std::unique_ptr<mlir::Pass>
 createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize);
 std::unique_ptr<mlir::Pass> createAnnotateConstantOperandsPass();
+std::unique_ptr<mlir::Pass> createSimplifyRegionLitePass();
 
 // declarative passes
 #define GEN_PASS_REGISTRATION

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -188,4 +188,12 @@ def MemoryAllocationOpt : Pass<"memory-allocation-opt", "mlir::FuncOp"> {
   let constructor = "::fir::createMemoryAllocationPass()";
 }
 
+def SimplifyRegionLite : Pass<"simplify-region-lite", "mlir::ModuleOp"> {
+  let summary = "Region simplification";
+  let description = [{
+    Run region DCE and erase unreachable blocks in regions.
+  }];
+  let constructor = "::fir::createSimplifyRegionLitePass()";
+}
+
 #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES

flang/include/flang/Tools/CLOptions.inc

@@ -143,6 +143,7 @@ inline void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm) {
   fir::addAVC(pm);
   pm.addNestedPass<mlir::FuncOp>(fir::createCharacterConversionPass());
   pm.addPass(mlir::createCanonicalizerPass(config));
+  pm.addPass(fir::createSimplifyRegionLitePass());
   fir::addMemoryAllocationOpt(pm);
 
   // The default inliner pass adds the canonicalizer pass with the default
@@ -157,6 +158,7 @@ inline void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm) {
   pm.addPass(mlir::createConvertSCFToCFPass());
 
   pm.addPass(mlir::createCanonicalizerPass(config));
+  pm.addPass(fir::createSimplifyRegionLitePass());
 }
 
 #if !defined(FLANG_EXCLUDE_CODEGEN)

flang/lib/Lower/Bridge.cpp

@@ -49,6 +49,68 @@ static llvm::cl::opt<bool> dumpBeforeFir(
     "fdebug-dump-pre-fir", llvm::cl::init(false),
     llvm::cl::desc("dump the Pre-FIR tree prior to FIR generation"));
 
+namespace {
+/// Helper class to generate the runtime type info global data. This data
+/// is required to describe the derived type to the runtime so that it can
+/// operate over it. It must be ensured this data will be generated for every
+/// derived type lowered in the current translated unit. However, this data
+/// cannot be generated before FuncOp have been created for functions since the
+/// initializers may take their address (e.g for type bound procedures). This
+/// class allows registering all the required runtime type info while it is not
+/// possible to create globals, and to generate this data after function
+/// lowering.
+class RuntimeTypeInfoConverter {
+  /// Store the location and symbols of derived type info to be generated.
+  /// The location of the derived type instantiation is also stored because
+  /// runtime type descriptor symbol are compiler generated and cannot be mapped
+  /// to user code on their own.
+  struct TypeInfoSymbol {
+    Fortran::semantics::SymbolRef symbol;
+    mlir::Location loc;
+  };
+
+public:
+  void registerTypeInfoSymbol(Fortran::lower::AbstractConverter &converter,
+                              mlir::Location loc,
+                              Fortran::semantics::SymbolRef typeInfoSym) {
+    if (seen.contains(typeInfoSym))
+      return;
+    seen.insert(typeInfoSym);
+    if (!skipRegistration) {
+      registeredTypeInfoSymbols.emplace_back(TypeInfoSymbol{typeInfoSym, loc});
+      return;
+    }
+    // Once the registration is closed, symbols cannot be added to the
+    // registeredTypeInfoSymbols list because it may be iterated over.
+    // However, after registration is closed, it is safe to directly generate
+    // the globals because all FuncOps whose addresses may be required by the
+    // initializers have been generated.
+    Fortran::lower::createRuntimeTypeInfoGlobal(converter, loc,
+                                                typeInfoSym.get());
+  }
+
+  void createTypeInfoGlobals(Fortran::lower::AbstractConverter &converter) {
+    skipRegistration = true;
+    for (const TypeInfoSymbol &info : registeredTypeInfoSymbols)
+      Fortran::lower::createRuntimeTypeInfoGlobal(converter, info.loc,
+                                                  info.symbol.get());
+    registeredTypeInfoSymbols.clear();
+  }
+
+private:
+  /// Store the runtime type descriptors that will be required for the
+  /// derived type that have been converted to FIR derived types.
+  llvm::SmallVector<TypeInfoSymbol> registeredTypeInfoSymbols;
+  /// Create derived type runtime info global immediately without storing the
+  /// symbol in registeredTypeInfoSymbols.
+  bool skipRegistration = false;
+  /// Track symbols symbols processed during and after the registration
+  /// to avoid infinite loops between type conversions and global variable
+  /// creation.
+  llvm::SmallSetVector<Fortran::semantics::SymbolRef, 64> seen;
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // FirConverter
 //===----------------------------------------------------------------------===//
@@ -101,6 +163,12 @@ class FirConverter : public Fortran::lower::AbstractConverter {
           },
           u);
     }
+
+    /// Once all the code has been translated, create runtime type info
+    /// global data structure for the derived types that have been
+    /// processed.
+    createGlobalOutsideOfFunctionLowering(
+        [&]() { runtimeTypeInfoConverter.createTypeInfoGlobals(*this); });
   }
 
   /// Declare a function.
@@ -689,6 +757,12 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     hostAssocTuple = val;
   }
 
+  void registerRuntimeTypeInfo(
+      mlir::Location loc,
+      Fortran::lower::SymbolRef typeInfoSym) override final {
+    runtimeTypeInfoConverter.registerTypeInfoSymbol(*this, loc, typeInfoSym);
+  }
+
 private:
   FirConverter() = delete;
   FirConverter(const FirConverter &) = delete;
@@ -2319,6 +2393,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   Fortran::lower::pft::Evaluation *evalPtr = nullptr;
   Fortran::lower::SymMap localSymbols;
   Fortran::parser::CharBlock currentPosition;
+  RuntimeTypeInfoConverter runtimeTypeInfoConverter;
 
   /// Tuple of host assoicated variables.
   mlir::Value hostAssocTuple;

flang/lib/Lower/CMakeLists.txt

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