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 @@ -160,6 +160,13 @@ /// otherwise. bool containsPoint(ArrayRef point) const; + /// Find pairs of indices of inequalities, using which an explicit + /// representation for each local variable can be computed. + /// Pairs are returned as indices of upperbound, lowerbound equalities. + /// If no such pair is found, pair is marked as None. + std::vector>> + computeLocalRepr() const; + // Clones this object. std::unique_ptr clone() const; 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 @@ -1250,6 +1250,112 @@ 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. +static Optional> +computeSingleVarRepr(const FlatAffineConstraints &cst, + const std::vector &foundRepr, unsigned pos) { + assert(pos < cst.getNumIds() && "invalid position"); + assert(foundRepr.size() == cst.getNumIds()); + + SmallVector lbIndices, ubIndices; + cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices); + + // `id` is equivalent to `expr floordiv divisor` if there + // are constraints of the form: + // expr - (divisor - 1) <= divisor * id <= expr + // Rearranging, we have: + // divisor * id - expr + (divisor - 1) >= 0 <-- Lower bound for 'id' + // -divisor * id + expr >= 0 <-- Upper bound for 'id' + // + // For example, if -32*k + 16*i + j >= 0 + // 32*k - 16*i - j + 31 >= 0 <=> + // k = ( 16*i + j ) floordiv 32 + for (unsigned ubPos : ubIndices) { + for (unsigned lbPos : lbIndices) { + // Due to the form of the inequalities, sum of constants of the + // inequalities is (divisor - 1) + int64_t divisor = cst.atIneq(lbPos, cst.getNumCols() - 1) + + cst.atIneq(ubPos, cst.getNumCols() - 1) + 1; + + // Divisor should be positive. + if (divisor <= 0) + continue; + + // Check if coeff of variable is equal to divisor. + if (divisor != cst.atIneq(lbPos, pos)) + continue; + + // Check if constraints are opposite of each other. Constant term + // is not required to be opposite and is not checked. + unsigned c = 0, f = 0; + for (c = 0, f = cst.getNumIds(); c < f; ++c) + if (cst.atIneq(ubPos, c) != -cst.atIneq(lbPos, c)) + break; + + if (c < f) + 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. + for (c = 0, f = cst.getNumIds(); c < f; ++c) { + if (c == pos) + continue; + if (!foundRepr[c] && cst.atIneq(lbPos, 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) + continue; + + return std::make_pair(ubPos, lbPos); + } + } + + return llvm::None; +} + +/// Find pairs of indices of inequalities, using which an explicit +/// representation for each local variable can be computed. +/// Pairs are returned as indices of upperbound, lowerbound equalities. +/// If no such pair is found, pair is marked as None. +std::vector>> +FlatAffineConstraints::computeLocalRepr() const { + + std::vector>> repr( + getNumLocalIds(), llvm::None); + std::vector foundRepr(getNumIds(), false); + + for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i) + foundRepr[i] = true; + + unsigned divOffset = getNumDimAndSymbolIds(); + bool changed; + do { + // Each time changed is true, at end of this iteration, one or more local + // vars have been detected as floor divs. + changed = false; + for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) { + if (!foundRepr[i + divOffset]) { + if (auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i)) { + foundRepr[i + divOffset] = true; + repr[i] = res; + changed = true; + } + } + } + } while (changed); + + return repr; +} + /// Tightens inequalities given that we are dealing with integer spaces. This is /// analogous to the GCD test but applied to inequalities. The constant term can /// be reduced to the preceding multiple of the GCD of the coefficients, i.e., @@ -1516,70 +1622,46 @@ SmallVectorImpl &exprs) { assert(pos < cst.getNumIds() && "invalid position"); - SmallVector lbIndices, ubIndices; - cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices); + // Get upper-lower bound pair for this variable + std::vector foundRepr(cst.getNumIds(), false); + for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i) + if (exprs[i]) + foundRepr[i] = true; + + auto ulPair = computeSingleVarRepr(cst, foundRepr, pos); + + // No upper-lower bound pair found for this var + if (!ulPair) + return false; + + unsigned ubPos = ulPair->first; - // Check if any lower bound, upper bound pair is of the form: - // divisor * id >= expr - (divisor - 1) <-- Lower bound for 'id' - // divisor * id <= expr <-- Upper bound for 'id' - // Then, 'id' is equivalent to 'expr floordiv divisor'. (where divisor > 1). + // Upper bound is of the form: + // -divisor * id + expr >= 0 + // where `id` is equivalent to `expr floordiv divisor` // - // For example, if -32*k + 16*i + j >= 0 - // 32*k - 16*i - j + 31 >= 0 <=> - // k = ( 16*i + j ) floordiv 32 - unsigned seenDividends = 0; - for (auto ubPos : ubIndices) { - for (auto lbPos : lbIndices) { - // Check if the lower bound's constant term is divisor - 1. The - // 'divisor' here is cst.atIneq(lbPos, pos) and we already know that it's - // positive (since cst.Ineq(lbPos, ...) is a lower bound expr for 'pos'. - int64_t divisor = cst.atIneq(lbPos, pos); - int64_t lbConstTerm = cst.atIneq(lbPos, cst.getNumCols() - 1); - if (lbConstTerm != divisor - 1) - continue; - // Check if upper bound's constant term is 0. - if (cst.atIneq(ubPos, cst.getNumCols() - 1) != 0) - continue; - // For the remaining part, check if the lower bound expr's coeff's are - // negations of corresponding upper bound ones'. - unsigned c, f; - for (c = 0, f = cst.getNumCols() - 1; c < f; c++) { - if (cst.atIneq(lbPos, c) != -cst.atIneq(ubPos, c)) - break; - if (c != pos && cst.atIneq(lbPos, c) != 0) - seenDividends++; - } - // Lb coeff's aren't negative of ub coeff's (for the non constant term - // part). - if (c < f) - continue; - if (seenDividends >= 1) { - // Construct the dividend expression. - auto dividendExpr = getAffineConstantExpr(0, context); - unsigned c, f; - for (c = 0, f = cst.getNumCols() - 1; c < f; c++) { - if (c == pos) - continue; - int64_t ubVal = cst.atIneq(ubPos, c); - if (ubVal == 0) - continue; - if (!exprs[c]) - break; - dividendExpr = dividendExpr + ubVal * exprs[c]; - } - // 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) - continue; - // Successfully detected the floordiv. - exprs[pos] = dividendExpr.floorDiv(divisor); - return true; - } - } + // Since the division cannot be dependent on itself, the coefficient of + // of `id` in `expr` is zero. The coefficient of `id` in the upperbound + // is -divisor. + int64_t divisor = -cst.atIneq(ubPos, pos); + int64_t constantTerm = cst.atIneq(ubPos, cst.getNumCols() - 1); + + // Construct the dividend expression. + auto dividendExpr = getAffineConstantExpr(constantTerm, context); + unsigned c, f; + for (c = 0, f = cst.getNumCols() - 1; c < f; c++) { + if (c == pos) + continue; + int64_t ubVal = cst.atIneq(ubPos, c); + if (ubVal == 0) + continue; + // computeSingleVarRepr gurantees that expr is known here + dividendExpr = dividendExpr + ubVal * exprs[c]; } - return false; + + // Successfully detected the floordiv. + exprs[pos] = dividendExpr.floorDiv(divisor); + return true; } // Fills an inequality row with the value 'val'. diff --git a/mlir/unittests/Analysis/AffineStructuresTest.cpp b/mlir/unittests/Analysis/AffineStructuresTest.cpp --- a/mlir/unittests/Analysis/AffineStructuresTest.cpp +++ b/mlir/unittests/Analysis/AffineStructuresTest.cpp @@ -587,4 +587,122 @@ EXPECT_EQ(fac.atIneq(0, 1), 0); } +/// Check if the expected division representation of local variables matches the +/// computed representation. The expected division representation is given as +/// a vector of expressions set in `divisions` and the corressponding +/// denominator in `denoms`. If expected denominator for a variable is +/// non-positive, the local variable is expected to not have a computed +/// representation. +void checkDivisionRepresentation( + FlatAffineConstraints &fac, + const std::vector> &divisions, + const std::vector &denoms) { + + assert(divisions.size() == fac.getNumLocalIds() && + "Size of expected divisions does not match number of local variables"); + assert( + denoms.size() == fac.getNumLocalIds() && + "Size of expected denominators does not match number of local variables"); + + auto res = fac.computeLocalRepr(); + + // Check if all expected divisions are computed. + for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i) + if (denoms[i] > 0) + EXPECT_NE(res[i], llvm::None); + else + EXPECT_EQ(res[i], llvm::None); + + unsigned divOffset = fac.getNumDimAndSymbolIds(); + for (unsigned i = 0, e = fac.getNumLocalIds(); i < e; ++i) { + if (!res[i]) + continue; + + // Check if bounds are of the form: + // divisor * id - expr + (divisor - 1) >= 0 <-- Lower bound for 'id' + // -divisor * id + expr >= 0 <-- Upper bound for 'id' + // where `id = expr floordiv divisor` + unsigned ubPos = res[i]->first, lbPos = res[i]->second; + const SmallVector &expr = divisions[i]; + + // Check if lower bound is of the correct form. + int64_t computedDivisorLb = fac.atIneq(lbPos, i + divOffset); + EXPECT_EQ(computedDivisorLb, denoms[i]); + for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) { + if (c == i + divOffset) + continue; + EXPECT_EQ(fac.atIneq(lbPos, c), -expr[c]); + } + // Check if constant term of lower bound matches expected constant term + EXPECT_EQ(fac.atIneq(lbPos, fac.getNumCols() - 1), + -expr.back() + (denoms[i] - 1)); + + // Check if upper bound is of the correct form. + int64_t computedDivisorUb = fac.atIneq(ubPos, i + divOffset); + EXPECT_EQ(computedDivisorUb, -denoms[i]); + for (unsigned c = 0, f = fac.getNumLocalIds(); c < f; ++c) { + if (c == i + divOffset) + continue; + EXPECT_EQ(fac.atIneq(ubPos, c), expr[c]); + } + // Check if constant term of upper bound matches expected constant term + EXPECT_EQ(fac.atIneq(ubPos, fac.getNumCols() - 1), expr.back()); + } +} + +TEST(FlatAffineConstraintsTest, computeLocalReprSimple) { + FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {}); + + fac.addLocalFloorDiv({1, 4}, 10); + fac.addLocalFloorDiv({1, 0, 100}, 10); + + std::vector> divisions = {{1, 0, 0, 4}, + {1, 0, 0, 100}}; + std::vector denoms = {10, 10}; + + // Check if floordivs can be computed when no other inequalities exist + // and floor divs do not depend on each other. + checkDivisionRepresentation(fac, divisions, denoms); +} + +TEST(FlatAffineConstraintsTest, computeLocalReprConstantFloorDiv) { + FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {}); + + fac.addInequality({1, 0, 3, 1, 2}); + fac.addInequality({1, 2, -8, 1, 10}); + fac.addEquality({1, 2, -4, 1, 10}); + + fac.addLocalFloorDiv({0, 0, 0, 0, 10}, 30); + fac.addLocalFloorDiv({0, 0, 0, 0, 0, 99}, 101); + + std::vector> divisions = {{0, 0, 0, 0, 0, 0, 10}, + {0, 0, 0, 0, 0, 0, 99}}; + std::vector denoms = {30, 101}; + + // Check if floordivs with constant numerator can be computed. + checkDivisionRepresentation(fac, divisions, denoms); +} + +TEST(FlatAffineConstraintsTest, computeLocalReprRecursive) { + FlatAffineConstraints fac = makeFACFromConstraints(4, {}, {}); + fac.addInequality({1, 0, 3, 1, 2}); + fac.addInequality({1, 2, -8, 1, 10}); + fac.addEquality({1, 2, -4, 1, 10}); + + fac.addLocalFloorDiv({0, -2, 7, 2, 10}, 3); + fac.addLocalFloorDiv({3, 0, 9, 2, 2, 10}, 5); + fac.addLocalFloorDiv({0, 1, -123, 2, 0, -4, 10}, 3); + + fac.addInequality({1, 2, -2, 1, -5, 0, 6, 100}); + fac.addInequality({1, 2, -8, 1, 3, 7, 0, -9}); + + std::vector> divisions = {{0, -2, 7, 2, 0, 0, 0, 10}, + {3, 0, 9, 2, 2, 0, 0, 10}, + {0, 1, -123, 2, 0, -4, 10}}; + std::vector denoms = {3, 5, 3}; + + // Check if floordivs which may depend on other floordivs can be computed. + checkDivisionRepresentation(fac, divisions, denoms); +} + } // namespace mlir