[MLIR] Extract division representation from equality expressions.

Extract the division representation from equality constraints.
For example:
    32*k == 16*i + j - 31                 <-- k is the localVariable
    expr = 16*i + j - 31, divisor = 32
    k = (16*i + j - 32) floordiv 32
The dividend of the division is set to [16, 1, -32] and the divisor is set
to 32.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D117959
This commit is contained in:
Prashant Kumar 2022-01-23 00:22:46 +05:30
parent 5281f0dab2
commit 1e7c464d2c
5 changed files with 155 additions and 38 deletions

View File

@ -265,10 +265,11 @@ public:
/// otherwise.
bool containsPoint(ArrayRef<int64_t> 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

View File

@ -43,10 +43,11 @@ struct MaybeLocalRepr {
/// 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<bool> foundRepr, unsigned pos,
SmallVector<int64_t, 8> &dividend,

View File

@ -856,13 +856,11 @@ void IntegerPolyhedron::getLocalReprs(
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);

View File

@ -140,14 +140,79 @@ static LogicalResult getDivRepr(const IntegerPolyhedron &cst, unsigned pos,
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<int64_t, 8> &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<bool> foundRepr,
SmallVectorImpl<int64_t> &dividend,
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<bool> foundRepr, unsigned pos,
SmallVector<int64_t, 8> &dividend, unsigned &divisor) {
@ -155,9 +220,9 @@ MaybeLocalRepr presburger_utils::computeSingleVarRepr(
assert(foundRepr.size() == cst.getNumIds() &&
"Size of foundRepr does not match total number of variables");
SmallVector<unsigned, 4> lbIndices, ubIndices;
cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
MaybeLocalRepr repr;
SmallVector<unsigned, 4> lbIndices, ubIndices, eqIndices;
cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices, &eqIndices);
MaybeLocalRepr repr{};
for (unsigned ubPos : ubIndices) {
for (unsigned lbPos : lbIndices) {
@ -165,22 +230,7 @@ MaybeLocalRepr presburger_utils::computeSingleVarRepr(
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,6 +238,18 @@ MaybeLocalRepr presburger_utils::computeSingleVarRepr(
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;
}

View File

@ -736,6 +736,61 @@ TEST(IntegerPolyhedronTest, computeLocalReprTightUpperBound) {
}
}
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<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
SmallVector<unsigned, 8> 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<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
SmallVector<unsigned, 8> 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<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, -2}};
SmallVector<unsigned, 8> 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<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, 0, 1},
{-1, -1, 0, 0, 0}};
SmallVector<unsigned, 8> denoms = {4, 3};
checkDivisionRepresentation(poly, divisions, denoms);
}
}
TEST(IntegerPolyhedronTest, computeLocalReprNoRepr) {
MLIRContext context;
IntegerPolyhedron poly =