diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h --- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h @@ -23,6 +23,11 @@ #include namespace mlir { +class AffineExpr; +class AffineMap; +class IntegerSet; +class MLIRContext; + namespace presburger { class IntegerRelation; @@ -851,9 +856,103 @@ return cst->getKind() == Kind::IntegerPolyhedron; } + /// Adds a bound for the variable at the specified position with constraints + /// being drawn from the specified bound map. 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. + /// + /// The bound can be added as open or closed by specifying isClosedBound. In + /// case of a LB/UB, isClosedBound = false means the bound is added internally + /// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can + /// only be added as a closed bound. + /// + /// Note: The dimensions/symbols of this IntegerRelation must match the + /// dimensions/symbols of the affine map. + LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, + bool isClosedBound); + + /// Adds a bound for the variable at the specified position with constraints + /// being drawn from the specified bound map. 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. + /// + /// Note: The dimensions/symbols of this IntegerRelation must match the + /// dimensions/symbols of the affine map. A lower bound is assumed to be + /// closed and an upper bound is assumed to be open. + LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap); + + /// The `addBound` overload above hides the inherited overloads by default, so + /// we explicitly introduce them here. + using IntegerRelation::addBound; + // Clones this object. std::unique_ptr clone() const; + /// Check if the pos^th variable can be expressed as a floordiv of an affine + /// function of other variables (where the divisor is a positive constant) + /// given the initial set of expressions in `exprs`. If it can be, the + /// corresponding position in `exprs` is set as the detected affine expr. For + /// eg: 4q <= i + j <= 4q + 3 <=> q = (i + j) floordiv 4. An equality can + /// also yield a floordiv: eg. 4q = i + j <=> q = (i + j) floordiv 4. 32q + + /// 28 + /// <= i <= 32q + 31 => q = i floordiv 32. + bool detectAsFloorDiv(unsigned pos, MLIRContext *context, + SmallVectorImpl &exprs) const; + + /// Determine whether the variable at 'pos' (say var_r) can be expressed as + /// modulo of another known variable (say var_n) w.r.t a constant. For + /// example, if the following constraints hold true: + /// ``` + /// 0 <= var_r <= divisor - 1 + /// var_n - (divisor * q_expr) = var_r + /// ``` + /// where `var_n` is a known variable (called dividend), and `q_expr` is an + /// `AffineExpr` (called the quotient expression), `var_r` can be written as: + /// + /// `var_r = var_n mod divisor`. + /// + /// Additionally, in a special case of the above constaints where `q_expr` is + /// an variable itself that is not yet known (say `var_q`), it can be written + /// as a floordiv in the following way: + /// + /// `var_q = var_n floordiv divisor`. + /// + /// Returns true if the above mod or floordiv are detected, updating 'memo' + /// with these new expressions. Returns false otherwise. + bool detectAsMod(unsigned pos, int64_t lbConst, int64_t ubConst, + SmallVectorImpl &memo, + MLIRContext *context) const; + + /// Gets the lower and upper bound of the `offset` + `pos`th variable + /// treating [0, offset) U [offset + num, symStartPos) as dimensions and + /// [symStartPos, getNumDimAndSymbolVars) as symbols, and `pos` lies in + /// [0, num). The multi-dimensional maps in the returned pair represent the + /// max and min of potentially multiple affine expressions. `localExprs` holds + /// pre-computed AffineExpr's for all local variables in the system. + /// + /// By default the returned lower bounds are closed and upper bounds are open. + /// If `closedUb` is true, the upper bound is closed. + std::pair + getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num, + unsigned symStartPos, ArrayRef localExprs, + MLIRContext *context, bool closedUB = false) const; + + /// Computes the lower and upper bounds of the first `num` dimensional + /// variables (starting at `offset`) as an affine map of the remaining + /// variables (dimensional and symbolic). This method is able to detect + /// variables as floordiv's and mod's of affine expressions of other + /// variables with respect to (positive) constants. Sets bound map to a + /// null AffineMap if such a bound can't be found (or yet unimplemented). + /// + /// By default the returned lower bounds are closed and upper bounds are open. + /// If `closedUb` is true, the upper bound is closed. + void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context, + SmallVectorImpl *lbMaps, + SmallVectorImpl *ubMaps, + bool closedUB = false); + /// Insert `num` variables of the specified kind at position `pos`. /// Positions are relative to the kind of variable. Return the absolute /// column position (i.e., not relative to the kind of variable) of the @@ -867,9 +966,54 @@ /// Return the set difference of this set and the given set, i.e., /// return `this \ set`. PresburgerSet subtract(const PresburgerSet &other) const; + +protected: + /// Given an affine map that is aligned with this constraint system: + /// * Flatten the map. + /// * Add newly introduced local columns at the beginning of this constraint + /// system (local column pos 0). + /// * Add equalities that define the new local columns to this constraint + /// system. + /// * Return the flattened expressions via `flattenedExprs`. + /// + /// Note: This is a shared helper function of `addLowerOrUpperBound` and + /// `composeMatchingMap`. + LogicalResult flattenAlignedMapAndMergeLocals( + AffineMap map, std::vector> *flattenedExprs); }; } // namespace presburger + +/// 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 variables +/// to existing dimensional and symbolic variables. See documentation for +/// AffineExprFlattener on how mod's and div's are flattened. +LogicalResult +getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols, + SmallVectorImpl *flattenedExpr, + presburger::IntegerPolyhedron *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 variables to +/// existing dimensional and / symbolic variables. 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, + presburger::IntegerPolyhedron *cst = nullptr); +LogicalResult +getFlattenedAffineExprs(IntegerSet set, + std::vector> *flattenedExprs, + presburger::IntegerPolyhedron *cst = nullptr); } // namespace mlir #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERRELATION_H diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -181,32 +181,6 @@ /// the columns in the current one regarding numbers and values. void addAffineIfOpDomain(AffineIfOp ifOp); - /// Adds a bound for the variable at the specified position with constraints - /// being drawn from the specified bound map. 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. - /// - /// The bound can be added as open or closed by specifying isClosedBound. In - /// case of a LB/UB, isClosedBound = false means the bound is added internally - /// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can - /// only be added as a closed bound. - /// - /// Note: The dimensions/symbols of this FlatAffineConstraints must match the - /// dimensions/symbols of the affine map. - LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap, - bool isClosedBound); - - /// Adds a bound for the variable at the specified position with constraints - /// being drawn from the specified bound map. 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. - /// Note: The dimensions/symbols of this FlatAffineConstraints must match the - /// dimensions/symbols of the affine map. By default the lower bound is closed - /// and the upper bound is open. - LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap); - /// Adds a bound for the variable 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 @@ -228,40 +202,12 @@ /// being known and such a local variable appearing in any of the constraints. IntegerSet getAsIntegerSet(MLIRContext *context) const; - /// Computes the lower and upper bounds of the first `num` dimensional - /// variables (starting at `offset`) as an affine map of the remaining - /// variables (dimensional and symbolic). This method is able to detect - /// variables as floordiv's and mod's of affine expressions of other - /// variables with respect to (positive) constants. Sets bound map to a - /// null AffineMap if such a bound can't be found (or yet unimplemented). - /// - /// By default the returned lower bounds are closed and upper bounds are open. - /// If `closedUb` is true, the upper bound is closed. - void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context, - SmallVectorImpl *lbMaps, - SmallVectorImpl *ubMaps, - bool closedUB = false); - /// Composes an affine map whose dimensions and symbols match one to one with /// the dimensions and symbols of this FlatAffineConstraints. The results of /// the map `other` are added as the leading dimensions of this constraint /// system. Returns failure if `other` is a semi-affine map. LogicalResult composeMatchingMap(AffineMap other); - /// Gets the lower and upper bound of the `offset` + `pos`th variable - /// treating [0, offset) U [offset + num, symStartPos) as dimensions and - /// [symStartPos, getNumDimAndSymbolVars) as symbols, and `pos` lies in - /// [0, num). The multi-dimensional maps in the returned pair represent the - /// max and min of potentially multiple affine expressions. `localExprs` holds - /// pre-computed AffineExpr's for all local variables in the system. - /// - /// By default the returned lower bounds are closed and upper bounds are open. - /// If `closedUb` is true, the upper bound is closed. - std::pair - getLowerAndUpperBound(unsigned pos, unsigned offset, unsigned num, - unsigned symStartPos, ArrayRef localExprs, - MLIRContext *context, bool closedUB = false) const; - /// Returns the bound for the variable 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 @@ -492,19 +438,6 @@ /// is meant to be used within an assert internally. bool hasConsistentState() const override; - /// Given an affine map that is aligned with this constraint system: - /// * Flatten the map. - /// * Add newly introduced local columns at the beginning of this constraint - /// system (local column pos 0). - /// * Add equalities that define the new local columns to this constraint - /// system. - /// * Return the flattened expressions via `flattenedExprs`. - /// - /// Note: This is a shared helper function of `addLowerOrUpperBound` and - /// `composeMatchingMap`. - LogicalResult flattenAlignedMapAndMergeLocals( - AffineMap map, std::vector> *flattenedExprs); - /// Eliminates the variable at the specified position using Fourier-Motzkin /// variable elimination, but uses Gaussian elimination if there is an /// equality involving that variable. If the result of the elimination is @@ -606,37 +539,6 @@ unsigned numRangeDims; }; -/// 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 variables -/// to existing dimensional and symbolic variables. See documentation for -/// AffineExprFlattener on how mod's and div's are flattened. -LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, - unsigned numSymbols, - SmallVectorImpl *flattenedExpr, - FlatAffineValueConstraints *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 variables to -/// existing dimensional and / symbolic variables. 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, - FlatAffineValueConstraints *cst = nullptr); -LogicalResult -getFlattenedAffineExprs(IntegerSet set, - std::vector> *flattenedExprs, - FlatAffineValueConstraints *cst = nullptr); - LogicalResult getMultiAffineFunctionFromMap(AffineMap map, presburger::MultiAffineFunction &multiAff); diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp --- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp @@ -18,6 +18,9 @@ #include "mlir/Analysis/Presburger/PresburgerRelation.h" #include "mlir/Analysis/Presburger/Simplex.h" #include "mlir/Analysis/Presburger/Utils.h" +#include "mlir/IR/AffineExprVisitor.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/IntegerSet.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/Support/Debug.h" @@ -383,6 +386,223 @@ inequalities.resizeVertically(0); } +bool IntegerPolyhedron::detectAsMod(unsigned pos, int64_t lbConst, + int64_t ubConst, + SmallVectorImpl &memo, + MLIRContext *context) const { + assert(pos < getNumVars() && "invalid position"); + + // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can + // be determined. + if (lbConst != 0 || ubConst < 1) + return false; + int64_t divisor = ubConst + 1; + + // Check for the aforementioned conditions in each equality. + for (unsigned curEquality = 0, numEqualities = getNumEqualities(); + curEquality < numEqualities; curEquality++) { + int64_t coefficientAtPos = atEq64(curEquality, pos); + // If current equality does not involve `var_r`, continue to the next + // equality. + if (coefficientAtPos == 0) + continue; + + // Constant term should be 0 in this equality. + if (atEq64(curEquality, getNumCols() - 1) != 0) + continue; + + // Traverse through the equality and construct the dividend expression + // `dividendExpr`, to contain all the variables which are known and are + // not divisible by `(coefficientAtPos * divisor)`. Hope here is that the + // `dividendExpr` gets simplified into a single variable `var_n` discussed + // above. + auto dividendExpr = getAffineConstantExpr(0, context); + + // Track the terms that go into quotient expression, later used to detect + // additional floordiv. + unsigned quotientCount = 0; + int quotientPosition = -1; + int quotientSign = 1; + + // Consider each term in the current equality. + unsigned curVar, e; + for (curVar = 0, e = getNumDimAndSymbolVars(); curVar < e; ++curVar) { + // Ignore var_r. + if (curVar == pos) + continue; + int64_t coefficientOfCurVar = atEq64(curEquality, curVar); + // Ignore vars that do not contribute to the current equality. + if (coefficientOfCurVar == 0) + continue; + // Check if the current var goes into the quotient expression. + if (coefficientOfCurVar % (divisor * coefficientAtPos) == 0) { + quotientCount++; + quotientPosition = curVar; + quotientSign = (coefficientOfCurVar * coefficientAtPos) > 0 ? 1 : -1; + continue; + } + // Variables that are part of dividendExpr should be known. + if (!memo[curVar]) + break; + // Append the current variable to the dividend expression. + dividendExpr = dividendExpr + memo[curVar] * coefficientOfCurVar; + } + + // Can't construct expression as it depends on a yet uncomputed var. + if (curVar < e) + continue; + + // Express `var_r` in terms of the other vars collected so far. + if (coefficientAtPos > 0) + dividendExpr = (-dividendExpr).floorDiv(coefficientAtPos); + else + dividendExpr = dividendExpr.floorDiv(-coefficientAtPos); + + // Simplify the expression. + dividendExpr = + simplifyAffineExpr(dividendExpr, getNumDimVars(), getNumSymbolVars()); + // Only if the final dividend expression is just a single var (which we call + // `var_n`), we can proceed. + // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it + // to dims themselves. + auto dimExpr = dividendExpr.dyn_cast(); + if (!dimExpr) + continue; + + // Express `var_r` as `var_n % divisor` and store the expression in `memo`. + if (quotientCount >= 1) { + auto ub = getConstantBound64(BoundType::UB, dimExpr.getPosition()); + // If `var_n` has an upperbound that is less than the divisor, mod can be + // eliminated altogether. + if (ub && *ub < divisor) + memo[pos] = dimExpr; + else + memo[pos] = dimExpr % divisor; + // If a unique quotient `var_q` was seen, it can be expressed as + // `var_n floordiv divisor`. + if (quotientCount == 1 && !memo[quotientPosition]) + memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign; + + return true; + } + } + return false; +} + +bool IntegerPolyhedron::detectAsFloorDiv( + unsigned pos, MLIRContext *context, + SmallVectorImpl &exprs) const { + assert(pos < getNumVars() && "invalid position"); + + // Get upper-lower bound pair for this variable. + SmallVector foundRepr(getNumVars(), false); + for (unsigned i = 0, e = getNumVars(); i < e; ++i) + if (exprs[i]) + foundRepr[i] = true; + + SmallVector dividend(getNumCols()); + unsigned divisor; + auto ulPair = computeSingleVarRepr(*this, foundRepr, pos, dividend, divisor); + + // No upper-lower bound pair found for this var. + if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality) + return false; + + // Construct the dividend expression. + auto dividendExpr = getAffineConstantExpr(dividend.back(), context); + for (unsigned c = 0, f = getNumVars(); c < f; c++) + if (dividend[c] != 0) + dividendExpr = dividendExpr + dividend[c] * exprs[c]; + + // Successfully detected the floordiv. + exprs[pos] = dividendExpr.floorDiv(divisor); + return true; +} + +std::pair IntegerPolyhedron::getLowerAndUpperBound( + unsigned pos, unsigned offset, unsigned num, unsigned symStartPos, + ArrayRef localExprs, MLIRContext *context, + bool closedUB) const { + assert(pos + offset < getNumDimVars() && "invalid dim start pos"); + assert(symStartPos >= (pos + offset) && "invalid sym start pos"); + assert(getNumLocalVars() == localExprs.size() && + "incorrect local exprs count"); + + SmallVector lbIndices, ubIndices, eqIndices; + getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices, + offset, num); + + /// Add to 'b' from 'a' in set [0, offset) U [offset + num, symbStartPos). + auto addCoeffs = [&](ArrayRef a, SmallVectorImpl &b) { + b.clear(); + for (unsigned i = 0, e = a.size(); i < e; ++i) { + if (i < offset || i >= offset + num) + b.push_back(a[i]); + } + }; + + SmallVector lb, ub; + SmallVector lbExprs; + unsigned dimCount = symStartPos - num; + unsigned symCount = getNumDimAndSymbolVars() - symStartPos; + lbExprs.reserve(lbIndices.size() + eqIndices.size()); + // Lower bound expressions. + for (auto idx : lbIndices) { + auto ineq = getInequality64(idx); + // Extract the lower bound (in terms of other coeff's + const), i.e., if + // i - j + 1 >= 0 is the constraint, 'pos' is for i the lower bound is j + // - 1. + addCoeffs(ineq, lb); + std::transform(lb.begin(), lb.end(), lb.begin(), std::negate()); + auto expr = + getAffineExprFromFlatForm(lb, dimCount, symCount, localExprs, context); + // expr ceildiv divisor is (expr + divisor - 1) floordiv divisor + int64_t divisor = std::abs(ineq[pos + offset]); + expr = (expr + divisor - 1).floorDiv(divisor); + lbExprs.push_back(expr); + } + + SmallVector ubExprs; + ubExprs.reserve(ubIndices.size() + eqIndices.size()); + // Upper bound expressions. + for (auto idx : ubIndices) { + auto ineq = getInequality64(idx); + // Extract the upper bound (in terms of other coeff's + const). + addCoeffs(ineq, ub); + auto expr = + getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context); + expr = expr.floorDiv(std::abs(ineq[pos + offset])); + int64_t ubAdjustment = closedUB ? 0 : 1; + ubExprs.push_back(expr + ubAdjustment); + } + + // Equalities. It's both a lower and a upper bound. + SmallVector b; + for (auto idx : eqIndices) { + auto eq = getEquality64(idx); + addCoeffs(eq, b); + if (eq[pos + offset] > 0) + std::transform(b.begin(), b.end(), b.begin(), std::negate()); + + // Extract the upper bound (in terms of other coeff's + const). + auto expr = + getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context); + expr = expr.floorDiv(std::abs(eq[pos + offset])); + // Upper bound is exclusive. + ubExprs.push_back(expr + 1); + // Lower bound. + expr = + getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context); + expr = expr.ceilDiv(std::abs(eq[pos + offset])); + lbExprs.push_back(expr); + } + + auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context); + auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context); + + return {lbMap, ubMap}; +} + /// Gather all lower and upper bounds of the variable at `pos`, and /// optionally any equalities on it. In addition, the bounds are to be /// independent of variables in position range [`offset`, `offset` + `num`). @@ -2264,6 +2484,63 @@ void IntegerRelation::dump() const { print(llvm::errs()); } +LogicalResult IntegerPolyhedron::addBound(BoundType type, unsigned pos, + AffineMap boundMap, + bool isClosedBound) { + assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch"); + assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch"); + assert(pos < getNumDimAndSymbolVars() && "invalid position"); + assert((type != BoundType::EQ || isClosedBound) && + "EQ bound must be closed."); + + // Equality follows the logic of lower bound except that we add an equality + // instead of an inequality. + assert((type != BoundType::EQ || boundMap.getNumResults() == 1) && + "single result expected"); + bool lower = type == BoundType::LB || type == BoundType::EQ; + + std::vector> flatExprs; + if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs))) + return failure(); + assert(flatExprs.size() == boundMap.getNumResults()); + + // Add one (in)equality for each result. + for (const auto &flatExpr : flatExprs) { + SmallVector ineq(getNumCols(), 0); + // Dims and symbols. + for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) { + ineq[j] = lower ? -flatExpr[j] : flatExpr[j]; + } + // Invalid bound: pos appears in `boundMap`. + // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or + // its callers to prevent invalid bounds from being added. + if (ineq[pos] != 0) + continue; + ineq[pos] = lower ? 1 : -1; + // Local columns of `ineq` are at the beginning. + unsigned j = getNumDimVars() + getNumSymbolVars(); + unsigned end = flatExpr.size() - 1; + for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) { + ineq[j] = lower ? -flatExpr[i] : flatExpr[i]; + } + // Make the bound closed in if flatExpr is open. The inequality is always + // created in the upper bound form, so the adjustment is -1. + int64_t boundAdjustment = (isClosedBound || type == BoundType::EQ) ? 0 : -1; + // Constant term. + ineq[getNumCols() - 1] = (lower ? -flatExpr[flatExpr.size() - 1] + : flatExpr[flatExpr.size() - 1]) + + boundAdjustment; + type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq); + } + + return success(); +} + +LogicalResult IntegerPolyhedron::addBound(BoundType type, unsigned pos, + AffineMap boundMap) { + return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB); +} + unsigned IntegerPolyhedron::insertVar(VarKind kind, unsigned pos, unsigned num) { assert((kind != VarKind::Domain || num == 0) && @@ -2275,6 +2552,319 @@ return IntegerPolyhedron(IntegerRelation::intersect(other)); } +/// Computes the lower and upper bounds of the first 'num' dimensional +/// variables (starting at 'offset') as affine maps of the remaining +/// variables (dimensional and symbolic variables). Local variables are +/// themselves explicitly computed as affine functions of other variables in +/// this process if needed. +void IntegerPolyhedron::getSliceBounds(unsigned offset, unsigned num, + MLIRContext *context, + SmallVectorImpl *lbMaps, + SmallVectorImpl *ubMaps, + bool closedUB) { + assert(num < getNumDimVars() && "invalid range"); + + // Basic simplification. + normalizeConstraintsByGCD(); + + LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num + << " variables\n"); + LLVM_DEBUG(dump()); + + // Record computed/detected variables. + SmallVector memo(getNumVars()); + // Initialize dimensional and symbolic variables. + for (unsigned i = 0, e = getNumDimVars(); i < e; i++) { + if (i < offset) + memo[i] = getAffineDimExpr(i, context); + else if (i >= offset + num) + memo[i] = getAffineDimExpr(i - num, context); + } + for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++) + memo[i] = getAffineSymbolExpr(i - getNumDimVars(), context); + + bool changed; + do { + changed = false; + // Identify yet unknown variables as constants or mod's / floordiv's of + // other variables if possible. + for (unsigned pos = 0; pos < getNumVars(); pos++) { + if (memo[pos]) + continue; + + auto lbConst = getConstantBound64(BoundType::LB, pos); + auto ubConst = getConstantBound64(BoundType::UB, pos); + if (lbConst.has_value() && ubConst.has_value()) { + // Detect equality to a constant. + if (*lbConst == *ubConst) { + memo[pos] = getAffineConstantExpr(*lbConst, context); + changed = true; + continue; + } + + // Detect an variable as modulo of another variable w.r.t a + // constant. + if (detectAsMod(pos, *lbConst, *ubConst, memo, context)) { + changed = true; + continue; + } + } + + // Detect an variable as a floordiv of an affine function of other + // variables (divisor is a positive constant). + if (detectAsFloorDiv(pos, context, memo)) { + changed = true; + continue; + } + + // Detect an variable as an expression of other variables. + unsigned idx; + if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) { + continue; + } + + // Build AffineExpr solving for variable 'pos' in terms of all others. + auto expr = getAffineConstantExpr(0, context); + unsigned j, e; + for (j = 0, e = getNumVars(); j < e; ++j) { + if (j == pos) + continue; + int64_t c = atEq64(idx, j); + if (c == 0) + continue; + // If any of the involved IDs hasn't been found yet, we can't proceed. + if (!memo[j]) + break; + expr = expr + memo[j] * c; + } + if (j < e) + // Can't construct expression as it depends on a yet uncomputed + // variable. + continue; + + // Add constant term to AffineExpr. + expr = expr + atEq64(idx, getNumVars()); + int64_t vPos = atEq64(idx, pos); + assert(vPos != 0 && "expected non-zero here"); + if (vPos > 0) + expr = (-expr).floorDiv(vPos); + else + // vPos < 0. + expr = expr.floorDiv(-vPos); + // Successfully constructed expression. + memo[pos] = expr; + changed = true; + } + // This loop is guaranteed to reach a fixed point - since once an + // variable's explicit form is computed (in memo[pos]), it's not updated + // again. + } while (changed); + + int64_t ubAdjustment = closedUB ? 0 : 1; + + // Set the lower and upper bound maps for all the variables that were + // computed as affine expressions of the rest as the "detected expr" and + // "detected expr + 1" respectively; set the undetected ones to null. + std::optional tmpClone; + for (unsigned pos = 0; pos < num; pos++) { + unsigned numMapDims = getNumDimVars() - num; + unsigned numMapSymbols = getNumSymbolVars(); + AffineExpr expr = memo[pos + offset]; + if (expr) + expr = simplifyAffineExpr(expr, numMapDims, numMapSymbols); + + AffineMap &lbMap = (*lbMaps)[pos]; + AffineMap &ubMap = (*ubMaps)[pos]; + + if (expr) { + lbMap = AffineMap::get(numMapDims, numMapSymbols, expr); + ubMap = AffineMap::get(numMapDims, numMapSymbols, expr + ubAdjustment); + } else { + // TODO: Whenever there are local variables in the dependence + // constraints, we'll conservatively over-approximate, since we don't + // always explicitly compute them above (in the while loop). + if (getNumLocalVars() == 0) { + // Work on a copy so that we don't update this constraint system. + if (!tmpClone) { + tmpClone.emplace(IntegerPolyhedron(*this)); + // Removing redundant inequalities is necessary so that we don't get + // redundant loop bounds. + tmpClone->removeRedundantInequalities(); + } + std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound( + pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context, + closedUB); + } + + // If the above fails, we'll just use the constant lower bound and the + // constant upper bound (if they exist) as the slice bounds. + // TODO: being conservative for the moment in cases that + // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is + // fixed (b/126426796). + if (!lbMap || lbMap.getNumResults() > 1) { + LLVM_DEBUG(llvm::dbgs() + << "WARNING: Potentially over-approximating slice lb\n"); + auto lbConst = getConstantBound64(BoundType::LB, pos + offset); + if (lbConst.has_value()) { + lbMap = AffineMap::get(numMapDims, numMapSymbols, + getAffineConstantExpr(*lbConst, context)); + } + } + if (!ubMap || ubMap.getNumResults() > 1) { + LLVM_DEBUG(llvm::dbgs() + << "WARNING: Potentially over-approximating slice ub\n"); + auto ubConst = getConstantBound64(BoundType::UB, pos + offset); + if (ubConst.has_value()) { + ubMap = AffineMap::get( + numMapDims, numMapSymbols, + getAffineConstantExpr(*ubConst + ubAdjustment, context)); + } + } + } + LLVM_DEBUG(llvm::dbgs() + << "lb map for pos = " << Twine(pos + offset) << ", expr: "); + LLVM_DEBUG(lbMap.dump();); + LLVM_DEBUG(llvm::dbgs() + << "ub map for pos = " << Twine(pos + offset) << ", expr: "); + LLVM_DEBUG(ubMap.dump();); + } +} + +LogicalResult IntegerPolyhedron::flattenAlignedMapAndMergeLocals( + AffineMap map, std::vector> *flattenedExprs) { + IntegerPolyhedron localCst(PresburgerSpace::getSetSpace()); + if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) { + LLVM_DEBUG(llvm::dbgs() + << "composition unimplemented for semi-affine maps\n"); + return failure(); + } + + // Add localCst information. + if (localCst.getNumLocalVars() > 0) { + unsigned numLocalVars = getNumLocalVars(); + // Insert local dims of localCst at the beginning. + insertVar(VarKind::Local, /*pos=*/0, /*num=*/localCst.getNumLocalVars()); + // Insert local dims of `this` at the end of localCst. + localCst.appendVar(VarKind::Local, /*num=*/numLocalVars); + // Dimensions of localCst and this constraint set match. Append localCst to + // this constraint set. + append(localCst); + } + + return success(); +} + PresburgerSet IntegerPolyhedron::subtract(const PresburgerSet &other) const { return PresburgerSet(IntegerRelation::subtract(other)); } + +namespace { + +// See comments for SimpleAffineExprFlattener. +// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording +// constraint information associated with mod's, floordiv's, and ceildiv's +// in FlatAffineValueConstraints 'localVarCst'. +struct AffineExprFlattener : public SimpleAffineExprFlattener { +public: + // Constraints connecting newly introduced local variables (for mod's and + // div's) to existing (dimensional and symbolic) ones. These are always + // inequalities. + IntegerPolyhedron localVarCst; + + AffineExprFlattener(unsigned nDims, unsigned nSymbols) + : SimpleAffineExprFlattener(nDims, nSymbols), + localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {} + +private: + // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr). + // The local variable added is always a floordiv of a pure add/mul affine + // function of other variables, coefficients of which are specified in + // `dividend' and with respect to the positive constant `divisor'. localExpr + // is the simplified tree expression (AffineExpr) corresponding to the + // quantifier. + void addLocalFloorDivId(ArrayRef dividend, int64_t divisor, + AffineExpr localExpr) override { + SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr); + // Update localVarCst. + localVarCst.addLocalFloorDiv(dividend, divisor); + } +}; + +} // namespace + +// Flattens the expressions in map. Returns failure if 'expr' was unable to be +// flattened (i.e., semi-affine expressions not handled yet). +static LogicalResult +getFlattenedAffineExprs(ArrayRef exprs, unsigned numDims, + unsigned numSymbols, + std::vector> *flattenedExprs, + IntegerPolyhedron *localVarCst) { + if (exprs.empty()) { + if (localVarCst) + *localVarCst = + IntegerPolyhedron(PresburgerSpace::getSetSpace(numDims, numSymbols)); + return success(); + } + + AffineExprFlattener flattener(numDims, numSymbols); + // Use the same flattener to simplify each expression successively. This way + // local variables / expressions are shared. + for (auto expr : exprs) { + if (!expr.isPureAffine()) + return failure(); + + flattener.walkPostOrder(expr); + } + + assert(flattener.operandExprStack.size() == exprs.size()); + flattenedExprs->clear(); + flattenedExprs->assign(flattener.operandExprStack.begin(), + flattener.operandExprStack.end()); + + if (localVarCst) + localVarCst->clearAndCopyFrom(flattener.localVarCst); + + return success(); +} + +// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to +// be flattened (semi-affine expressions not handled yet). +LogicalResult mlir::getFlattenedAffineExpr( + AffineExpr expr, unsigned numDims, unsigned numSymbols, + SmallVectorImpl *flattenedExpr, IntegerPolyhedron *localVarCst) { + std::vector> flattenedExprs; + LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols, + &flattenedExprs, localVarCst); + *flattenedExpr = flattenedExprs[0]; + return ret; +} + +/// Flattens the expressions in map. Returns failure if 'expr' was unable to be +/// flattened (i.e., semi-affine expressions not handled yet). +LogicalResult mlir::getFlattenedAffineExprs( + AffineMap map, std::vector> *flattenedExprs, + IntegerPolyhedron *localVarCst) { + if (map.getNumResults() == 0) { + if (localVarCst) + *localVarCst = IntegerPolyhedron( + PresburgerSpace::getSetSpace(map.getNumDims(), map.getNumSymbols())); + return success(); + } + return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(), + map.getNumSymbols(), flattenedExprs, + localVarCst); +} + +LogicalResult mlir::getFlattenedAffineExprs( + IntegerSet set, std::vector> *flattenedExprs, + IntegerPolyhedron *localVarCst) { + if (set.getNumConstraints() == 0) { + if (localVarCst) + *localVarCst = IntegerPolyhedron( + PresburgerSpace::getSetSpace(set.getNumDims(), set.getNumSymbols())); + return success(); + } + return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(), + set.getNumSymbols(), flattenedExprs, + localVarCst); +} diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -33,120 +33,8 @@ using namespace mlir; using namespace presburger; -namespace { - -// See comments for SimpleAffineExprFlattener. -// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording -// constraint information associated with mod's, floordiv's, and ceildiv's -// in FlatAffineValueConstraints 'localVarCst'. -struct AffineExprFlattener : public SimpleAffineExprFlattener { -public: - // Constraints connecting newly introduced local variables (for mod's and - // div's) to existing (dimensional and symbolic) ones. These are always - // inequalities. - IntegerPolyhedron localVarCst; - - AffineExprFlattener(unsigned nDims, unsigned nSymbols) - : SimpleAffineExprFlattener(nDims, nSymbols), - localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {} - -private: - // Add a local variable (needed to flatten a mod, floordiv, ceildiv expr). - // The local variable added is always a floordiv of a pure add/mul affine - // function of other variables, coefficients of which are specified in - // `dividend' and with respect to the positive constant `divisor'. localExpr - // is the simplified tree expression (AffineExpr) corresponding to the - // quantifier. - void addLocalFloorDivId(ArrayRef dividend, int64_t divisor, - AffineExpr localExpr) override { - SimpleAffineExprFlattener::addLocalFloorDivId(dividend, divisor, localExpr); - // Update localVarCst. - localVarCst.addLocalFloorDiv(dividend, divisor); - } -}; - -} // namespace - -// Flattens the expressions in map. Returns failure if 'expr' was unable to be -// flattened (i.e., semi-affine expressions not handled yet). -static LogicalResult -getFlattenedAffineExprs(ArrayRef exprs, unsigned numDims, - unsigned numSymbols, - std::vector> *flattenedExprs, - FlatAffineValueConstraints *localVarCst) { - if (exprs.empty()) { - if (localVarCst) - *localVarCst = FlatAffineValueConstraints(numDims, numSymbols); - return success(); - } - - AffineExprFlattener flattener(numDims, numSymbols); - // Use the same flattener to simplify each expression successively. This way - // local variables / expressions are shared. - for (auto expr : exprs) { - if (!expr.isPureAffine()) - return failure(); - - flattener.walkPostOrder(expr); - } - - assert(flattener.operandExprStack.size() == exprs.size()); - flattenedExprs->clear(); - flattenedExprs->assign(flattener.operandExprStack.begin(), - flattener.operandExprStack.end()); - - if (localVarCst) - localVarCst->clearAndCopyFrom(flattener.localVarCst); - - return success(); -} - -// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to -// be flattened (semi-affine expressions not handled yet). -LogicalResult -mlir::getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, - unsigned numSymbols, - SmallVectorImpl *flattenedExpr, - FlatAffineValueConstraints *localVarCst) { - std::vector> flattenedExprs; - LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols, - &flattenedExprs, localVarCst); - *flattenedExpr = flattenedExprs[0]; - return ret; -} - -/// Flattens the expressions in map. Returns failure if 'expr' was unable to be -/// flattened (i.e., semi-affine expressions not handled yet). -LogicalResult mlir::getFlattenedAffineExprs( - AffineMap map, std::vector> *flattenedExprs, - FlatAffineValueConstraints *localVarCst) { - if (map.getNumResults() == 0) { - if (localVarCst) - *localVarCst = - FlatAffineValueConstraints(map.getNumDims(), map.getNumSymbols()); - return success(); - } - return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(), - map.getNumSymbols(), flattenedExprs, - localVarCst); -} - -LogicalResult mlir::getFlattenedAffineExprs( - IntegerSet set, std::vector> *flattenedExprs, - FlatAffineValueConstraints *localVarCst) { - if (set.getNumConstraints() == 0) { - if (localVarCst) - *localVarCst = - FlatAffineValueConstraints(set.getNumDims(), set.getNumSymbols()); - return success(); - } - return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(), - set.getNumSymbols(), flattenedExprs, - localVarCst); -} - //===----------------------------------------------------------------------===// -// FlatAffineConstraints / FlatAffineValueConstraints. +// FlatAffineValueConstraints. //===----------------------------------------------------------------------===// std::unique_ptr @@ -717,510 +605,6 @@ } } -// Determine whether the variable at 'pos' (say var_r) can be expressed as -// modulo of another known variable (say var_n) w.r.t a constant. For example, -// if the following constraints hold true: -// ``` -// 0 <= var_r <= divisor - 1 -// var_n - (divisor * q_expr) = var_r -// ``` -// where `var_n` is a known variable (called dividend), and `q_expr` is an -// `AffineExpr` (called the quotient expression), `var_r` can be written as: -// -// `var_r = var_n mod divisor`. -// -// Additionally, in a special case of the above constaints where `q_expr` is an -// variable itself that is not yet known (say `var_q`), it can be written as a -// floordiv in the following way: -// -// `var_q = var_n floordiv divisor`. -// -// Returns true if the above mod or floordiv are detected, updating 'memo' with -// these new expressions. Returns false otherwise. -static bool detectAsMod(const FlatAffineValueConstraints &cst, unsigned pos, - int64_t lbConst, int64_t ubConst, - SmallVectorImpl &memo, - MLIRContext *context) { - assert(pos < cst.getNumVars() && "invalid position"); - - // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can - // be determined. - if (lbConst != 0 || ubConst < 1) - return false; - int64_t divisor = ubConst + 1; - - // Check for the aforementioned conditions in each equality. - for (unsigned curEquality = 0, numEqualities = cst.getNumEqualities(); - curEquality < numEqualities; curEquality++) { - int64_t coefficientAtPos = cst.atEq64(curEquality, pos); - // If current equality does not involve `var_r`, continue to the next - // equality. - if (coefficientAtPos == 0) - continue; - - // Constant term should be 0 in this equality. - if (cst.atEq64(curEquality, cst.getNumCols() - 1) != 0) - continue; - - // Traverse through the equality and construct the dividend expression - // `dividendExpr`, to contain all the variables which are known and are - // not divisible by `(coefficientAtPos * divisor)`. Hope here is that the - // `dividendExpr` gets simplified into a single variable `var_n` discussed - // above. - auto dividendExpr = getAffineConstantExpr(0, context); - - // Track the terms that go into quotient expression, later used to detect - // additional floordiv. - unsigned quotientCount = 0; - int quotientPosition = -1; - int quotientSign = 1; - - // Consider each term in the current equality. - unsigned curVar, e; - for (curVar = 0, e = cst.getNumDimAndSymbolVars(); curVar < e; ++curVar) { - // Ignore var_r. - if (curVar == pos) - continue; - int64_t coefficientOfCurVar = cst.atEq64(curEquality, curVar); - // Ignore vars that do not contribute to the current equality. - if (coefficientOfCurVar == 0) - continue; - // Check if the current var goes into the quotient expression. - if (coefficientOfCurVar % (divisor * coefficientAtPos) == 0) { - quotientCount++; - quotientPosition = curVar; - quotientSign = (coefficientOfCurVar * coefficientAtPos) > 0 ? 1 : -1; - continue; - } - // Variables that are part of dividendExpr should be known. - if (!memo[curVar]) - break; - // Append the current variable to the dividend expression. - dividendExpr = dividendExpr + memo[curVar] * coefficientOfCurVar; - } - - // Can't construct expression as it depends on a yet uncomputed var. - if (curVar < e) - continue; - - // Express `var_r` in terms of the other vars collected so far. - if (coefficientAtPos > 0) - dividendExpr = (-dividendExpr).floorDiv(coefficientAtPos); - else - dividendExpr = dividendExpr.floorDiv(-coefficientAtPos); - - // Simplify the expression. - dividendExpr = simplifyAffineExpr(dividendExpr, cst.getNumDimVars(), - cst.getNumSymbolVars()); - // Only if the final dividend expression is just a single var (which we call - // `var_n`), we can proceed. - // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it - // to dims themselves. - auto dimExpr = dividendExpr.dyn_cast(); - if (!dimExpr) - continue; - - // Express `var_r` as `var_n % divisor` and store the expression in `memo`. - if (quotientCount >= 1) { - auto ub = cst.getConstantBound64( - FlatAffineValueConstraints::BoundType::UB, dimExpr.getPosition()); - // If `var_n` has an upperbound that is less than the divisor, mod can be - // eliminated altogether. - if (ub && *ub < divisor) - memo[pos] = dimExpr; - else - memo[pos] = dimExpr % divisor; - // If a unique quotient `var_q` was seen, it can be expressed as - // `var_n floordiv divisor`. - if (quotientCount == 1 && !memo[quotientPosition]) - memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign; - - return true; - } - } - return false; -} - -/// Check if the pos^th variable can be expressed as a floordiv of an affine -/// function of other variables (where the divisor is a positive constant) -/// given the initial set of expressions in `exprs`. If it can be, the -/// corresponding position in `exprs` is set as the detected affine expr. For -/// eg: 4q <= i + j <= 4q + 3 <=> q = (i + j) floordiv 4. An equality can -/// also yield a floordiv: eg. 4q = i + j <=> q = (i + j) floordiv 4. 32q + 28 -/// <= i <= 32q + 31 => q = i floordiv 32. -static bool detectAsFloorDiv(const FlatAffineValueConstraints &cst, - unsigned pos, MLIRContext *context, - SmallVectorImpl &exprs) { - assert(pos < cst.getNumVars() && "invalid position"); - - // Get upper-lower bound pair for this variable. - SmallVector foundRepr(cst.getNumVars(), false); - for (unsigned i = 0, e = cst.getNumVars(); i < e; ++i) - if (exprs[i]) - foundRepr[i] = true; - - SmallVector dividend(cst.getNumCols()); - unsigned divisor; - auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor); - - // No upper-lower bound pair found for this var. - if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality) - return false; - - // Construct the dividend expression. - auto dividendExpr = getAffineConstantExpr(dividend.back(), context); - for (unsigned c = 0, f = cst.getNumVars(); c < f; c++) - if (dividend[c] != 0) - dividendExpr = dividendExpr + dividend[c] * exprs[c]; - - // Successfully detected the floordiv. - exprs[pos] = dividendExpr.floorDiv(divisor); - return true; -} - -std::pair -FlatAffineValueConstraints::getLowerAndUpperBound( - unsigned pos, unsigned offset, unsigned num, unsigned symStartPos, - ArrayRef localExprs, MLIRContext *context, - bool closedUB) const { - assert(pos + offset < getNumDimVars() && "invalid dim start pos"); - assert(symStartPos >= (pos + offset) && "invalid sym start pos"); - assert(getNumLocalVars() == localExprs.size() && - "incorrect local exprs count"); - - SmallVector lbIndices, ubIndices, eqIndices; - getLowerAndUpperBoundIndices(pos + offset, &lbIndices, &ubIndices, &eqIndices, - offset, num); - - /// Add to 'b' from 'a' in set [0, offset) U [offset + num, symbStartPos). - auto addCoeffs = [&](ArrayRef a, SmallVectorImpl &b) { - b.clear(); - for (unsigned i = 0, e = a.size(); i < e; ++i) { - if (i < offset || i >= offset + num) - b.push_back(a[i]); - } - }; - - SmallVector lb, ub; - SmallVector lbExprs; - unsigned dimCount = symStartPos - num; - unsigned symCount = getNumDimAndSymbolVars() - symStartPos; - lbExprs.reserve(lbIndices.size() + eqIndices.size()); - // Lower bound expressions. - for (auto idx : lbIndices) { - auto ineq = getInequality64(idx); - // Extract the lower bound (in terms of other coeff's + const), i.e., if - // i - j + 1 >= 0 is the constraint, 'pos' is for i the lower bound is j - // - 1. - addCoeffs(ineq, lb); - std::transform(lb.begin(), lb.end(), lb.begin(), std::negate()); - auto expr = - getAffineExprFromFlatForm(lb, dimCount, symCount, localExprs, context); - // expr ceildiv divisor is (expr + divisor - 1) floordiv divisor - int64_t divisor = std::abs(ineq[pos + offset]); - expr = (expr + divisor - 1).floorDiv(divisor); - lbExprs.push_back(expr); - } - - SmallVector ubExprs; - ubExprs.reserve(ubIndices.size() + eqIndices.size()); - // Upper bound expressions. - for (auto idx : ubIndices) { - auto ineq = getInequality64(idx); - // Extract the upper bound (in terms of other coeff's + const). - addCoeffs(ineq, ub); - auto expr = - getAffineExprFromFlatForm(ub, dimCount, symCount, localExprs, context); - expr = expr.floorDiv(std::abs(ineq[pos + offset])); - int64_t ubAdjustment = closedUB ? 0 : 1; - ubExprs.push_back(expr + ubAdjustment); - } - - // Equalities. It's both a lower and a upper bound. - SmallVector b; - for (auto idx : eqIndices) { - auto eq = getEquality64(idx); - addCoeffs(eq, b); - if (eq[pos + offset] > 0) - std::transform(b.begin(), b.end(), b.begin(), std::negate()); - - // Extract the upper bound (in terms of other coeff's + const). - auto expr = - getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context); - expr = expr.floorDiv(std::abs(eq[pos + offset])); - // Upper bound is exclusive. - ubExprs.push_back(expr + 1); - // Lower bound. - expr = - getAffineExprFromFlatForm(b, dimCount, symCount, localExprs, context); - expr = expr.ceilDiv(std::abs(eq[pos + offset])); - lbExprs.push_back(expr); - } - - auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context); - auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context); - - return {lbMap, ubMap}; -} - -/// Computes the lower and upper bounds of the first 'num' dimensional -/// variables (starting at 'offset') as affine maps of the remaining -/// variables (dimensional and symbolic variables). Local variables are -/// themselves explicitly computed as affine functions of other variables in -/// this process if needed. -void FlatAffineValueConstraints::getSliceBounds( - unsigned offset, unsigned num, MLIRContext *context, - SmallVectorImpl *lbMaps, SmallVectorImpl *ubMaps, - bool closedUB) { - assert(num < getNumDimVars() && "invalid range"); - - // Basic simplification. - normalizeConstraintsByGCD(); - - LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num - << " variables\n"); - LLVM_DEBUG(dump()); - - // Record computed/detected variables. - SmallVector memo(getNumVars()); - // Initialize dimensional and symbolic variables. - for (unsigned i = 0, e = getNumDimVars(); i < e; i++) { - if (i < offset) - memo[i] = getAffineDimExpr(i, context); - else if (i >= offset + num) - memo[i] = getAffineDimExpr(i - num, context); - } - for (unsigned i = getNumDimVars(), e = getNumDimAndSymbolVars(); i < e; i++) - memo[i] = getAffineSymbolExpr(i - getNumDimVars(), context); - - bool changed; - do { - changed = false; - // Identify yet unknown variables as constants or mod's / floordiv's of - // other variables if possible. - for (unsigned pos = 0; pos < getNumVars(); pos++) { - if (memo[pos]) - continue; - - auto lbConst = getConstantBound64(BoundType::LB, pos); - auto ubConst = getConstantBound64(BoundType::UB, pos); - if (lbConst.has_value() && ubConst.has_value()) { - // Detect equality to a constant. - if (*lbConst == *ubConst) { - memo[pos] = getAffineConstantExpr(*lbConst, context); - changed = true; - continue; - } - - // Detect an variable as modulo of another variable w.r.t a - // constant. - if (detectAsMod(*this, pos, *lbConst, *ubConst, memo, context)) { - changed = true; - continue; - } - } - - // Detect an variable as a floordiv of an affine function of other - // variables (divisor is a positive constant). - if (detectAsFloorDiv(*this, pos, context, memo)) { - changed = true; - continue; - } - - // Detect an variable as an expression of other variables. - unsigned idx; - if (!findConstraintWithNonZeroAt(pos, /*isEq=*/true, &idx)) { - continue; - } - - // Build AffineExpr solving for variable 'pos' in terms of all others. - auto expr = getAffineConstantExpr(0, context); - unsigned j, e; - for (j = 0, e = getNumVars(); j < e; ++j) { - if (j == pos) - continue; - int64_t c = atEq64(idx, j); - if (c == 0) - continue; - // If any of the involved IDs hasn't been found yet, we can't proceed. - if (!memo[j]) - break; - expr = expr + memo[j] * c; - } - if (j < e) - // Can't construct expression as it depends on a yet uncomputed - // variable. - continue; - - // Add constant term to AffineExpr. - expr = expr + atEq64(idx, getNumVars()); - int64_t vPos = atEq64(idx, pos); - assert(vPos != 0 && "expected non-zero here"); - if (vPos > 0) - expr = (-expr).floorDiv(vPos); - else - // vPos < 0. - expr = expr.floorDiv(-vPos); - // Successfully constructed expression. - memo[pos] = expr; - changed = true; - } - // This loop is guaranteed to reach a fixed point - since once an - // variable's explicit form is computed (in memo[pos]), it's not updated - // again. - } while (changed); - - int64_t ubAdjustment = closedUB ? 0 : 1; - - // Set the lower and upper bound maps for all the variables that were - // computed as affine expressions of the rest as the "detected expr" and - // "detected expr + 1" respectively; set the undetected ones to null. - std::optional tmpClone; - for (unsigned pos = 0; pos < num; pos++) { - unsigned numMapDims = getNumDimVars() - num; - unsigned numMapSymbols = getNumSymbolVars(); - AffineExpr expr = memo[pos + offset]; - if (expr) - expr = simplifyAffineExpr(expr, numMapDims, numMapSymbols); - - AffineMap &lbMap = (*lbMaps)[pos]; - AffineMap &ubMap = (*ubMaps)[pos]; - - if (expr) { - lbMap = AffineMap::get(numMapDims, numMapSymbols, expr); - ubMap = AffineMap::get(numMapDims, numMapSymbols, expr + ubAdjustment); - } else { - // TODO: Whenever there are local variables in the dependence - // constraints, we'll conservatively over-approximate, since we don't - // always explicitly compute them above (in the while loop). - if (getNumLocalVars() == 0) { - // Work on a copy so that we don't update this constraint system. - if (!tmpClone) { - tmpClone.emplace(FlatAffineValueConstraints(*this)); - // Removing redundant inequalities is necessary so that we don't get - // redundant loop bounds. - tmpClone->removeRedundantInequalities(); - } - std::tie(lbMap, ubMap) = tmpClone->getLowerAndUpperBound( - pos, offset, num, getNumDimVars(), /*localExprs=*/{}, context, - closedUB); - } - - // If the above fails, we'll just use the constant lower bound and the - // constant upper bound (if they exist) as the slice bounds. - // TODO: being conservative for the moment in cases that - // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is - // fixed (b/126426796). - if (!lbMap || lbMap.getNumResults() > 1) { - LLVM_DEBUG(llvm::dbgs() - << "WARNING: Potentially over-approximating slice lb\n"); - auto lbConst = getConstantBound64(BoundType::LB, pos + offset); - if (lbConst.has_value()) { - lbMap = AffineMap::get(numMapDims, numMapSymbols, - getAffineConstantExpr(*lbConst, context)); - } - } - if (!ubMap || ubMap.getNumResults() > 1) { - LLVM_DEBUG(llvm::dbgs() - << "WARNING: Potentially over-approximating slice ub\n"); - auto ubConst = getConstantBound64(BoundType::UB, pos + offset); - if (ubConst.has_value()) { - ubMap = AffineMap::get( - numMapDims, numMapSymbols, - getAffineConstantExpr(*ubConst + ubAdjustment, context)); - } - } - } - LLVM_DEBUG(llvm::dbgs() - << "lb map for pos = " << Twine(pos + offset) << ", expr: "); - LLVM_DEBUG(lbMap.dump();); - LLVM_DEBUG(llvm::dbgs() - << "ub map for pos = " << Twine(pos + offset) << ", expr: "); - LLVM_DEBUG(ubMap.dump();); - } -} - -LogicalResult FlatAffineValueConstraints::flattenAlignedMapAndMergeLocals( - AffineMap map, std::vector> *flattenedExprs) { - FlatAffineValueConstraints localCst; - if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) { - LLVM_DEBUG(llvm::dbgs() - << "composition unimplemented for semi-affine maps\n"); - return failure(); - } - - // Add localCst information. - if (localCst.getNumLocalVars() > 0) { - unsigned numLocalVars = getNumLocalVars(); - // Insert local dims of localCst at the beginning. - insertLocalVar(/*pos=*/0, /*num=*/localCst.getNumLocalVars()); - // Insert local dims of `this` at the end of localCst. - localCst.appendLocalVar(/*num=*/numLocalVars); - // Dimensions of localCst and this constraint set match. Append localCst to - // this constraint set. - append(localCst); - } - - return success(); -} - -LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos, - AffineMap boundMap, - bool isClosedBound) { - assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch"); - assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch"); - assert(pos < getNumDimAndSymbolVars() && "invalid position"); - assert((type != BoundType::EQ || isClosedBound) && - "EQ bound must be closed."); - - // Equality follows the logic of lower bound except that we add an equality - // instead of an inequality. - assert((type != BoundType::EQ || boundMap.getNumResults() == 1) && - "single result expected"); - bool lower = type == BoundType::LB || type == BoundType::EQ; - - std::vector> flatExprs; - if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs))) - return failure(); - assert(flatExprs.size() == boundMap.getNumResults()); - - // Add one (in)equality for each result. - for (const auto &flatExpr : flatExprs) { - SmallVector ineq(getNumCols(), 0); - // Dims and symbols. - for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) { - ineq[j] = lower ? -flatExpr[j] : flatExpr[j]; - } - // Invalid bound: pos appears in `boundMap`. - // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or - // its callers to prevent invalid bounds from being added. - if (ineq[pos] != 0) - continue; - ineq[pos] = lower ? 1 : -1; - // Local columns of `ineq` are at the beginning. - unsigned j = getNumDimVars() + getNumSymbolVars(); - unsigned end = flatExpr.size() - 1; - for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) { - ineq[j] = lower ? -flatExpr[i] : flatExpr[i]; - } - // Make the bound closed in if flatExpr is open. The inequality is always - // created in the upper bound form, so the adjustment is -1. - int64_t boundAdjustment = (isClosedBound || type == BoundType::EQ) ? 0 : -1; - // Constant term. - ineq[getNumCols() - 1] = (lower ? -flatExpr[flatExpr.size() - 1] - : flatExpr[flatExpr.size() - 1]) + - boundAdjustment; - type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq); - } - - return success(); -} - -LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos, - AffineMap boundMap) { - return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB); -} - AffineMap FlatAffineValueConstraints::computeAlignedMap(AffineMap map, ValueRange operands) const { @@ -1247,7 +631,7 @@ AffineMap alignedMap = alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr); - // All symbols are already part of this FlatAffineConstraints. + // All symbols are already part of this FlatAffineValueConstraints. assert(syms.size() == newSymsPtr->size() && "unexpected new/missing symbols"); assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()) && "unexpected new/missing symbols"); @@ -1456,7 +840,7 @@ changed = false; for (unsigned i = 0, e = cst.getNumLocalVars(); i < e; ++i) if (!memo[numDims + numSyms + i] && - detectAsFloorDiv(cst, /*pos=*/numDims + numSyms + i, context, memo)) + cst.detectAsFloorDiv(/*pos=*/numDims + numSyms + i, context, memo)) changed = true; } while (changed);