diff --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h --- a/mlir/include/mlir/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Analysis/AffineStructures.h @@ -573,16 +573,21 @@ numLocals + 1, numDims, numSymbols, numLocals, valArgs) {} - /// Create a flat affine constraint system from an AffineValueMap or a list of - /// these. The constructed system will only include equalities. - explicit FlatAffineValueConstraints(const AffineValueMap &avm); - explicit FlatAffineValueConstraints(ArrayRef avmRef); - /// Creates an affine constraint system from an IntegerSet. explicit FlatAffineValueConstraints(IntegerSet set); - FlatAffineValueConstraints(ArrayRef avmRef, - IntegerSet set); + // Construct a hyperrectangular constraint set from ValueRanges that represent + // induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are + // expected to match one to one. The order of variables and constraints is: + // + // ivs | lbs | ubs | eq/ineq + // ----+-----+-----+--------- + // 1 -1 0 >= 0 + // ----+-----+-----+--------- + // -1 0 1 >= 0 + // + // All dimensions as set as DimId. + FlatAffineValueConstraints(ValueRange ivs, ValueRange lbs, ValueRange ubs); /// Return the kind of this FlatAffineConstraints. Kind getKind() const override { return Kind::FlatAffineValueConstraints; } @@ -608,8 +613,8 @@ /// Adds constraints (lower and upper bounds) for the specified 'affine.for' /// operation's Value using IR information stored in its bound maps. The - /// right identifier is first looked up using `forOp`'s Value. Asserts if the - /// Value corresponding to the 'affine.for' operation isn't found in the + /// right identifier is first looked up using `forOp`'s Value. Asserts if + /// the Value corresponding to the 'affine.for' operation isn't found in the /// constraint system. Returns failure for the yet unimplemented/unsupported /// cases. Any new identifiers that are found in the bound operands of the /// 'affine.for' operation are added as trailing identifiers (either @@ -619,16 +624,16 @@ LogicalResult addAffineForOpDomain(AffineForOp forOp); /// Adds constraints (lower and upper bounds) for each loop in the loop nest - /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice. - /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in - /// the nest, sorted outer-to-inner. `operands` contains the bound operands - /// for a single bound map. All the bound maps will use the same bound - /// operands. Note that some loops described by a computation slice might not - /// exist yet in the IR so the Value attached to those dimension identifiers - /// might be empty. For that reason, this method doesn't perform Value - /// look-ups to retrieve the dimension identifier positions. Instead, it - /// assumes the position of the dim identifiers in the constraint system is - /// the same as the position of the loop in the loop nest. + /// described by the bound maps `lbMaps` and `ubMaps` of a computation + /// slice. Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a + /// loop in the nest, sorted outer-to-inner. `operands` contains the bound + /// operands for a single bound map. All the bound maps will use the same + /// bound operands. Note that some loops described by a computation slice + /// might not exist yet in the IR so the Value attached to those dimension + /// identifiers might be empty. For that reason, this method doesn't perform + /// Value look-ups to retrieve the dimension identifier positions. Instead, + /// it assumes the position of the dim identifiers in the constraint system + /// is the same as the position of the loop in the loop nest. LogicalResult addDomainFromSliceMaps(ArrayRef lbMaps, ArrayRef ubMaps, ArrayRef operands); @@ -642,42 +647,43 @@ /// the columns in the current one regarding numbers and values. void addAffineIfOpDomain(AffineIfOp ifOp); - /// Adds a bound for the identifier at the specified position with constraints - /// being drawn from the specified bound map and operands. In case of an - /// EQ bound, the bound map is expected to have exactly one result. In case - /// of a LB/UB, the bound map may have more than one result, for each of which - /// an inequality is added. + /// Adds a bound for the identifier at the specified position with + /// constraints being drawn from the specified bound map and operands. In + /// case of an EQ bound, the bound map is expected to have exactly one + /// result. In case of a LB/UB, the bound map may have more than one result, + /// for each of which an inequality is added. LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, ValueRange operands); - /// Adds a constant bound for the identifier associated with the given Value. + /// Adds a constant bound for the identifier associated with the given + /// Value. void addBound(BoundType type, Value val, int64_t value); using FlatAffineConstraints::addBound; /// Returns the bound for the identifier at `pos` from the inequality at - /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned - /// affine value map can either be a lower bound or an upper bound depending - /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does - /// not involve the `pos`th identifier. + /// `ineqPos` as a 1-d affine value map (affine map + operands). The + /// returned affine value map can either be a lower bound or an upper bound + /// depending on the sign of atIneq(ineqPos, pos). Asserts if the row at + /// `ineqPos` does not involve the `pos`th identifier. void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos, AffineValueMap &vmap, MLIRContext *context) const; - /// Adds slice lower bounds represented by lower bounds in `lbMaps` and upper - /// bounds in `ubMaps` to each identifier in the constraint system which has - /// a value in `values`. Note that both lower/upper bounds share the same - /// operand list `operands`. - /// This function assumes `values.size` == `lbMaps.size` == `ubMaps.size`. - /// Note that both lower/upper bounds use operands from `operands`. + /// Adds slice lower bounds represented by lower bounds in `lbMaps` and + /// upper bounds in `ubMaps` to each identifier in the constraint system + /// which has a value in `values`. Note that both lower/upper bounds share + /// the same operand list `operands`. This function assumes `values.size` == + /// `lbMaps.size` == `ubMaps.size`. Note that both lower/upper bounds use + /// operands from `operands`. LogicalResult addSliceBounds(ArrayRef values, ArrayRef lbMaps, ArrayRef ubMaps, ArrayRef operands); - /// Looks up the position of the identifier with the specified Value. Returns - /// true if found (false otherwise). `pos` is set to the (column) position of - /// the identifier. + /// Looks up the position of the identifier with the specified Value. + /// Returns true if found (false otherwise). `pos` is set to the (column) + /// position of the identifier. bool findId(Value val, unsigned *pos) const; /// Returns true if an identifier with the specified Value exists, false @@ -687,12 +693,12 @@ /// Swap the posA^th identifier with the posB^th identifier. void swapId(unsigned posA, unsigned posB) override; - /// Insert identifiers of the specified kind at position `pos`. Positions are - /// relative to the kind of identifier. The coefficient columns corresponding - /// to the added identifiers are initialized to zero. `vals` are the Values - /// corresponding to the identifiers. Return the absolute column position - /// (i.e., not relative to the kind of identifier) of the first added - /// identifier. + /// Insert identifiers of the specified kind at position `pos`. Positions + /// are relative to the kind of identifier. The coefficient columns + /// corresponding to the added identifiers are initialized to zero. `vals` + /// are the Values corresponding to the identifiers. Return the absolute + /// column position (i.e., not relative to the kind of identifier) of the + /// first added identifier. /// /// Note: Empty Values are allowed in `vals`. unsigned insertDimId(unsigned pos, ValueRange vals); @@ -702,10 +708,10 @@ unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override; unsigned insertId(IdKind kind, unsigned pos, ValueRange vals); - /// Append identifiers of the specified kind after the last identifier of that - /// kind. The coefficient columns corresponding to the added identifiers are - /// initialized to zero. `vals` are the Values corresponding to the - /// identifiers. Return the position of the first added column. + /// Append identifiers of the specified kind after the last identifier of + /// that kind. The coefficient columns corresponding to the added + /// identifiers are initialized to zero. `vals` are the Values corresponding + /// to the identifiers. Return the position of the first added column. /// /// Note: Empty Values are allowed in `vals`. unsigned appendDimId(ValueRange vals); @@ -713,28 +719,29 @@ unsigned appendSymbolId(ValueRange vals); using FlatAffineConstraints::appendSymbolId; - /// Add the specified values as a dim or symbol id depending on its nature, if - /// it already doesn't exist in the system. `val` has to be either a terminal - /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any - /// symbols or loop IVs. The identifier is added to the end of the existing - /// dims or symbols. Additional information on the identifier is extracted - /// from the IR and added to the constraint system. + /// Add the specified values as a dim or symbol id depending on its nature, + /// if it already doesn't exist in the system. `val` has to be either a + /// terminal symbol or a loop IV, i.e., it cannot be the result affine.apply + /// of any symbols or loop IVs. The identifier is added to the end of the + /// existing dims or symbols. Additional information on the identifier is + /// extracted from the IR and added to the constraint system. void addInductionVarOrTerminalSymbol(Value val); - /// Align `map` with this constraint system based on `operands`. Each operand - /// must already have a corresponding dim/symbol in this constraint system. + /// Align `map` with this constraint system based on `operands`. Each + /// operand must already have a corresponding dim/symbol in this constraint + /// system. AffineMap computeAlignedMap(AffineMap map, ValueRange operands) const; - /// Composes the affine value map with this FlatAffineValueConstrains, adding - /// the results of the map as dimensions at the front - /// [0, vMap->getNumResults()) and with the dimensions set to the equalities + /// Composes the affine value map with this FlatAffineValueConstrains, + /// adding the results of the map as dimensions at the front [0, + /// vMap->getNumResults()) and with the dimensions set to the equalities /// specified by the value map. /// - /// Returns failure if the composition fails (when vMap is a semi-affine map). - /// The vMap's operand Value's are used to look up the right positions in - /// the FlatAffineConstraints with which to associate. Every operand of vMap - /// should have a matching dim/symbol column in this constraint system (with - /// the same associated Value). + /// Returns failure if the composition fails (when vMap is a semi-affine + /// map). The vMap's operand Value's are used to look up the right positions + /// in the FlatAffineConstraints with which to associate. Every operand of + /// vMap should have a matching dim/symbol column in this constraint system + /// (with the same associated Value). LogicalResult composeMap(const AffineValueMap *vMap); /// Projects out the identifier that is associate with Value. @@ -746,10 +753,11 @@ /// Updates the constraints to be the smallest bounding (enclosing) box that /// contains the points of `this` set and that of `other`, with the symbols - /// being treated specially. For each of the dimensions, the min of the lower - /// bounds (symbolic) and the max of the upper bounds (symbolic) is computed - /// to determine such a bounding box. `other` is expected to have the same - /// dimensional identifiers as this constraint system (in the same order). + /// being treated specially. For each of the dimensions, the min of the + /// lower bounds (symbolic) and the max of the upper bounds (symbolic) is + /// computed to determine such a bounding box. `other` is expected to have + /// the same dimensional identifiers as this constraint system (in the same + /// order). /// /// E.g.: /// 1) this = {0 <= d0 <= 127}, @@ -837,7 +845,8 @@ values[pos] = val; } - /// Sets the Values associated with the identifiers in the range [start, end). + /// Sets the Values associated with the identifiers in the range [start, + /// end). void setValues(unsigned start, unsigned end, ArrayRef values) { assert((start < numIds || end == start) && "invalid start position"); assert(end <= numIds && "invalid end position"); @@ -846,28 +855,28 @@ setValue(i, values[i - start]); } - /// Merge and align symbols of `this` and `other` such that both get union of - /// of symbols that are unique. Symbols with Value as `None` are considered - /// to be inequal to all other symbols. + /// Merge and align symbols of `this` and `other` such that both get union + /// of of symbols that are unique. Symbols with Value as `None` are + /// considered to be inequal to all other symbols. void mergeSymbolIds(FlatAffineValueConstraints &other); protected: - /// Returns false if the fields corresponding to various identifier counts, or - /// equality/inequality buffer sizes aren't consistent; true otherwise. This - /// is meant to be used within an assert internally. + /// Returns false if the fields corresponding to various identifier counts, + /// or equality/inequality buffer sizes aren't consistent; true otherwise. + /// This is meant to be used within an assert internally. bool hasConsistentState() const override; - /// Removes identifiers in the column range [idStart, idLimit), and copies any - /// remaining valid data into place, updates member variables, and resizes - /// arrays as needed. + /// Removes identifiers in the column range [idStart, idLimit), and copies + /// any remaining valid data into place, updates member variables, and + /// resizes arrays as needed. void removeIdRange(unsigned idStart, unsigned idLimit) override; - /// Eliminates the identifier at the specified position using Fourier-Motzkin - /// variable elimination, but uses Gaussian elimination if there is an - /// equality involving that identifier. If the result of the elimination is - /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is - /// set to true, a potential under approximation (subset) of the rational - /// shadow / exact integer shadow is computed. + /// Eliminates the identifier at the specified position using + /// Fourier-Motzkin variable elimination, but uses Gaussian elimination if + /// there is an equality involving that identifier. If the result of the + /// elimination is integer exact, `*isResultIntegerExact` is set to true. If + /// `darkShadow` is set to true, a potential under approximation (subset) of + /// the rational shadow / exact integer shadow is computed. // See implementation comments for more details. void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false, bool *isResultIntegerExact = nullptr) override; @@ -879,28 +888,29 @@ SmallVector, 8> values; }; -/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the -/// dimensions, symbols, and additional variables that represent floor divisions -/// of dimensions, symbols, and in turn other floor divisions. Returns failure -/// if 'expr' could not be flattened (i.e., semi-affine is not yet handled). -/// 'cst' contains constraints that connect newly introduced local identifiers -/// to existing dimensional and symbolic identifiers. See documentation for -/// AffineExprFlattener on how mod's and div's are flattened. +/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of +/// the dimensions, symbols, and additional variables that represent floor +/// divisions of dimensions, symbols, and in turn other floor divisions. +/// Returns failure if 'expr' could not be flattened (i.e., semi-affine is not +/// yet handled). 'cst' contains constraints that connect newly introduced +/// local identifiers to existing dimensional and symbolic identifiers. See +/// documentation for AffineExprFlattener on how mod's and div's are +/// flattened. LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, SmallVectorImpl *flattenedExpr, FlatAffineConstraints *cst = nullptr); -/// Flattens the result expressions of the map to their corresponding flattened -/// forms and set in 'flattenedExprs'. Returns failure if any expression in the -/// map could not be flattened (i.e., semi-affine is not yet handled). 'cst' -/// contains constraints that connect newly introduced local identifiers to -/// existing dimensional and / symbolic identifiers. See documentation for -/// AffineExprFlattener on how mod's and div's are flattened. For all affine -/// expressions that share the same operands (like those of an affine map), this -/// method should be used instead of repeatedly calling getFlattenedAffineExpr -/// since local variables added to deal with div's and mod's will be reused -/// across expressions. +/// Flattens the result expressions of the map to their corresponding +/// flattened forms and set in 'flattenedExprs'. Returns failure if any +/// expression in the map could not be flattened (i.e., semi-affine is not yet +/// handled). 'cst' contains constraints that connect newly introduced local +/// identifiers to existing dimensional and / symbolic identifiers. See +/// documentation for AffineExprFlattener on how mod's and div's are +/// flattened. For all affine expressions that share the same operands (like +/// those of an affine map), this method should be used instead of repeatedly +/// calling getFlattenedAffineExpr since local variables added to deal with +/// div's and mod's will be reused across expressions. LogicalResult getFlattenedAffineExprs(AffineMap map, std::vector> *flattenedExprs, @@ -910,27 +920,29 @@ std::vector> *flattenedExprs, FlatAffineConstraints *cst = nullptr); -/// Re-indexes the dimensions and symbols of an affine map with given `operands` -/// values to align with `dims` and `syms` values. +/// Re-indexes the dimensions and symbols of an affine map with given +/// `operands` values to align with `dims` and `syms` values. /// -/// Each dimension/symbol of the map, bound to an operand `o`, is replaced with -/// dimension `i`, where `i` is the position of `o` within `dims`. If `o` is not -/// in `dims`, replace it with symbol `i`, where `i` is the position of `o` -/// within `syms`. If `o` is not in `syms` either, replace it with a new symbol. +/// Each dimension/symbol of the map, bound to an operand `o`, is replaced +/// with dimension `i`, where `i` is the position of `o` within `dims`. If `o` +/// is not in `dims`, replace it with symbol `i`, where `i` is the position of +/// `o` within `syms`. If `o` is not in `syms` either, replace it with a new +/// symbol. /// -/// Note: If a value appears multiple times as a dimension/symbol (or both), all -/// corresponding dim/sym expressions are replaced with the first dimension -/// bound to that value (or first symbol if no such dimension exists). +/// Note: If a value appears multiple times as a dimension/symbol (or both), +/// all corresponding dim/sym expressions are replaced with the first +/// dimension bound to that value (or first symbol if no such dimension +/// exists). /// /// The resulting affine map has `dims.size()` many dimensions and at least /// `syms.size()` many symbols. /// /// The SSA values of the symbols of the resulting map are optionally returned -/// via `newSyms`. This is a concatenation of `syms` with the SSA values of the -/// newly added symbols. +/// via `newSyms`. This is a concatenation of `syms` with the SSA values of +/// the newly added symbols. /// -/// Note: As part of this re-indexing, dimensions may turn into symbols, or vice -/// versa. +/// Note: As part of this re-indexing, dimensions may turn into symbols, or +/// vice versa. AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands, ValueRange dims, ValueRange syms, SmallVector *newSyms = nullptr); diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -189,6 +189,46 @@ values.resize(numIds, None); } +// Construct a hyperrectangular constraint set from ValueRanges that represent +// induction variables, lower and upper bounds. `ivs`, `lbs` and `ubs` are +// expected to match one to one. The order of variables and constraints is: +// +// ivs | lbs | ubs | eq/ineq +// ----+-----+-----+--------- +// 1 -1 0 >= 0 +// ----+-----+-----+--------- +// -1 0 1 >= 0 +// +// All dimensions as set as DimId. +FlatAffineValueConstraints::FlatAffineValueConstraints(ValueRange ivs, + ValueRange lbs, + ValueRange ubs) { + int nIvs = ivs.size(); + assert(nIvs == lbs.size() && "expected as may lower bounds that ivs"); + assert(nIvs == ubs.size() && "expected as may upper bounds that ivs"); + + if (nIvs == 0) + return; + + insertDimId(/*pos=*/0, ivs); + insertDimId(/*pos=*/nIvs, lbs); + insertDimId(/*pos=*/2 * nIvs, ubs); + + MLIRContext *ctx = ivs.front().getContext(); + for (int ivIdx = 0, e = nIvs; ivIdx < e; ++ivIdx) { + // iv - lb >= 0 + AffineMap lb = + AffineMap::get(3 * nIvs, 0, getAffineDimExpr(nIvs + ivIdx, ctx)); + if (failed(addBound(BoundType::LB, ivIdx, lb))) + llvm_unreachable("Unexpected FlatAffineValueConstraints creation error"); + // -iv + ub >= 0 + AffineMap ub = + AffineMap::get(3 * nIvs, 0, getAffineDimExpr(2 * nIvs + ivIdx, ctx)); + if (failed(addBound(BoundType::UB, ivIdx, ub))) + llvm_unreachable("Unexpected FlatAffineValueConstraints creation error"); + } +} + void FlatAffineConstraints::reset(unsigned numReservedInequalities, unsigned numReservedEqualities, unsigned newNumReservedCols, diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp @@ -178,43 +178,6 @@ valid = true; } -/// Given a set of loops, assumed to be scf::ForOp, create a constraint set -/// containing the inequalities `iv - lb >= 0` and `-iv + ub - 1 >= 0` for each -/// loop. The order of the constraints follows: -/// -/// ivs | lbs | ubs | eq/ineq -/// ----+-----+-----+--------- -/// 1 -1 0 >= 0 -/// ----+-----+-----+--------- -/// -1 0 1 >= 0 -/// -static FlatAffineValueConstraints -initLoopIvsAndBounds(ArrayRef loops) { - FlatAffineValueConstraints constraints; - // Append dims for all ivs, lbs, ubs: the order is important. - for (scf::ForOp op : loops) - constraints.appendDimId(op.getInductionVar()); - for (scf::ForOp op : loops) - constraints.appendDimId(op.lowerBound()); - for (scf::ForOp op : loops) - constraints.appendDimId(op.upperBound()); - int numLoops = loops.size(); - for (int ivIdx = 0, e = numLoops; ivIdx < e; ++ivIdx) { - // iv - lb >= 0 - SmallVector ineqLb(constraints.getNumCols(), 0); - ineqLb[ivIdx] = 1; - ineqLb[ivIdx + numLoops] = -1; - // -iv + ub >= 0 - SmallVector ineqUb(constraints.getNumCols(), 0); - ineqUb[ivIdx] = -1; - ineqUb[ivIdx + 2 * numLoops] = 1; - ineqUb[constraints.getNumCols() - 1] = -1; - constraints.addInequality(ineqLb); - constraints.addInequality(ineqUb); - } - return constraints; -} - static bool isDefinedOutsideOrConstant(scf::ForOp outer, Value v) { return outer.isDefinedOutsideOfLoop(v) || v.getDefiningOp(); } @@ -306,8 +269,16 @@ // `backwardSlice`. FailureOr> HoistingAnalysis::getPackedTensorSizes(ImplicitLocOpBuilder &b) { - FlatAffineValueConstraints constraints = - initLoopIvsAndBounds(packingLoops.getArrayRef()); + // Create the base affine constaints for the packedLoops. + FlatAffineValueConstraints constraints( + llvm::to_vector<8>(llvm::map_range( + packingLoops, [](scf::ForOp op) { return op.getInductionVar(); })), + llvm::to_vector<8>(llvm::map_range( + packingLoops, [](scf::ForOp op) { return op.lowerBound(); })), + llvm::to_vector<8>(llvm::map_range( + packingLoops, [](scf::ForOp op) { return op.upperBound(); }))); + + // Iteratively try to fold the upper bounds into the constraints set. if (failed(foldUpperBoundsIntoConstraintsSet( constraints, outermostEnclosingForOp, packingLoops.getArrayRef()))) return failure();