xref: /llvm-project/flang/include/flang/Lower/VectorSubscripts.h (revision 9aeb7f035bdde83501e5eddd9e6ad175b8ed697f)
1 //===-- VectorSubscripts.h -- vector subscripts tools -----------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 ///
9 /// \file
10 /// \brief Defines a compiler internal representation for lowered designators
11 ///  containing vector subscripts. This representation allows working on such
12 ///  designators in custom ways while ensuring the designator subscripts are
13 ///  only evaluated once. It is mainly intended for cases that do not fit in
14 ///  the array expression lowering framework like input IO in presence of
15 ///  vector subscripts.
16 ///
17 //===----------------------------------------------------------------------===//
18 
19 #ifndef FORTRAN_LOWER_VECTORSUBSCRIPTS_H
20 #define FORTRAN_LOWER_VECTORSUBSCRIPTS_H
21 
22 #include "flang/Optimizer/Builder/BoxValue.h"
23 
24 namespace fir {
25 class FirOpBuilder;
26 }
27 
28 namespace Fortran {
29 
30 namespace evaluate {
31 template <typename>
32 class Expr;
33 struct SomeType;
34 } // namespace evaluate
35 
36 namespace lower {
37 
38 class AbstractConverter;
39 class StatementContext;
40 
41 /// VectorSubscriptBox is a lowered representation for any Designator<T> that
42 /// contain at least one vector subscript.
43 ///
44 /// A designator `x%a(i,j)%b(1:foo():1, vector, k)%c%d(m)%e1
45 /// Is lowered into:
46 ///   - an ExtendedValue for ranked base (x%a(i,j)%b)
47 ///   - mlir:Values and ExtendedValues for the triplet, vector subscript and
48 ///     scalar subscripts of the ranked array reference (1:foo():1, vector, k)
49 ///   - a list of fir.field_index and scalar integers mlir::Value for the
50 ///   component
51 ///     path at the right of the ranked array ref (%c%d(m)%e).
52 ///
53 /// This representation allows later creating loops over the designator elements
54 /// and fir.array_coor to get the element addresses without re-evaluating any
55 /// sub-expressions.
56 class VectorSubscriptBox {
57 public:
58   /// Type of the callbacks that can be passed to work with the element
59   /// addresses.
60   using ElementalGenerator = std::function<void(const fir::ExtendedValue &)>;
61   using ElementalGeneratorWithBoolReturn =
62       std::function<mlir::Value(const fir::ExtendedValue &)>;
63   struct LoweredVectorSubscript {
LoweredVectorSubscriptLoweredVectorSubscript64     LoweredVectorSubscript(fir::ExtendedValue &&vector, mlir::Value size)
65         : vector{std::move(vector)}, size{size} {}
66     fir::ExtendedValue vector;
67     // Vector size, guaranteed to be of indexType.
68     mlir::Value size;
69   };
70   struct LoweredTriplet {
71     // Triplets value, guaranteed to be of indexType.
72     mlir::Value lb;
73     mlir::Value ub;
74     mlir::Value stride;
75   };
76   using LoweredSubscript =
77       std::variant<mlir::Value, LoweredTriplet, LoweredVectorSubscript>;
78   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)79   VectorSubscriptBox(
80       fir::ExtendedValue &&loweredBase,
81       llvm::SmallVector<LoweredSubscript, 16> &&loweredSubscripts,
82       llvm::SmallVector<mlir::Value> &&componentPath,
83       MaybeSubstring substringBounds, mlir::Type elementType)
84       : loweredBase{std::move(loweredBase)}, loweredSubscripts{std::move(
85                                                  loweredSubscripts)},
86         componentPath{std::move(componentPath)},
87         substringBounds{substringBounds}, elementType{elementType} {};
88 
89   /// Loop over the elements described by the VectorSubscriptBox, and call
90   /// \p elementalGenerator inside the loops with the element addresses.
91   void loopOverElements(fir::FirOpBuilder &builder, mlir::Location loc,
92                         const ElementalGenerator &elementalGenerator);
93 
94   /// Loop over the elements described by the VectorSubscriptBox while a
95   /// condition is true, and call \p elementalGenerator inside the loops with
96   /// the element addresses. The initial condition value is \p initialCondition,
97   /// and then it is the result of \p elementalGenerator. The value of the
98   /// condition after the loops is returned.
99   mlir::Value loopOverElementsWhile(
100       fir::FirOpBuilder &builder, mlir::Location loc,
101       const ElementalGeneratorWithBoolReturn &elementalGenerator,
102       mlir::Value initialCondition);
103 
104   /// Return the type of the elements of the array section.
getElementType()105   mlir::Type getElementType() { return elementType; }
106 
107 private:
108   /// Common implementation for DoLoop and IterWhile loop creations.
109   template <typename LoopType, typename Generator>
110   mlir::Value loopOverElementsBase(fir::FirOpBuilder &builder,
111                                    mlir::Location loc,
112                                    const Generator &elementalGenerator,
113                                    mlir::Value initialCondition);
114   /// Create sliceOp for the designator.
115   mlir::Value createSlice(fir::FirOpBuilder &builder, mlir::Location loc);
116 
117   /// Create ExtendedValue the element inside the loop.
118   fir::ExtendedValue getElementAt(fir::FirOpBuilder &builder,
119                                   mlir::Location loc, mlir::Value shape,
120                                   mlir::Value slice,
121                                   mlir::ValueRange inductionVariables);
122 
123   /// Generate the [lb, ub, step] to loop over the section (in loop order, not
124   /// Fortran dimension order).
125   llvm::SmallVector<std::tuple<mlir::Value, mlir::Value, mlir::Value>>
126   genLoopBounds(fir::FirOpBuilder &builder, mlir::Location loc);
127 
128   /// Lowered base of the ranked array ref.
129   fir::ExtendedValue loweredBase;
130   /// Subscripts values of the rank arrayRef part.
131   llvm::SmallVector<LoweredSubscript, 16> loweredSubscripts;
132   /// Scalar subscripts and components at the right of the ranked
133   /// array ref part of any.
134   llvm::SmallVector<mlir::Value> componentPath;
135   /// List of substring bounds if this is a substring (only the lower bound if
136   /// the upper is implicit).
137   MaybeSubstring substringBounds;
138   /// Type of the elements described by this array section.
139   mlir::Type elementType;
140 };
141 
142 /// Lower \p expr, that must be an designator containing vector subscripts, to a
143 /// VectorSubscriptBox representation. This causes evaluation of all the
144 /// subscripts. Any required clean-ups from subscript expression are added to \p
145 /// stmtCtx.
146 VectorSubscriptBox genVectorSubscriptBox(
147     mlir::Location loc, Fortran::lower::AbstractConverter &converter,
148     Fortran::lower::StatementContext &stmtCtx,
149     const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &expr);
150 
151 } // namespace lower
152 } // namespace Fortran
153 
154 #endif // FORTRAN_LOWER_VECTORSUBSCRIPTS_H
155