xref: /llvm-project/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h (revision db791b278a414fb6df1acc1799adcf11d8fb9169)
1 //===- AffineStructures.h - MLIR Affine Structures Class --------*- 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 // Structures for affine/polyhedral analysis of ML functions.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_AFFINESTRUCTURES_H
14 #define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINESTRUCTURES_H
15 
16 #include "mlir/Analysis/FlatLinearValueConstraints.h"
17 #include "mlir/Analysis/Presburger/IntegerRelation.h"
18 #include "mlir/Analysis/Presburger/Matrix.h"
19 #include "mlir/IR/AffineExpr.h"
20 #include "mlir/IR/OpDefinition.h"
21 #include <optional>
22 
23 namespace mlir {
24 class AffineMap;
25 class IntegerSet;
26 class MemRefType;
27 class MLIRContext;
28 struct MutableAffineMap;
29 class Value;
30 
31 namespace presburger {
32 class MultiAffineFunction;
33 } // namespace presburger
34 
35 namespace affine {
36 class AffineCondition;
37 class AffineForOp;
38 class AffineIfOp;
39 class AffineParallelOp;
40 class AffineValueMap;
41 
42 /// FlatAffineValueConstraints is an extension of FlatLinearValueConstraints
43 /// with helper functions for Affine dialect ops.
44 class FlatAffineValueConstraints : public FlatLinearValueConstraints {
45 public:
46   using FlatLinearValueConstraints::FlatLinearValueConstraints;
47 
48   /// Return the kind of this object.
getKind()49   Kind getKind() const override { return Kind::FlatAffineValueConstraints; }
50 
classof(const IntegerRelation * cst)51   static bool classof(const IntegerRelation *cst) {
52     return cst->getKind() >= Kind::FlatAffineValueConstraints &&
53            cst->getKind() <= Kind::FlatAffineRelation;
54   }
55 
56   /// Adds constraints (lower and upper bounds) for the specified 'affine.for'
57   /// operation's Value using IR information stored in its bound maps. The
58   /// right variable is first looked up using `forOp`'s Value. Asserts if the
59   /// Value corresponding to the 'affine.for' operation isn't found in the
60   /// constraint system. Returns failure for the yet unimplemented/unsupported
61   /// cases.  Any new variables that are found in the bound operands of the
62   /// 'affine.for' operation are added as trailing variables (either
63   /// dimensional or symbolic depending on whether the operand is a valid
64   /// symbol).
65   LogicalResult addAffineForOpDomain(AffineForOp forOp);
66 
67   /// Add constraints (lower and upper bounds) for the specified
68   /// 'affine.parallel' operation's Value using IR information stored in its
69   /// bound maps. Returns failure for the yet unimplemented/unsupported cases.
70   /// Asserts if the Value corresponding to the 'affine.parallel' operation
71   /// isn't found in the constraint system.
72   LogicalResult addAffineParallelOpDomain(AffineParallelOp parallelOp);
73 
74   /// Adds constraints (lower and upper bounds) for each loop in the loop nest
75   /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice.
76   /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in
77   /// the nest, sorted outer-to-inner. `operands` contains the bound operands
78   /// for a single bound map. All the bound maps will use the same bound
79   /// operands. Note that some loops described by a computation slice might not
80   /// exist yet in the IR so the Value attached to those dimension variables
81   /// might be empty. For that reason, this method doesn't perform Value
82   /// look-ups to retrieve the dimension variable positions. Instead, it
83   /// assumes the position of the dim variables in the constraint system is
84   /// the same as the position of the loop in the loop nest.
85   LogicalResult addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
86                                        ArrayRef<AffineMap> ubMaps,
87                                        ArrayRef<Value> operands);
88 
89   /// Adds constraints imposed by the `affine.if` operation. These constraints
90   /// are collected from the IntegerSet attached to the given `affine.if`
91   /// instance argument (`ifOp`). It is asserted that:
92   /// 1) The IntegerSet of the given `affine.if` instance should not contain
93   /// semi-affine expressions,
94   /// 2) The columns of the constraint system created from `ifOp` should match
95   /// the columns in the current one regarding numbers and values.
96   void addAffineIfOpDomain(AffineIfOp ifOp);
97 
98   /// Adds a bound for the variable at the specified position with constraints
99   /// being drawn from the specified bound map and operands. In case of an
100   /// EQ bound, the  bound map is expected to have exactly one result. In case
101   /// of a LB/UB, the bound map may have more than one result, for each of which
102   /// an inequality is added.
103   LogicalResult addBound(presburger::BoundType type, unsigned pos,
104                          AffineMap boundMap, ValueRange operands);
105   using FlatLinearValueConstraints::addBound;
106 
107   /// Add the specified values as a dim or symbol var depending on its nature,
108   /// if it already doesn't exist in the system. `val` has to be either a
109   /// terminal symbol or a loop IV, i.e., it cannot be the result affine.apply
110   /// of any symbols or loop IVs. The variable is added to the end of the
111   /// existing dims or symbols. Additional information on the variable is
112   /// extracted from the IR and added to the constraint system.
113   void addInductionVarOrTerminalSymbol(Value val);
114 
115   /// Adds slice lower bounds represented by lower bounds in `lbMaps` and upper
116   /// bounds in `ubMaps` to each variable in the constraint system which has
117   /// a value in `values`. Note that both lower/upper bounds share the same
118   /// operand list `operands`.
119   /// This function assumes `values.size` == `lbMaps.size` == `ubMaps.size`.
120   /// Note that both lower/upper bounds use operands from `operands`.
121   LogicalResult addSliceBounds(ArrayRef<Value> values,
122                                ArrayRef<AffineMap> lbMaps,
123                                ArrayRef<AffineMap> ubMaps,
124                                ArrayRef<Value> operands);
125 
126   /// Changes all symbol variables which are loop IVs to dim variables.
127   void convertLoopIVSymbolsToDims();
128 
129   /// Returns the bound for the variable at `pos` from the inequality at
130   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
131   /// affine value map can either be a lower bound or an upper bound depending
132   /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does
133   /// not involve the `pos`th variable.
134   void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos,
135                                AffineValueMap &vmap,
136                                MLIRContext *context) const;
137 
138   /// Composes the affine value map with this FlatAffineValueConstrains, adding
139   /// the results of the map as dimensions at the front
140   /// [0, vMap->getNumResults()) and with the dimensions set to the equalities
141   /// specified by the value map.
142   ///
143   /// Returns failure if the composition fails (when vMap is a semi-affine map).
144   /// The vMap's operand Value's are used to look up the right positions in
145   /// the FlatAffineValueConstraints with which to associate. Every operand of
146   /// vMap should have a matching dim/symbol column in this constraint system
147   /// (with the same associated Value).
148   LogicalResult composeMap(const AffineValueMap *vMap);
149 };
150 
151 /// A FlatAffineRelation represents a set of ordered pairs (domain -> range)
152 /// where "domain" and "range" are tuples of variables. The relation is
153 /// represented as a FlatAffineValueConstraints with separation of dimension
154 /// variables into domain and  range. The variables are stored as:
155 /// [domainVars, rangeVars, symbolVars, localVars, constant].
156 ///
157 /// Deprecated: use IntegerRelation and store SSA Values in the PresburgerSpace
158 /// of the relation using PresburgerSpace::identifiers. Note that
159 /// FlatAffineRelation::numDomainDims and FlatAffineRelation::numRangeDims are
160 /// independent of numDomain and numRange of the relation's space. In
161 /// particular, operations such as FlatAffineRelation::compose do not ensure
162 /// consistency between numDomainDims/numRangeDims and numDomain/numRange which
163 /// may lead to unexpected behaviour.
164 class FlatAffineRelation : public FlatAffineValueConstraints {
165 public:
166   FlatAffineRelation(unsigned numReservedInequalities,
167                      unsigned numReservedEqualities, unsigned numReservedCols,
168                      unsigned numDomainDims, unsigned numRangeDims,
169                      unsigned numSymbols, unsigned numLocals,
170                      ArrayRef<std::optional<Value>> valArgs = {})
171       : FlatAffineValueConstraints(
172             numReservedInequalities, numReservedEqualities, numReservedCols,
173             numDomainDims + numRangeDims, numSymbols, numLocals, valArgs),
174         numDomainDims(numDomainDims), numRangeDims(numRangeDims) {}
175 
176   FlatAffineRelation(unsigned numDomainDims = 0, unsigned numRangeDims = 0,
177                      unsigned numSymbols = 0, unsigned numLocals = 0)
178       : FlatAffineValueConstraints(numDomainDims + numRangeDims, numSymbols,
179                                    numLocals),
180         numDomainDims(numDomainDims), numRangeDims(numRangeDims) {}
181 
FlatAffineRelation(unsigned numDomainDims,unsigned numRangeDims,FlatAffineValueConstraints & fac)182   FlatAffineRelation(unsigned numDomainDims, unsigned numRangeDims,
183                      FlatAffineValueConstraints &fac)
184       : FlatAffineValueConstraints(fac), numDomainDims(numDomainDims),
185         numRangeDims(numRangeDims) {}
186 
FlatAffineRelation(unsigned numDomainDims,unsigned numRangeDims,IntegerPolyhedron & fac)187   FlatAffineRelation(unsigned numDomainDims, unsigned numRangeDims,
188                      IntegerPolyhedron &fac)
189       : FlatAffineValueConstraints(fac), numDomainDims(numDomainDims),
190         numRangeDims(numRangeDims) {}
191 
192   /// Return the kind of this object.
getKind()193   Kind getKind() const override { return Kind::FlatAffineRelation; }
194 
classof(const IntegerRelation * cst)195   static bool classof(const IntegerRelation *cst) {
196     return cst->getKind() == Kind::FlatAffineRelation;
197   }
198 
199   /// Returns a set corresponding to the domain/range of the affine relation.
200   FlatAffineValueConstraints getDomainSet() const;
201   FlatAffineValueConstraints getRangeSet() const;
202 
203   /// Returns the number of variables corresponding to domain/range of
204   /// relation.
getNumDomainDims()205   inline unsigned getNumDomainDims() const { return numDomainDims; }
getNumRangeDims()206   inline unsigned getNumRangeDims() const { return numRangeDims; }
207 
208   /// Given affine relation `other: (domainOther -> rangeOther)`, this operation
209   /// takes the composition of `other` on `this: (domainThis -> rangeThis)`.
210   /// The resulting relation represents tuples of the form: `domainOther ->
211   /// rangeThis`.
212   void compose(const FlatAffineRelation &other);
213 
214   /// Swap domain and range of the relation.
215   /// `(domain -> range)` is converted to `(range -> domain)`.
216   void inverse();
217 
218   /// Insert `num` variables of the specified kind after the `pos` variable
219   /// of that kind. The coefficient columns corresponding to the added
220   /// variables are initialized to zero.
221   void insertDomainVar(unsigned pos, unsigned num = 1);
222   void insertRangeVar(unsigned pos, unsigned num = 1);
223 
224   /// Append `num` variables of the specified kind after the last variable
225   /// of that kind. The coefficient columns corresponding to the added
226   /// variables are initialized to zero.
227   void appendDomainVar(unsigned num = 1);
228   void appendRangeVar(unsigned num = 1);
229 
230   /// Removes variables in the column range [varStart, varLimit), and copies any
231   /// remaining valid data into place, updates member variables, and resizes
232   /// arrays as needed.
233   void removeVarRange(VarKind kind, unsigned varStart,
234                       unsigned varLimit) override;
235   using IntegerRelation::removeVarRange;
236 
237 protected:
238   // Number of dimension variables corresponding to domain variables.
239   unsigned numDomainDims;
240 
241   // Number of dimension variables corresponding to range variables.
242   unsigned numRangeDims;
243 };
244 
245 /// Builds a relation from the given AffineMap/AffineValueMap `map`, containing
246 /// all pairs of the form `operands -> result` that satisfy `map`. `rel` is set
247 /// to the relation built. For example, give the AffineMap:
248 ///
249 ///   (d0, d1)[s0] -> (d0 + s0, d0 - s0)
250 ///
251 /// the resulting relation formed is:
252 ///
253 ///   (d0, d1) -> (r1, r2)
254 ///   [d0  d1  r1  r2  s0  const]
255 ///    1   0   -1   0  1     0     = 0
256 ///    0   1    0  -1  -1    0     = 0
257 ///
258 /// For AffineValueMap, the domain and symbols have Value set corresponding to
259 /// the Value in `map`. Returns failure if the AffineMap could not be flattened
260 /// (i.e., semi-affine is not yet handled).
261 LogicalResult getRelationFromMap(AffineMap &map,
262                                  presburger::IntegerRelation &rel);
263 LogicalResult getRelationFromMap(const AffineValueMap &map,
264                                  presburger::IntegerRelation &rel);
265 
266 } // namespace affine
267 } // namespace mlir
268 
269 #endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINESTRUCTURES_H
270