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_ANALYSIS_AFFINESTRUCTURES_H 14 #define MLIR_ANALYSIS_AFFINESTRUCTURES_H 15 16 #include "mlir/Analysis/Presburger/Matrix.h" 17 #include "mlir/IR/AffineExpr.h" 18 #include "mlir/IR/OpDefinition.h" 19 #include "mlir/Support/LogicalResult.h" 20 21 namespace mlir { 22 23 class AffineCondition; 24 class AffineForOp; 25 class AffineIfOp; 26 class AffineMap; 27 class AffineValueMap; 28 class IntegerSet; 29 class MLIRContext; 30 class Value; 31 class MemRefType; 32 struct MutableAffineMap; 33 34 /// A flat list of affine equalities and inequalities in the form. 35 /// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} >= 0 36 /// Equality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} == 0 37 /// 38 /// FlatAffineConstraints stores coefficients in a contiguous buffer (one buffer 39 /// for equalities and one for inequalities). The size of each buffer is 40 /// numReservedCols * number of inequalities (or equalities). The reserved size 41 /// is numReservedCols * numReservedInequalities (or numReservedEqualities). A 42 /// coefficient (r, c) lives at the location numReservedCols * r + c in the 43 /// buffer. The extra space between getNumCols() and numReservedCols exists to 44 /// prevent frequent movement of data when adding columns, especially at the 45 /// end. 46 /// 47 /// The identifiers x_0, x_1, ... appear in the order: dimensional identifiers, 48 /// symbolic identifiers, and local identifiers. The local identifiers 49 /// correspond to local/internal variables created when converting from 50 /// AffineExpr's containing mod's and div's; they are thus needed to increase 51 /// representational power. Each local identifier is always (by construction) a 52 /// floordiv of a pure add/mul affine function of dimensional, symbolic, and 53 /// other local identifiers, in a non-mutually recursive way. Hence, every local 54 /// identifier can ultimately always be recovered as an affine function of 55 /// dimensional and symbolic identifiers (involving floordiv's); note however 56 /// that some floordiv combinations are converted to mod's by AffineExpr 57 /// construction. 58 /// 59 class FlatAffineConstraints { 60 public: 61 /// All derived classes of FlatAffineConstraints. 62 enum class Kind { FlatAffineConstraints, FlatAffineValueConstraints }; 63 64 /// Kind of identifier (column). 65 enum IdKind { Dimension, Symbol, Local }; 66 67 /// Constructs a constraint system reserving memory for the specified number 68 /// of constraints and identifiers. FlatAffineConstraints(unsigned numReservedInequalities,unsigned numReservedEqualities,unsigned numReservedCols,unsigned numDims,unsigned numSymbols,unsigned numLocals)69 FlatAffineConstraints(unsigned numReservedInequalities, 70 unsigned numReservedEqualities, 71 unsigned numReservedCols, unsigned numDims, 72 unsigned numSymbols, unsigned numLocals) 73 : numIds(numDims + numSymbols + numLocals), numDims(numDims), 74 numSymbols(numSymbols), 75 equalities(0, numIds + 1, numReservedEqualities, numReservedCols), 76 inequalities(0, numIds + 1, numReservedInequalities, numReservedCols) { 77 assert(numReservedCols >= numIds + 1); 78 } 79 80 /// Constructs a constraint system with the specified number of 81 /// dimensions and symbols. 82 FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0, 83 unsigned numLocals = 0) 84 : FlatAffineConstraints(/*numReservedInequalities=*/0, 85 /*numReservedEqualities=*/0, 86 /*numReservedCols=*/numDims + numSymbols + 87 numLocals + 1, 88 numDims, numSymbols, numLocals) {} 89 90 /// Return a system with no constraints, i.e., one which is satisfied by all 91 /// points. 92 static FlatAffineConstraints getUniverse(unsigned numDims = 0, 93 unsigned numSymbols = 0) { 94 return FlatAffineConstraints(numDims, numSymbols); 95 } 96 97 /// Creates an affine constraint system from an IntegerSet. 98 explicit FlatAffineConstraints(IntegerSet set); 99 100 FlatAffineConstraints(const MutableAffineMap &map); 101 102 virtual ~FlatAffineConstraints() = default; 103 104 /// Return the kind of this FlatAffineConstraints. getKind()105 virtual Kind getKind() const { return Kind::FlatAffineConstraints; } 106 classof(const FlatAffineConstraints * cst)107 static bool classof(const FlatAffineConstraints *cst) { return true; } 108 109 /// Clears any existing data and reserves memory for the specified 110 /// constraints. 111 virtual void reset(unsigned numReservedInequalities, 112 unsigned numReservedEqualities, unsigned numReservedCols, 113 unsigned numDims, unsigned numSymbols, 114 unsigned numLocals = 0); 115 116 void reset(unsigned numDims = 0, unsigned numSymbols = 0, 117 unsigned numLocals = 0); 118 119 /// Appends constraints from `other` into `this`. This is equivalent to an 120 /// intersection with no simplification of any sort attempted. 121 void append(const FlatAffineConstraints &other); 122 123 /// Checks for emptiness by performing variable elimination on all 124 /// identifiers, running the GCD test on each equality constraint, and 125 /// checking for invalid constraints. Returns true if the GCD test fails for 126 /// any equality, or if any invalid constraints are discovered on any row. 127 /// Returns false otherwise. 128 bool isEmpty() const; 129 130 /// Runs the GCD test on all equality constraints. Returns true if this test 131 /// fails on any equality. Returns false otherwise. 132 /// This test can be used to disprove the existence of a solution. If it 133 /// returns true, no integer solution to the equality constraints can exist. 134 bool isEmptyByGCDTest() const; 135 136 /// Returns true if the set of constraints is found to have no solution, 137 /// false if a solution exists. Uses the same algorithm as 138 /// `findIntegerSample`. 139 bool isIntegerEmpty() const; 140 141 /// Returns a matrix where each row is a vector along which the polytope is 142 /// bounded. The span of the returned vectors is guaranteed to contain all 143 /// such vectors. The returned vectors are NOT guaranteed to be linearly 144 /// independent. This function should not be called on empty sets. 145 Matrix getBoundedDirections() const; 146 147 /// Find an integer sample point satisfying the constraints using a 148 /// branch and bound algorithm with generalized basis reduction, with some 149 /// additional processing using Simplex for unbounded sets. 150 /// 151 /// Returns an integer sample point if one exists, or an empty Optional 152 /// otherwise. 153 Optional<SmallVector<int64_t, 8>> findIntegerSample() const; 154 155 /// Returns true if the given point satisfies the constraints, or false 156 /// otherwise. 157 bool containsPoint(ArrayRef<int64_t> point) const; 158 159 /// Find pairs of inequalities identified by their position indices, using 160 /// which an explicit representation for each local variable can be computed 161 /// The pairs are stored as indices of upperbound, lowerbound 162 /// inequalities. If no such pair can be found, it is stored as llvm::None. 163 void getLocalReprLbUbPairs( 164 std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const; 165 166 // Clones this object. 167 std::unique_ptr<FlatAffineConstraints> clone() const; 168 169 /// Returns the value at the specified equality row and column. atEq(unsigned i,unsigned j)170 inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); } atEq(unsigned i,unsigned j)171 inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); } 172 173 /// Returns the value at the specified inequality row and column. atIneq(unsigned i,unsigned j)174 inline int64_t atIneq(unsigned i, unsigned j) const { 175 return inequalities(i, j); 176 } atIneq(unsigned i,unsigned j)177 inline int64_t &atIneq(unsigned i, unsigned j) { return inequalities(i, j); } 178 179 /// Returns the number of columns in the constraint system. getNumCols()180 inline unsigned getNumCols() const { return numIds + 1; } 181 getNumEqualities()182 inline unsigned getNumEqualities() const { return equalities.getNumRows(); } 183 getNumInequalities()184 inline unsigned getNumInequalities() const { 185 return inequalities.getNumRows(); 186 } 187 getNumReservedEqualities()188 inline unsigned getNumReservedEqualities() const { 189 return equalities.getNumReservedRows(); 190 } 191 getNumReservedInequalities()192 inline unsigned getNumReservedInequalities() const { 193 return inequalities.getNumReservedRows(); 194 } 195 getEquality(unsigned idx)196 inline ArrayRef<int64_t> getEquality(unsigned idx) const { 197 return equalities.getRow(idx); 198 } 199 getInequality(unsigned idx)200 inline ArrayRef<int64_t> getInequality(unsigned idx) const { 201 return inequalities.getRow(idx); 202 } 203 204 /// The type of bound: equal, lower bound or upper bound. 205 enum BoundType { EQ, LB, UB }; 206 207 /// Adds a bound for the identifier at the specified position with constraints 208 /// being drawn from the specified bound map. In case of an EQ bound, the 209 /// bound map is expected to have exactly one result. In case of a LB/UB, the 210 /// bound map may have more than one result, for each of which an inequality 211 /// is added. 212 /// Note: The dimensions/symbols of this FlatAffineConstraints must match the 213 /// dimensions/symbols of the affine map. 214 LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap); 215 216 /// Adds a constant bound for the specified identifier. 217 void addBound(BoundType type, unsigned pos, int64_t value); 218 219 /// Adds a constant bound for the specified expression. 220 void addBound(BoundType type, ArrayRef<int64_t> expr, int64_t value); 221 222 /// Returns the constraint system as an integer set. Returns a null integer 223 /// set if the system has no constraints, or if an integer set couldn't be 224 /// constructed as a result of a local variable's explicit representation not 225 /// being known and such a local variable appearing in any of the constraints. 226 IntegerSet getAsIntegerSet(MLIRContext *context) const; 227 228 /// Computes the lower and upper bounds of the first `num` dimensional 229 /// identifiers (starting at `offset`) as an affine map of the remaining 230 /// identifiers (dimensional and symbolic). This method is able to detect 231 /// identifiers as floordiv's and mod's of affine expressions of other 232 /// identifiers with respect to (positive) constants. Sets bound map to a 233 /// null AffineMap if such a bound can't be found (or yet unimplemented). 234 void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context, 235 SmallVectorImpl<AffineMap> *lbMaps, 236 SmallVectorImpl<AffineMap> *ubMaps); 237 238 /// Adds an inequality (>= 0) from the coefficients specified in `inEq`. 239 void addInequality(ArrayRef<int64_t> inEq); 240 /// Adds an equality from the coefficients specified in `eq`. 241 void addEquality(ArrayRef<int64_t> eq); 242 243 /// Adds a new local identifier as the floordiv of an affine function of other 244 /// identifiers, the coefficients of which are provided in `dividend` and with 245 /// respect to a positive constant `divisor`. Two constraints are added to the 246 /// system to capture equivalence with the floordiv: 247 /// q = dividend floordiv c <=> c*q <= dividend <= c*q + c - 1. 248 void addLocalFloorDiv(ArrayRef<int64_t> dividend, int64_t divisor); 249 250 /// Swap the posA^th identifier with the posB^th identifier. 251 virtual void swapId(unsigned posA, unsigned posB); 252 253 /// Insert `num` identifiers of the specified kind at position `pos`. 254 /// Positions are relative to the kind of identifier. The coefficient columns 255 /// corresponding to the added identifiers are initialized to zero. Return the 256 /// absolute column position (i.e., not relative to the kind of identifier) 257 /// of the first added identifier. 258 unsigned insertDimId(unsigned pos, unsigned num = 1); 259 unsigned insertSymbolId(unsigned pos, unsigned num = 1); 260 unsigned insertLocalId(unsigned pos, unsigned num = 1); 261 virtual unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1); 262 263 /// Append `num` identifiers of the specified kind after the last identifier. 264 /// of that kind. Return the position of the first appended column. The 265 /// coefficient columns corresponding to the added identifiers are initialized 266 /// to zero. 267 unsigned appendDimId(unsigned num = 1); 268 unsigned appendSymbolId(unsigned num = 1); 269 unsigned appendLocalId(unsigned num = 1); 270 271 /// Composes an affine map whose dimensions and symbols match one to one with 272 /// the dimensions and symbols of this FlatAffineConstraints. The results of 273 /// the map `other` are added as the leading dimensions of this constraint 274 /// system. Returns failure if `other` is a semi-affine map. 275 LogicalResult composeMatchingMap(AffineMap other); 276 277 /// Projects out (aka eliminates) `num` identifiers starting at position 278 /// `pos`. The resulting constraint system is the shadow along the dimensions 279 /// that still exist. This method may not always be integer exact. 280 // TODO: deal with integer exactness when necessary - can return a value to 281 // mark exactness for example. 282 void projectOut(unsigned pos, unsigned num); projectOut(unsigned pos)283 inline void projectOut(unsigned pos) { return projectOut(pos, 1); } 284 285 /// Removes identifiers of the specified kind with the specified pos (or 286 /// within the specified range) from the system. The specified location is 287 /// relative to the first identifier of the specified kind. 288 void removeId(IdKind kind, unsigned pos); 289 void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit); 290 291 /// Removes the specified identifier from the system. 292 void removeId(unsigned pos); 293 294 void removeEquality(unsigned pos); 295 void removeInequality(unsigned pos); 296 297 /// Remove the (in)equalities at positions [start, end). 298 void removeEqualityRange(unsigned start, unsigned end); 299 void removeInequalityRange(unsigned start, unsigned end); 300 301 /// Sets the `values.size()` identifiers starting at `po`s to the specified 302 /// values and removes them. 303 void setAndEliminate(unsigned pos, ArrayRef<int64_t> values); 304 305 /// Changes the partition between dimensions and symbols. Depending on the new 306 /// symbol count, either a chunk of trailing dimensional identifiers becomes 307 /// symbols, or some of the leading symbols become dimensions. 308 void setDimSymbolSeparation(unsigned newSymbolCount); 309 310 /// Tries to fold the specified identifier to a constant using a trivial 311 /// equality detection; if successful, the constant is substituted for the 312 /// identifier everywhere in the constraint system and then removed from the 313 /// system. 314 LogicalResult constantFoldId(unsigned pos); 315 316 /// This method calls `constantFoldId` for the specified range of identifiers, 317 /// `num` identifiers starting at position `pos`. 318 void constantFoldIdRange(unsigned pos, unsigned num); 319 320 /// Updates the constraints to be the smallest bounding (enclosing) box that 321 /// contains the points of `this` set and that of `other`, with the symbols 322 /// being treated specially. For each of the dimensions, the min of the lower 323 /// bounds (symbolic) and the max of the upper bounds (symbolic) is computed 324 /// to determine such a bounding box. `other` is expected to have the same 325 /// dimensional identifiers as this constraint system (in the same order). 326 /// 327 /// E.g.: 328 /// 1) this = {0 <= d0 <= 127}, 329 /// other = {16 <= d0 <= 192}, 330 /// output = {0 <= d0 <= 192} 331 /// 2) this = {s0 + 5 <= d0 <= s0 + 20}, 332 /// other = {s0 + 1 <= d0 <= s0 + 9}, 333 /// output = {s0 + 1 <= d0 <= s0 + 20} 334 /// 3) this = {0 <= d0 <= 5, 1 <= d1 <= 9} 335 /// other = {2 <= d0 <= 6, 5 <= d1 <= 15}, 336 /// output = {0 <= d0 <= 6, 1 <= d1 <= 15} 337 LogicalResult unionBoundingBox(const FlatAffineConstraints &other); 338 getNumConstraints()339 unsigned getNumConstraints() const { 340 return getNumInequalities() + getNumEqualities(); 341 } getNumIds()342 inline unsigned getNumIds() const { return numIds; } getNumDimIds()343 inline unsigned getNumDimIds() const { return numDims; } getNumSymbolIds()344 inline unsigned getNumSymbolIds() const { return numSymbols; } getNumDimAndSymbolIds()345 inline unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; } getNumLocalIds()346 inline unsigned getNumLocalIds() const { 347 return numIds - numDims - numSymbols; 348 } 349 350 /// Replaces the contents of this FlatAffineConstraints with `other`. 351 virtual void clearAndCopyFrom(const FlatAffineConstraints &other); 352 353 /// Returns the smallest known constant bound for the extent of the specified 354 /// identifier (pos^th), i.e., the smallest known constant that is greater 355 /// than or equal to 'exclusive upper bound' - 'lower bound' of the 356 /// identifier. This constant bound is guaranteed to be non-negative. Returns 357 /// None if it's not a constant. This method employs trivial (low complexity / 358 /// cost) checks and detection. Symbolic identifiers are treated specially, 359 /// i.e., it looks for constant differences between affine expressions 360 /// involving only the symbolic identifiers. `lb` and `ub` (along with the 361 /// `boundFloorDivisor`) are set to represent the lower and upper bound 362 /// associated with the constant difference: `lb`, `ub` have the coefficients, 363 /// and `boundFloorDivisor`, their divisor. `minLbPos` and `minUbPos` if 364 /// non-null are set to the position of the constant lower bound and upper 365 /// bound respectively (to the same if they are from an equality). Ex: if the 366 /// lower bound is [(s0 + s2 - 1) floordiv 32] for a system with three 367 /// symbolic identifiers, *lb = [1, 0, 1], lbDivisor = 32. See comments at 368 /// function definition for examples. 369 Optional<int64_t> getConstantBoundOnDimSize( 370 unsigned pos, SmallVectorImpl<int64_t> *lb = nullptr, 371 int64_t *boundFloorDivisor = nullptr, 372 SmallVectorImpl<int64_t> *ub = nullptr, unsigned *minLbPos = nullptr, 373 unsigned *minUbPos = nullptr) const; 374 375 /// Returns the constant bound for the pos^th identifier if there is one; 376 /// None otherwise. 377 // TODO: Support EQ bounds. 378 Optional<int64_t> getConstantBound(BoundType type, unsigned pos) const; 379 380 /// Gets the lower and upper bound of the `offset` + `pos`th identifier 381 /// treating [0, offset) U [offset + num, symStartPos) as dimensions and 382 /// [symStartPos, getNumDimAndSymbolIds) as symbols, and `pos` lies in 383 /// [0, num). The multi-dimensional maps in the returned pair represent the 384 /// max and min of potentially multiple affine expressions. The upper bound is 385 /// exclusive. `localExprs` holds pre-computed AffineExpr's for all local 386 /// identifiers in the system. 387 std::pair<AffineMap, AffineMap> 388 getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num, 389 unsigned symStartPos, ArrayRef<AffineExpr> localExprs, 390 MLIRContext *context) const; 391 392 /// Gather positions of all lower and upper bounds of the identifier at `pos`, 393 /// and optionally any equalities on it. In addition, the bounds are to be 394 /// independent of identifiers in position range [`offset`, `offset` + `num`). 395 void 396 getLowerAndUpperBoundIndices(unsigned pos, 397 SmallVectorImpl<unsigned> *lbIndices, 398 SmallVectorImpl<unsigned> *ubIndices, 399 SmallVectorImpl<unsigned> *eqIndices = nullptr, 400 unsigned offset = 0, unsigned num = 0) const; 401 402 /// Removes constraints that are independent of (i.e., do not have a 403 /// coefficient) identifiers in the range [pos, pos + num). 404 void removeIndependentConstraints(unsigned pos, unsigned num); 405 406 /// Returns true if the set can be trivially detected as being 407 /// hyper-rectangular on the specified contiguous set of identifiers. 408 bool isHyperRectangular(unsigned pos, unsigned num) const; 409 410 /// Removes duplicate constraints, trivially true constraints, and constraints 411 /// that can be detected as redundant as a result of differing only in their 412 /// constant term part. A constraint of the form <non-negative constant> >= 0 413 /// is considered trivially true. This method is a linear time method on the 414 /// constraints, does a single scan, and updates in place. It also normalizes 415 /// constraints by their GCD and performs GCD tightening on inequalities. 416 void removeTrivialRedundancy(); 417 418 /// A more expensive check than `removeTrivialRedundancy` to detect redundant 419 /// inequalities. 420 void removeRedundantInequalities(); 421 422 /// Removes redundant constraints using Simplex. Although the algorithm can 423 /// theoretically take exponential time in the worst case (rare), it is known 424 /// to perform much better in the average case. If V is the number of vertices 425 /// in the polytope and C is the number of constraints, the algorithm takes 426 /// O(VC) time. 427 void removeRedundantConstraints(); 428 429 /// Merge local ids of `this` and `other`. This is done by appending local ids 430 /// of `other` to `this` and inserting local ids of `this` to `other` at start 431 /// of its local ids. 432 void mergeLocalIds(FlatAffineConstraints &other); 433 434 /// Removes all equalities and inequalities. 435 void clearConstraints(); 436 437 void print(raw_ostream &os) const; 438 void dump() const; 439 440 protected: 441 /// Return the index at which the specified kind of id starts. 442 unsigned getIdKindOffset(IdKind kind) const; 443 444 /// Assert that `value` is at most the number of ids of the specified kind. 445 void assertAtMostNumIdKind(unsigned value, IdKind kind) const; 446 447 /// Returns false if the fields corresponding to various identifier counts, or 448 /// equality/inequality buffer sizes aren't consistent; true otherwise. This 449 /// is meant to be used within an assert internally. 450 virtual bool hasConsistentState() const; 451 452 /// Checks all rows of equality/inequality constraints for trivial 453 /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced 454 /// after elimination. Returns true if an invalid constraint is found; 455 /// false otherwise. 456 bool hasInvalidConstraint() const; 457 458 /// Returns the constant lower bound bound if isLower is true, and the upper 459 /// bound if isLower is false. 460 template <bool isLower> 461 Optional<int64_t> computeConstantLowerOrUpperBound(unsigned pos); 462 463 /// Given an affine map that is aligned with this constraint system: 464 /// * Flatten the map. 465 /// * Add newly introduced local columns at the beginning of this constraint 466 /// system (local column pos 0). 467 /// * Add equalities that define the new local columns to this constraint 468 /// system. 469 /// * Return the flattened expressions via `flattenedExprs`. 470 /// 471 /// Note: This is a shared helper function of `addLowerOrUpperBound` and 472 /// `composeMatchingMap`. 473 LogicalResult flattenAlignedMapAndMergeLocals( 474 AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs); 475 476 /// Eliminates a single identifier at `position` from equality and inequality 477 /// constraints. Returns `success` if the identifier was eliminated, and 478 /// `failure` otherwise. gaussianEliminateId(unsigned position)479 inline LogicalResult gaussianEliminateId(unsigned position) { 480 return success(gaussianEliminateIds(position, position + 1) == 1); 481 } 482 483 /// Removes local variables using equalities. Each equality is checked if it 484 /// can be reduced to the form: `e = affine-expr`, where `e` is a local 485 /// variable and `affine-expr` is an affine expression not containing `e`. 486 /// If an equality satisfies this form, the local variable is replaced in 487 /// each constraint and then removed. The equality used to replace this local 488 /// variable is also removed. 489 void removeRedundantLocalVars(); 490 491 /// Eliminates identifiers from equality and inequality constraints 492 /// in column range [posStart, posLimit). 493 /// Returns the number of variables eliminated. 494 unsigned gaussianEliminateIds(unsigned posStart, unsigned posLimit); 495 496 /// Eliminates the identifier at the specified position using Fourier-Motzkin 497 /// variable elimination, but uses Gaussian elimination if there is an 498 /// equality involving that identifier. If the result of the elimination is 499 /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is 500 /// set to true, a potential under approximation (subset) of the rational 501 /// shadow / exact integer shadow is computed. 502 // See implementation comments for more details. 503 virtual void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false, 504 bool *isResultIntegerExact = nullptr); 505 506 /// Tightens inequalities given that we are dealing with integer spaces. This 507 /// is similar to the GCD test but applied to inequalities. The constant term 508 /// can be reduced to the preceding multiple of the GCD of the coefficients, 509 /// i.e., 510 /// 64*i - 100 >= 0 => 64*i - 128 >= 0 (since 'i' is an integer). This is a 511 /// fast method (linear in the number of coefficients). 512 void gcdTightenInequalities(); 513 514 /// Normalized each constraints by the GCD of its coefficients. 515 void normalizeConstraintsByGCD(); 516 517 /// Removes identifiers in the column range [idStart, idLimit), and copies any 518 /// remaining valid data into place, updates member variables, and resizes 519 /// arrays as needed. 520 virtual void removeIdRange(unsigned idStart, unsigned idLimit); 521 522 /// Total number of identifiers. 523 unsigned numIds; 524 525 /// Number of identifiers corresponding to real dimensions. 526 unsigned numDims; 527 528 /// Number of identifiers corresponding to symbols (unknown but constant for 529 /// analysis). 530 unsigned numSymbols; 531 532 /// Coefficients of affine equalities (in == 0 form). 533 Matrix equalities; 534 535 /// Coefficients of affine inequalities (in >= 0 form). 536 Matrix inequalities; 537 538 /// A parameter that controls detection of an unrealistic number of 539 /// constraints. If the number of constraints is this many times the number of 540 /// variables, we consider such a system out of line with the intended use 541 /// case of FlatAffineConstraints. 542 // The rationale for 32 is that in the typical simplest of cases, an 543 // identifier is expected to have one lower bound and one upper bound 544 // constraint. With a level of tiling or a connection to another identifier 545 // through a div or mod, an extra pair of bounds gets added. As a limit, we 546 // don't expect an identifier to have more than 32 lower/upper/equality 547 // constraints. This is conservatively set low and can be raised if needed. 548 constexpr static unsigned kExplosionFactor = 32; 549 }; 550 551 /// An extension of FlatAffineConstraints in which dimensions and symbols can 552 /// optionally be associated with an SSA value. 553 class FlatAffineValueConstraints : public FlatAffineConstraints { 554 public: 555 /// Constructs a constraint system reserving memory for the specified number 556 /// of constraints and identifiers. 557 FlatAffineValueConstraints(unsigned numReservedInequalities, 558 unsigned numReservedEqualities, 559 unsigned numReservedCols, unsigned numDims, 560 unsigned numSymbols, unsigned numLocals, 561 ArrayRef<Optional<Value>> valArgs = {}) FlatAffineConstraints(numReservedInequalities,numReservedEqualities,numReservedCols,numDims,numSymbols,numLocals)562 : FlatAffineConstraints(numReservedInequalities, numReservedEqualities, 563 numReservedCols, numDims, numSymbols, numLocals) { 564 assert(numReservedCols >= numIds + 1); 565 assert(valArgs.empty() || valArgs.size() == numIds); 566 values.reserve(numReservedCols); 567 if (valArgs.empty()) 568 values.resize(numIds, None); 569 else 570 values.append(valArgs.begin(), valArgs.end()); 571 } 572 573 /// Constructs a constraint system with the specified number of 574 /// dimensions and symbols. 575 FlatAffineValueConstraints(unsigned numDims = 0, unsigned numSymbols = 0, 576 unsigned numLocals = 0, 577 ArrayRef<Optional<Value>> valArgs = {}) 578 : FlatAffineValueConstraints(/*numReservedInequalities=*/0, 579 /*numReservedEqualities=*/0, 580 /*numReservedCols=*/numDims + numSymbols + 581 numLocals + 1, 582 numDims, numSymbols, numLocals, valArgs) {} 583 584 /// Create a flat affine constraint system from an AffineValueMap or a list of 585 /// these. The constructed system will only include equalities. 586 explicit FlatAffineValueConstraints(const AffineValueMap &avm); 587 explicit FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef); 588 589 /// Creates an affine constraint system from an IntegerSet. 590 explicit FlatAffineValueConstraints(IntegerSet set); 591 592 FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef, 593 IntegerSet set); 594 595 // Construct a hyperrectangular constraint set from ValueRanges that represent 596 // induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are 597 // expected to match one to one. The order of variables and constraints is: 598 // 599 // ivs | lbs | ubs | eq/ineq 600 // ----+-----+-----+--------- 601 // 1 -1 0 >= 0 602 // ----+-----+-----+--------- 603 // -1 0 1 >= 0 604 // 605 // All dimensions as set as DimId. 606 static FlatAffineValueConstraints 607 getHyperrectangular(ValueRange ivs, ValueRange lbs, ValueRange ubs); 608 609 /// Return the kind of this FlatAffineConstraints. getKind()610 Kind getKind() const override { return Kind::FlatAffineValueConstraints; } 611 classof(const FlatAffineConstraints * cst)612 static bool classof(const FlatAffineConstraints *cst) { 613 return cst->getKind() == Kind::FlatAffineValueConstraints; 614 } 615 616 /// Clears any existing data and reserves memory for the specified 617 /// constraints. 618 void reset(unsigned numReservedInequalities, unsigned numReservedEqualities, 619 unsigned numReservedCols, unsigned numDims, unsigned numSymbols, 620 unsigned numLocals = 0) override; 621 void reset(unsigned numReservedInequalities, unsigned numReservedEqualities, 622 unsigned numReservedCols, unsigned numDims, unsigned numSymbols, 623 unsigned numLocals, ArrayRef<Value> valArgs); 624 void reset(unsigned numDims, unsigned numSymbols, unsigned numLocals, 625 ArrayRef<Value> valArgs); 626 using FlatAffineConstraints::reset; 627 628 /// Clones this object. 629 std::unique_ptr<FlatAffineValueConstraints> clone() const; 630 631 /// Adds constraints (lower and upper bounds) for the specified 'affine.for' 632 /// operation's Value using IR information stored in its bound maps. The 633 /// right identifier is first looked up using `forOp`'s Value. Asserts if the 634 /// Value corresponding to the 'affine.for' operation isn't found in the 635 /// constraint system. Returns failure for the yet unimplemented/unsupported 636 /// cases. Any new identifiers that are found in the bound operands of the 637 /// 'affine.for' operation are added as trailing identifiers (either 638 /// dimensional or symbolic depending on whether the operand is a valid 639 /// symbol). 640 // TODO: add support for non-unit strides. 641 LogicalResult addAffineForOpDomain(AffineForOp forOp); 642 643 /// Adds constraints (lower and upper bounds) for each loop in the loop nest 644 /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice. 645 /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in 646 /// the nest, sorted outer-to-inner. `operands` contains the bound operands 647 /// for a single bound map. All the bound maps will use the same bound 648 /// operands. Note that some loops described by a computation slice might not 649 /// exist yet in the IR so the Value attached to those dimension identifiers 650 /// might be empty. For that reason, this method doesn't perform Value 651 /// look-ups to retrieve the dimension identifier positions. Instead, it 652 /// assumes the position of the dim identifiers in the constraint system is 653 /// the same as the position of the loop in the loop nest. 654 LogicalResult addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps, 655 ArrayRef<AffineMap> ubMaps, 656 ArrayRef<Value> operands); 657 658 /// Adds constraints imposed by the `affine.if` operation. These constraints 659 /// are collected from the IntegerSet attached to the given `affine.if` 660 /// instance argument (`ifOp`). It is asserted that: 661 /// 1) The IntegerSet of the given `affine.if` instance should not contain 662 /// semi-affine expressions, 663 /// 2) The columns of the constraint system created from `ifOp` should match 664 /// the columns in the current one regarding numbers and values. 665 void addAffineIfOpDomain(AffineIfOp ifOp); 666 667 /// Adds a bound for the identifier at the specified position with constraints 668 /// being drawn from the specified bound map and operands. In case of an 669 /// EQ bound, the bound map is expected to have exactly one result. In case 670 /// of a LB/UB, the bound map may have more than one result, for each of which 671 /// an inequality is added. 672 LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, 673 ValueRange operands); 674 675 /// Adds a constant bound for the identifier associated with the given Value. 676 void addBound(BoundType type, Value val, int64_t value); 677 678 using FlatAffineConstraints::addBound; 679 680 /// Returns the bound for the identifier at `pos` from the inequality at 681 /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned 682 /// affine value map can either be a lower bound or an upper bound depending 683 /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does 684 /// not involve the `pos`th identifier. 685 void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos, 686 AffineValueMap &vmap, 687 MLIRContext *context) const; 688 689 /// Adds slice lower bounds represented by lower bounds in `lbMaps` and upper 690 /// bounds in `ubMaps` to each identifier in the constraint system which has 691 /// a value in `values`. Note that both lower/upper bounds share the same 692 /// operand list `operands`. 693 /// This function assumes `values.size` == `lbMaps.size` == `ubMaps.size`. 694 /// Note that both lower/upper bounds use operands from `operands`. 695 LogicalResult addSliceBounds(ArrayRef<Value> values, 696 ArrayRef<AffineMap> lbMaps, 697 ArrayRef<AffineMap> ubMaps, 698 ArrayRef<Value> operands); 699 700 /// Looks up the position of the identifier with the specified Value. Returns 701 /// true if found (false otherwise). `pos` is set to the (column) position of 702 /// the identifier. 703 bool findId(Value val, unsigned *pos) const; 704 705 /// Returns true if an identifier with the specified Value exists, false 706 /// otherwise. 707 bool containsId(Value val) const; 708 709 /// Swap the posA^th identifier with the posB^th identifier. 710 void swapId(unsigned posA, unsigned posB) override; 711 712 /// Insert identifiers of the specified kind at position `pos`. Positions are 713 /// relative to the kind of identifier. The coefficient columns corresponding 714 /// to the added identifiers are initialized to zero. `vals` are the Values 715 /// corresponding to the identifiers. Return the absolute column position 716 /// (i.e., not relative to the kind of identifier) of the first added 717 /// identifier. 718 /// 719 /// Note: Empty Values are allowed in `vals`. 720 unsigned insertDimId(unsigned pos, ValueRange vals); 721 using FlatAffineConstraints::insertDimId; 722 unsigned insertSymbolId(unsigned pos, ValueRange vals); 723 using FlatAffineConstraints::insertSymbolId; 724 unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override; 725 unsigned insertId(IdKind kind, unsigned pos, ValueRange vals); 726 727 /// Append identifiers of the specified kind after the last identifier of that 728 /// kind. The coefficient columns corresponding to the added identifiers are 729 /// initialized to zero. `vals` are the Values corresponding to the 730 /// identifiers. Return the position of the first added column. 731 /// 732 /// Note: Empty Values are allowed in `vals`. 733 unsigned appendDimId(ValueRange vals); 734 using FlatAffineConstraints::appendDimId; 735 unsigned appendSymbolId(ValueRange vals); 736 using FlatAffineConstraints::appendSymbolId; 737 738 /// Add the specified values as a dim or symbol id depending on its nature, if 739 /// it already doesn't exist in the system. `val` has to be either a terminal 740 /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any 741 /// symbols or loop IVs. The identifier is added to the end of the existing 742 /// dims or symbols. Additional information on the identifier is extracted 743 /// from the IR and added to the constraint system. 744 void addInductionVarOrTerminalSymbol(Value val); 745 746 /// Align `map` with this constraint system based on `operands`. Each operand 747 /// must already have a corresponding dim/symbol in this constraint system. 748 AffineMap computeAlignedMap(AffineMap map, ValueRange operands) const; 749 750 /// Composes the affine value map with this FlatAffineValueConstrains, adding 751 /// the results of the map as dimensions at the front 752 /// [0, vMap->getNumResults()) and with the dimensions set to the equalities 753 /// specified by the value map. 754 /// 755 /// Returns failure if the composition fails (when vMap is a semi-affine map). 756 /// The vMap's operand Value's are used to look up the right positions in 757 /// the FlatAffineConstraints with which to associate. Every operand of vMap 758 /// should have a matching dim/symbol column in this constraint system (with 759 /// the same associated Value). 760 LogicalResult composeMap(const AffineValueMap *vMap); 761 762 /// Projects out the identifier that is associate with Value. 763 void projectOut(Value val); 764 using FlatAffineConstraints::projectOut; 765 766 /// Changes all symbol identifiers which are loop IVs to dim identifiers. 767 void convertLoopIVSymbolsToDims(); 768 769 /// Updates the constraints to be the smallest bounding (enclosing) box that 770 /// contains the points of `this` set and that of `other`, with the symbols 771 /// being treated specially. For each of the dimensions, the min of the lower 772 /// bounds (symbolic) and the max of the upper bounds (symbolic) is computed 773 /// to determine such a bounding box. `other` is expected to have the same 774 /// dimensional identifiers as this constraint system (in the same order). 775 /// 776 /// E.g.: 777 /// 1) this = {0 <= d0 <= 127}, 778 /// other = {16 <= d0 <= 192}, 779 /// output = {0 <= d0 <= 192} 780 /// 2) this = {s0 + 5 <= d0 <= s0 + 20}, 781 /// other = {s0 + 1 <= d0 <= s0 + 9}, 782 /// output = {s0 + 1 <= d0 <= s0 + 20} 783 /// 3) this = {0 <= d0 <= 5, 1 <= d1 <= 9} 784 /// other = {2 <= d0 <= 6, 5 <= d1 <= 15}, 785 /// output = {0 <= d0 <= 6, 1 <= d1 <= 15} 786 LogicalResult unionBoundingBox(const FlatAffineValueConstraints &other); 787 using FlatAffineConstraints::unionBoundingBox; 788 789 /// Merge and align the identifiers of `this` and `other` starting at 790 /// `offset`, so that both constraint systems get the union of the contained 791 /// identifiers that is dimension-wise and symbol-wise unique; both 792 /// constraint systems are updated so that they have the union of all 793 /// identifiers, with `this`'s original identifiers appearing first followed 794 /// by any of `other`'s identifiers that didn't appear in `this`. Local 795 /// identifiers of each system are by design separate/local and are placed 796 /// one after other (`this`'s followed by `other`'s). 797 // E.g.: Input: `this` has (%i, %j) [%M, %N] 798 // `other` has (%k, %j) [%P, %N, %M] 799 // Output: both `this`, `other` have (%i, %j, %k) [%M, %N, %P] 800 // 801 void mergeAndAlignIdsWithOther(unsigned offset, 802 FlatAffineValueConstraints *other); 803 804 /// Returns true if this constraint system and `other` are in the same 805 /// space, i.e., if they are associated with the same set of identifiers, 806 /// appearing in the same order. Returns false otherwise. 807 bool areIdsAlignedWithOther(const FlatAffineValueConstraints &other); 808 809 /// Replaces the contents of this FlatAffineValueConstraints with `other`. 810 void clearAndCopyFrom(const FlatAffineConstraints &other) override; 811 812 /// Returns the Value associated with the pos^th identifier. Asserts if 813 /// no Value identifier was associated. getValue(unsigned pos)814 inline Value getValue(unsigned pos) const { 815 assert(hasValue(pos) && "identifier's Value not set"); 816 return values[pos].getValue(); 817 } 818 819 /// Returns true if the pos^th identifier has an associated Value. hasValue(unsigned pos)820 inline bool hasValue(unsigned pos) const { return values[pos].hasValue(); } 821 822 /// Returns true if at least one identifier has an associated Value. 823 bool hasValues() const; 824 825 /// Returns the Values associated with identifiers in range [start, end). 826 /// Asserts if no Value was associated with one of these identifiers. getValues(unsigned start,unsigned end,SmallVectorImpl<Value> * values)827 inline void getValues(unsigned start, unsigned end, 828 SmallVectorImpl<Value> *values) const { 829 assert((start < numIds || start == end) && "invalid start position"); 830 assert(end <= numIds && "invalid end position"); 831 values->clear(); 832 values->reserve(end - start); 833 for (unsigned i = start; i < end; i++) 834 values->push_back(getValue(i)); 835 } getAllValues(SmallVectorImpl<Value> * values)836 inline void getAllValues(SmallVectorImpl<Value> *values) const { 837 getValues(0, numIds, values); 838 } 839 getMaybeValues()840 inline ArrayRef<Optional<Value>> getMaybeValues() const { 841 return {values.data(), values.size()}; 842 } 843 getMaybeDimValues()844 inline ArrayRef<Optional<Value>> getMaybeDimValues() const { 845 return {values.data(), getNumDimIds()}; 846 } 847 getMaybeSymbolValues()848 inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const { 849 return {values.data() + getNumDimIds(), getNumSymbolIds()}; 850 } 851 getMaybeDimAndSymbolValues()852 inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const { 853 return {values.data(), getNumDimIds() + getNumSymbolIds()}; 854 } 855 856 /// Sets the Value associated with the pos^th identifier. setValue(unsigned pos,Value val)857 inline void setValue(unsigned pos, Value val) { 858 assert(pos < numIds && "invalid id position"); 859 values[pos] = val; 860 } 861 862 /// Sets the Values associated with the identifiers in the range [start, end). setValues(unsigned start,unsigned end,ArrayRef<Value> values)863 void setValues(unsigned start, unsigned end, ArrayRef<Value> values) { 864 assert((start < numIds || end == start) && "invalid start position"); 865 assert(end <= numIds && "invalid end position"); 866 assert(values.size() == end - start); 867 for (unsigned i = start; i < end; ++i) 868 setValue(i, values[i - start]); 869 } 870 871 /// Merge and align symbols of `this` and `other` such that both get union of 872 /// of symbols that are unique. Symbols with Value as `None` are considered 873 /// to be inequal to all other symbols. 874 void mergeSymbolIds(FlatAffineValueConstraints &other); 875 876 protected: 877 /// Returns false if the fields corresponding to various identifier counts, or 878 /// equality/inequality buffer sizes aren't consistent; true otherwise. This 879 /// is meant to be used within an assert internally. 880 bool hasConsistentState() const override; 881 882 /// Removes identifiers in the column range [idStart, idLimit), and copies any 883 /// remaining valid data into place, updates member variables, and resizes 884 /// arrays as needed. 885 void removeIdRange(unsigned idStart, unsigned idLimit) override; 886 887 /// Eliminates the identifier at the specified position using Fourier-Motzkin 888 /// variable elimination, but uses Gaussian elimination if there is an 889 /// equality involving that identifier. If the result of the elimination is 890 /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is 891 /// set to true, a potential under approximation (subset) of the rational 892 /// shadow / exact integer shadow is computed. 893 // See implementation comments for more details. 894 void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false, 895 bool *isResultIntegerExact = nullptr) override; 896 897 /// Values corresponding to the (column) identifiers of this constraint 898 /// system appearing in the order the identifiers correspond to columns. 899 /// Temporary ones or those that aren't associated with any Value are set to 900 /// None. 901 SmallVector<Optional<Value>, 8> values; 902 }; 903 904 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the 905 /// dimensions, symbols, and additional variables that represent floor divisions 906 /// of dimensions, symbols, and in turn other floor divisions. Returns failure 907 /// if 'expr' could not be flattened (i.e., semi-affine is not yet handled). 908 /// 'cst' contains constraints that connect newly introduced local identifiers 909 /// to existing dimensional and symbolic identifiers. See documentation for 910 /// AffineExprFlattener on how mod's and div's are flattened. 911 LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, 912 unsigned numSymbols, 913 SmallVectorImpl<int64_t> *flattenedExpr, 914 FlatAffineConstraints *cst = nullptr); 915 916 /// Flattens the result expressions of the map to their corresponding flattened 917 /// forms and set in 'flattenedExprs'. Returns failure if any expression in the 918 /// map could not be flattened (i.e., semi-affine is not yet handled). 'cst' 919 /// contains constraints that connect newly introduced local identifiers to 920 /// existing dimensional and / symbolic identifiers. See documentation for 921 /// AffineExprFlattener on how mod's and div's are flattened. For all affine 922 /// expressions that share the same operands (like those of an affine map), this 923 /// method should be used instead of repeatedly calling getFlattenedAffineExpr 924 /// since local variables added to deal with div's and mod's will be reused 925 /// across expressions. 926 LogicalResult 927 getFlattenedAffineExprs(AffineMap map, 928 std::vector<SmallVector<int64_t, 8>> *flattenedExprs, 929 FlatAffineConstraints *cst = nullptr); 930 LogicalResult 931 getFlattenedAffineExprs(IntegerSet set, 932 std::vector<SmallVector<int64_t, 8>> *flattenedExprs, 933 FlatAffineConstraints *cst = nullptr); 934 935 /// Re-indexes the dimensions and symbols of an affine map with given `operands` 936 /// values to align with `dims` and `syms` values. 937 /// 938 /// Each dimension/symbol of the map, bound to an operand `o`, is replaced with 939 /// dimension `i`, where `i` is the position of `o` within `dims`. If `o` is not 940 /// in `dims`, replace it with symbol `i`, where `i` is the position of `o` 941 /// within `syms`. If `o` is not in `syms` either, replace it with a new symbol. 942 /// 943 /// Note: If a value appears multiple times as a dimension/symbol (or both), all 944 /// corresponding dim/sym expressions are replaced with the first dimension 945 /// bound to that value (or first symbol if no such dimension exists). 946 /// 947 /// The resulting affine map has `dims.size()` many dimensions and at least 948 /// `syms.size()` many symbols. 949 /// 950 /// The SSA values of the symbols of the resulting map are optionally returned 951 /// via `newSyms`. This is a concatenation of `syms` with the SSA values of the 952 /// newly added symbols. 953 /// 954 /// Note: As part of this re-indexing, dimensions may turn into symbols, or vice 955 /// versa. 956 AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands, 957 ValueRange dims, ValueRange syms, 958 SmallVector<Value> *newSyms = nullptr); 959 960 } // end namespace mlir. 961 962 #endif // MLIR_ANALYSIS_AFFINESTRUCTURES_H 963