diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h b/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h --- a/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h @@ -216,12 +216,12 @@ /// Gather positions of all lower and upper bounds of the identifier at `pos`, /// and optionally any equalities on it. In addition, the bounds are to be /// independent of identifiers in position range [`offset`, `offset` + `num`). - void - getLowerAndUpperBoundIndices(unsigned pos, - SmallVectorImpl *lbIndices, - SmallVectorImpl *ubIndices, - SmallVectorImpl *eqIndices = nullptr, - unsigned offset = 0, unsigned num = 0) const; + void getLowerAndUpperBoundIndices(unsigned pos, + SmallVectorImpl *lbIndices, + SmallVectorImpl *ubIndices, + SmallVectorImpl *eqIndices, + unsigned offset = 0, + unsigned num = 0) const; /// Checks for emptiness by performing variable elimination on all /// identifiers, running the GCD test on each equality constraint, and diff --git a/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp b/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp --- a/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp @@ -842,6 +842,11 @@ repr[i].repr.inEqualityPair = {res.repr.inEqualityPair.lowerBoundIdx, res.repr.inEqualityPair.upperBoundIdx}; changed = true; + } else if (res.kind == ReprKind::Equality) { + foundRepr[i + divOffset] = true; + repr[i].kind = ReprKind::Equality; + repr[i].repr.equalityIdx = res.repr.equalityIdx; + changed = true; } } } diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp --- a/mlir/lib/Analysis/Presburger/Utils.cpp +++ b/mlir/lib/Analysis/Presburger/Utils.cpp @@ -140,6 +140,68 @@ return success(); } +/// Check if the pos^th identifier can be represented as a division using +/// equality at position `eqInd`. +/// +/// For example: +/// 32*k == 16*i + j - 31 <-- `eqInd` for 'k' +/// expr = 16*i + j - 31, divisor = 32 +/// k = (16*i + j - 31) floordiv 32 +/// +/// If successful, `expr` is set to dividend of the division and `divisor` is +/// set to the denominator of the division. The final division expression is +/// normalized by GCD. +static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos, + unsigned eqInd, SmallVector &expr, + unsigned &divisor) { + + assert(pos <= cst.getNumIds() && "Invalid identifier position"); + assert(eqInd <= cst.getNumEqualities() && "Invalid equality position"); + + // Extract divisor, the divisor can be negative and hence its sign information + // is stored in `signDiv` to reverse the sign of dividend's coefficients. + int64_t temp_div = cst.atEq(eqInd, pos); + int64_t signDiv = temp_div < 0 ? -1 : 1; + + // Divisor is always positive integer. + divisor = temp_div * signDiv; + + expr.resize(cst.getNumCols(), 0); + for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i) + if (i != pos) + expr[i] = signDiv * cst.atEq(eqInd, i); + + expr.back() = signDiv * cst.atEq(eqInd, cst.getNumCols() - 1); + normalizeDivisionByGCD(expr, divisor); + + return success(); +} + +// `checkExplicitRepresentation` returns false if the constraints depends on a +// variable for which an explicit representation has not been found yet, +// otherwise returns true. +static bool checkExplicitRepresentation(const IntegerPolyhedron &cst, + ArrayRef foundRepr, + SmallVectorImpl ÷nd, + const unsigned &pos) { + // Exit to avoid circular dependencies between divisions. + unsigned c, f; + for (c = 0, f = cst.getNumIds(); c < f; ++c) { + if (c == pos) + continue; + if (!foundRepr[c] && dividend[c] != 0) + break; + } + + // Expression can't be constructed as it depends on a yet unknown + // identifier. + // TODO: Visit/compute the identifiers in an order so that this doesn't + // happen. More complex but much more efficient. + if (c < f) + return false; + return true; +} + /// Check if the pos^th identifier can be expressed as a floordiv of an affine /// function of other identifiers (where the divisor is a positive constant). /// `foundRepr` contains a boolean for each identifier indicating if the @@ -155,8 +217,8 @@ assert(foundRepr.size() == cst.getNumIds() && "Size of foundRepr does not match total number of variables"); - SmallVector lbIndices, ubIndices; - cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices); + SmallVector lbIndices, ubIndices, eqIndices; + cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices, &eqIndices); MaybeLocalRepr repr; for (unsigned ubPos : ubIndices) { @@ -165,22 +227,7 @@ if (failed(getDivRepr(cst, pos, ubPos, lbPos, dividend, divisor))) continue; - // Check if the inequalities depend on a variable for which - // an explicit representation has not been found yet. - // Exit to avoid circular dependencies between divisions. - unsigned c, f; - for (c = 0, f = cst.getNumIds(); c < f; ++c) { - if (c == pos) - continue; - if (!foundRepr[c] && dividend[c] != 0) - break; - } - - // Expression can't be constructed as it depends on a yet unknown - // identifier. - // TODO: Visit/compute the identifiers in an order so that this doesn't - // happen. More complex but much more efficient. - if (c < f) + if (!checkExplicitRepresentation(cst, foundRepr, dividend, pos)) continue; repr.kind = ReprKind::Inequality; @@ -188,5 +235,17 @@ return repr; } } + for (unsigned eqPos : eqIndices) { + // Attempt to get divison representation from eqPos. + if (failed(getDivRepr(cst, pos, eqPos, dividend, divisor))) + continue; + + if (!checkExplicitRepresentation(cst, foundRepr, dividend, pos)) + continue; + + repr.kind = ReprKind::Equality; + repr.repr.equalityIdx = eqPos; + return repr; + } return repr; } diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -3289,7 +3289,8 @@ } SmallVector lbIndices, ubIndices; - cst.getLowerAndUpperBoundIndices(/*pos=*/0, &lbIndices, &ubIndices); + cst.getLowerAndUpperBoundIndices(/*pos=*/0, &lbIndices, &ubIndices, + /*eqIndices=*/nullptr); auto fLb = cst.getInequality(fullTileLbPos); auto fUb = cst.getInequality(fullTileUbPos); diff --git a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp --- a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp +++ b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp @@ -734,6 +734,43 @@ } } +TEST(IntegerPolyhedronTest, computeLocalReprFromEquality) { + MLIRContext context; + { + IntegerPolyhedron poly = + parsePoly("(i, j, q) : (-4*q + i + j == 0)", &context); + // Convert `q` to a local variable. + poly.convertDimToLocal(2, 3); + + std::vector> divisions = {{-1, -1, 0, 0}}; + SmallVector denoms = {4}; + + checkDivisionRepresentation(poly, divisions, denoms); + } + { + IntegerPolyhedron poly = + parsePoly("(i, j, q) : (4*q - i - j == 0)", &context); + // Convert `q` to a local variable. + poly.convertDimToLocal(2, 3); + + std::vector> divisions = {{-1, -1, 0, 0}}; + SmallVector denoms = {4}; + + checkDivisionRepresentation(poly, divisions, denoms); + } + { + IntegerPolyhedron poly = + parsePoly("(i, j, q) : (3*q + i + j - 2 == 0)", &context); + // Convert `q` to a local variable. + poly.convertDimToLocal(2, 3); + + std::vector> divisions = {{1, 1, 0, -2}}; + SmallVector denoms = {3}; + + checkDivisionRepresentation(poly, divisions, denoms); + } +} + TEST(IntegerPolyhedronTest, computeLocalReprNoRepr) { MLIRContext context; IntegerPolyhedron poly =