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