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 @@ -259,10 +259,11 @@ /// otherwise. bool containsPoint(ArrayRef point) const; - /// Find pairs of inequalities identified by their position indices, using - /// which an explicit representation for each local variable can be computed. - /// The pairs are stored as indices of upperbound, lowerbound inequalities. If - /// no such pair can be found, it is stored as llvm::None. + /// Find equality and pairs of inequality contraints identified by their + /// position indices, using which an explicit representation for each local + /// variable can be computed. The indices of the constraints are stored in + /// `MaybeLocalRepr` struct. If no such pair can be found, the kind attribute + /// in `MaybeLocalRepr` is set to None. /// /// The dividends of the explicit representations are stored in `dividends` /// and the denominators in `denominators`. If no explicit representation diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h --- a/mlir/include/mlir/Analysis/Presburger/Utils.h +++ b/mlir/include/mlir/Analysis/Presburger/Utils.h @@ -42,10 +42,11 @@ /// function of other identifiers (where the divisor is a positive constant). /// `foundRepr` contains a boolean for each identifier indicating if the /// explicit representation for that identifier has already been computed. -/// Returns the upper and lower bound inequalities using which the floordiv -/// can be computed. If the representation could be computed, `dividend` and -/// `denominator` are set. If the representation could not be computed, -/// `llvm::None` is returned. +/// Returns the `MaybeLocalRepr` struct which contains the indices of the +/// constraints that can be expressed as a floordiv of an affine function. If +/// the representation could be computed, `dividend` and `denominator` are set. +/// If the representation could not be computed, the kind attribute in +/// `MaybeLocalRepr` is set to None. MaybeLocalRepr computeSingleVarRepr(const IntegerPolyhedron &cst, ArrayRef foundRepr, unsigned pos, SmallVector ÷nd, 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 @@ -836,13 +836,11 @@ if (!foundRepr[i + divOffset]) { auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i, dividends[i], denominators[i]); - if (res.kind == ReprKind::Inequality) { - foundRepr[i + divOffset] = true; - repr[i].kind = ReprKind::Inequality; - repr[i].repr.inEqualityPair = {res.repr.inEqualityPair.lowerBoundIdx, - res.repr.inEqualityPair.upperBoundIdx}; - changed = true; - } + if (res.kind == ReprKind::None) + continue; + foundRepr[i + divOffset] = true; + repr[i] = res; + changed = true; } } } while (changed); 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,14 +140,79 @@ 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. + // Equality must involve the pos-th variable and hence `temp_div` != 0. + int64_t temp_div = cst.atEq(eqInd, pos); + if (temp_div == 0) + return failure(); + int64_t signDiv = temp_div < 0 ? -1 : 1; + + // The divisor is always a 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(); +} + +// 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, + 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 /// explicit representation for that identifier has already been computed. -/// Returns the upper and lower bound inequalities using which the floordiv can -/// be computed. If the representation could be computed, `dividend` and -/// `denominator` are set. If the representation could not be computed, -/// `llvm::None` is returned. +/// Returns the `MaybeLocalRepr` struct which contains the indices of the +/// constraints that can be expressed as a floordiv of an affine function. If +/// the representation could be computed, `dividend` and `denominator` are set. +/// If the representation could not be computed, the kind attribute in +/// `MaybeLocalRepr` is set to None. MaybeLocalRepr presburger_utils::computeSingleVarRepr( const IntegerPolyhedron &cst, ArrayRef foundRepr, unsigned pos, SmallVector ÷nd, unsigned &divisor) { @@ -155,9 +220,9 @@ assert(foundRepr.size() == cst.getNumIds() && "Size of foundRepr does not match total number of variables"); - SmallVector lbIndices, ubIndices; - cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices); - MaybeLocalRepr repr; + SmallVector lbIndices, ubIndices, eqIndices; + cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices, &eqIndices); + MaybeLocalRepr repr{}; for (unsigned ubPos : ubIndices) { for (unsigned lbPos : lbIndices) { @@ -165,22 +230,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 +238,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/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,61 @@ } } +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, computeLocalReprFromEqualityAndInequality) { + MLIRContext context; + { + IntegerPolyhedron poly = + parsePoly("(i, j, q, k) : (-3*k + i + j == 0, 4*q - " + "i - j + 2 >= 0, -4*q + i + j >= 0)", + &context); + // Convert `q` and `k` to local variables. + poly.convertDimToLocal(2, 4); + + std::vector> divisions = {{1, 1, 0, 0, 1}, + {-1, -1, 0, 0, 0}}; + SmallVector denoms = {4, 3}; + + checkDivisionRepresentation(poly, divisions, denoms); + } +} + TEST(IntegerPolyhedronTest, computeLocalReprNoRepr) { MLIRContext context; IntegerPolyhedron poly =