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 @@ -147,11 +147,8 @@ /// returns true, no integer solution to the equality constraints can exist. bool isEmptyByGCDTest() const; - /// Runs the GCD test heuristic. If it proves inconclusive, falls back to - /// generalized basis reduction if the set is bounded. - /// /// Returns true if the set of constraints is found to have no solution, - /// false if a solution exists or all tests were inconclusive. + /// false if a solution exists. Uses the same algorithm as findIntegerSample. bool isIntegerEmpty() const; // Returns a matrix where each row is a vector along which the polytope is @@ -160,11 +157,12 @@ // independent. This function should not be called on empty sets. Matrix getBoundedDirections() const; - /// Find a sample point satisfying the constraints. This uses a branch and - /// bound algorithm with generalized basis reduction, which always works if - /// the set is bounded. This should not be called for unbounded sets. + /// Find an integer sample point satisfying the constraints using a + /// branch and bound algorithm with generalized basis reduction, with some + /// additional processing using Simplex for unbounded sets. /// - /// Returns such a point if one exists, or an empty Optional otherwise. + /// Returns an integer sample point if one exists, or an empty Optional + /// otherwise. Optional> findIntegerSample() const; /// Returns true if the given point satisfies the constraints, or false @@ -387,8 +385,9 @@ /// Changes all symbol identifiers which are loop IVs to dim identifiers. void convertLoopIVSymbolsToDims(); - /// Sets the specified identifier to a constant and removes it. - void setAndEliminate(unsigned pos, int64_t constVal); + /// Sets the values.size() identifiers starting at pos to the specified values + /// and removes them. + void setAndEliminate(unsigned pos, ArrayRef values); /// Tries to fold the specified identifier to a constant using a trivial /// equality detection; if successful, the constant is substituted for the diff --git a/mlir/include/mlir/Analysis/LinearTransform.h b/mlir/include/mlir/Analysis/LinearTransform.h --- a/mlir/include/mlir/Analysis/LinearTransform.h +++ b/mlir/include/mlir/Analysis/LinearTransform.h @@ -35,10 +35,15 @@ // Returns a FlatAffineConstraints having a constraint vector vT for every // constraint vector v in fac, where T is this transform. - FlatAffineConstraints applyTo(const FlatAffineConstraints &fac); + FlatAffineConstraints applyTo(const FlatAffineConstraints &fac) const; - // Post-multiply the given vector v with this transform, say T, returning vT. - SmallVector applyTo(ArrayRef v); + // The given vector is interpreted as a row vector v. Post-multiply v with + // this transform, say T, and return vT. + SmallVector postMultiplyRow(ArrayRef rowVec) const; + + // The given vector is interpreted as a column vector v. Pre-multiply v with + // this transform, say T, and return Tv. + SmallVector preMultiplyColumn(ArrayRef colVec) const; private: Matrix matrix; diff --git a/mlir/include/mlir/Analysis/Presburger/Simplex.h b/mlir/include/mlir/Analysis/Presburger/Simplex.h --- a/mlir/include/mlir/Analysis/Presburger/Simplex.h +++ b/mlir/include/mlir/Analysis/Presburger/Simplex.h @@ -218,7 +218,11 @@ /// tableau A and one in B. static Simplex makeProduct(const Simplex &a, const Simplex &b); - /// Returns the current sample point if it is integral. Otherwise, returns an + /// Returns a rational sample point. This should not be called when Simplex is + /// empty. + SmallVector getRationalSample() const; + + /// Returns the current sample point if it is integral. Otherwise, returns /// None. Optional> getSamplePointIfIntegral() 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 @@ -20,6 +20,7 @@ #include "mlir/IR/IntegerSet.h" #include "mlir/Support/LLVM.h" #include "mlir/Support/MathExtras.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Debug.h" @@ -1111,9 +1112,13 @@ fac.removeInequality(i - 1); } +bool FlatAffineConstraints::isIntegerEmpty() const { + return !findIntegerSample().hasValue(); +} + /// Let this set be S. If S is bounded then we directly call into the GBR /// sampling algorithm. Otherwise, there are some unbounded directions, i.e., -/// vectors v such that S extends to infininty along v or -v. In this case we +/// vectors v such that S extends to infinity along v or -v. In this case we /// use an algorithm described in the integer set library (isl) manual and used /// by the isl_set_sample function in that library. The algorithm is: /// @@ -1121,44 +1126,54 @@ /// dimensions in which S*T is bounded lie in the linear span of a prefix of the /// dimensions. /// -/// 2) Construct a set transformedSet by removing all constraints that involve -/// the unbounded dimensions and also deleting the unbounded dimensions. Note -/// that this is a bounded set. +/// 2) Construct a set B by removing all constraints that involve +/// the unbounded dimensions and then deleting the unbounded dimensions. Note +/// that B is a Bounded set. /// -/// 3) Check if transformedSet is empty using the GBR sampling algorithm. +/// 3) Try to obtain a sample from B using the GBR sampling +/// algorithm. If no sample is found, return that S is empty. /// -/// 4) return S is empty iff transformedSet is empty. +/// 4) Otherwise, substitute the obtained sample into S*T to obtain a set +/// C. C is a full-dimensional Cone and always contains a sample. /// -/// Since T is unimodular, a vector v is a solution to S*T iff T*v is a -/// solution to S. The following is a sketch of a proof that S*T is empty -/// iff transformedSet is empty: +/// 5) Obtain an integer sample from C. /// -/// If transformedSet is empty, then S*T is certainly empty since transformedSet -/// was obtained by removing constraints and deleting dimensions from S*T. +/// 6) Return T*v, where v is the concatenation of the samples from B and C. /// -/// If transformedSet contains a sample, consider the set C obtained by -/// substituting the sample for the bounded dimensions of S*T. All the -/// constraints of S*T that did not involve unbounded dimensions are -/// satisfied by this substitution. +/// The following is a sketch of a proof that +/// a) If the algorithm returns empty, then S is empty. +/// b) If the algorithm returns a sample, it is a valid sample in S. /// -/// In step 1, all dimensions in the linear span of the dimensions outside the -/// prefix are unbounded in S*T. Substituting values for the bounded dimensions -/// cannot makes these dimensions bounded, and these are the only remaining -/// dimensions in C, so C is unbounded along every vector. C is hence a -/// full-dimensional cone and therefore always contains an integer point, which -/// we can then substitute to get a full solution to S*T. -bool FlatAffineConstraints::isIntegerEmpty() const { +/// The algorithm returns empty only if B is empty, in which case S*T is +/// certainly empty since B was obtained by removing constraints and then +/// deleting unconstrained dimensions from S*T. Since T is unimodular, a vector +/// v is in S*T iff T*v is in S. So in this case, since +/// S*T is empty, S is empty too. +/// +/// Otherwise, the algorithm substitutes the sample from B into S*T. All the +/// constraints of S*T that did not involve unbounded dimensions are satisfied +/// by this substitution. All dimensions in the linear span of the dimensions +/// outside the prefix are unbounded in S*T (step 1). Substituting values for +/// the bounded dimensions cannot make these dimensions bounded, and these are +/// the only remaining dimensions in C, so C is unbounded along every vector (in +/// the positive or negative direction, or both). C is hence a full-dimensional +/// cone and therefore always contains an integer point. +/// +/// Concatenating the samples from B and C gives a sample v in S*T, so the +/// returned sample T*v is a sample in S. +Optional> +FlatAffineConstraints::findIntegerSample() const { // First, try the GCD test heuristic. if (isEmptyByGCDTest()) - return true; + return {}; Simplex simplex(*this); if (simplex.isEmpty()) - return true; + return {}; // For a bounded set, we directly call into the GBR sampling algorithm. if (!simplex.isUnbounded()) - return !simplex.findIntegerSample().hasValue(); + return simplex.findIntegerSample(); // The set is unbounded. We cannot directly use the GBR algorithm. // @@ -1172,21 +1187,79 @@ // transform to use in step 1 of the algorithm. std::pair result = LinearTransform::makeTransformToColumnEchelon(std::move(m)); - FlatAffineConstraints transformedSet = result.second.applyTo(*this); + const LinearTransform &transform = result.second; + // 1) Apply T to S to obtain S*T. + FlatAffineConstraints transformedSet = transform.applyTo(*this); + // 2) Remove the unbounded dimensions and constraints involving them to + // obtain a bounded set. + FlatAffineConstraints boundedSet = transformedSet; unsigned numBoundedDims = result.first; unsigned numUnboundedDims = getNumIds() - numBoundedDims; - removeConstraintsInvolvingSuffixDims(transformedSet, numUnboundedDims); + removeConstraintsInvolvingSuffixDims(boundedSet, numUnboundedDims); + boundedSet.removeIdRange(numBoundedDims, boundedSet.getNumIds()); + + // 3) Try to obtain a sample from the bounded set. + Optional> boundedSample = + Simplex(boundedSet).findIntegerSample(); + if (!boundedSample) + return {}; + assert(boundedSet.containsPoint(*boundedSample) && + "Simplex returned an invalid sample!"); + + // 4) Substitute the values of the bounded dimensions into S*T to obtain a + // full-dimensional cone, which necessarily contains an integer sample. + transformedSet.setAndEliminate(0, *boundedSample); + FlatAffineConstraints &cone = transformedSet; + + // 5) Obtain an integer sample from the cone. + // + // We shrink the cone such that for any rational point in the shrunken cone, + // rounding up each of the point's coordinates produces a point that still + // lies in the original cone. + // + // Rounding up a point x adds a number e_i in [0, 1) to each coordinate x_i. + // For each inequality sum_i a_i x_i + c >= 0 in the original cone, the + // shrunken cone will have the inequality tightened by some amount s, such + // that if x satisfies the shrunken cone's tightened inequality, then x + e + // satisfies the original inequality, i.e., + // + // sum_i a_i x_i + c + s >= 0 implies sum_i a_i (x_i + e_i) + c >= 0 + // + // for any e_i values in [0, 1). In fact, we will handle the slightly more + // general case where e_i can be in [0, 1]. For example, consider the + // inequality 2x_1 - 3x_2 - 7x_3 - 6 >= 0, and let x = (3, 0, 0). How low + // could the LHS go if we added a number in [0, 1] to each coordinate? The LHS + // is minimized when we add 1 to the x_i with negative coefficient a_i and + // keep the other x_i the same. In the example, we would get x = (3, 1, 1), + // changing the value of the LHS by -3 + -7 = -10. + // + // In general, the value of the LHS can change by at most the sum of the + // negative a_i, so we accomodate this by shifting the inequality by this + // amount for the shrunken cone. + for (unsigned i = 0, e = cone.getNumInequalities(); i < e; ++i) { + for (unsigned j = 0; j < cone.numIds; ++j) { + int64_t coeff = cone.atIneq(i, j); + if (coeff < 0) + cone.atIneq(i, cone.numIds) += coeff; + } + } - // Remove all the unbounded dimensions. - transformedSet.removeIdRange(numBoundedDims, transformedSet.getNumIds()); + // Obtain an integer sample in the cone by rounding up a rational point from + // the shrunken cone. Shrinking the cone amounts to shifting its apex + // "inwards" without changing its "shape"; the shrunken cone is still a + // full-dimensional cone and is hence non-empty. + Simplex shrunkenConeSimplex(cone); + assert(!shrunkenConeSimplex.isEmpty() && "Shrunken cone cannot be empty!"); + SmallVector shrunkenConeSample = + shrunkenConeSimplex.getRationalSample(); - return !Simplex(transformedSet).findIntegerSample().hasValue(); -} + SmallVector coneSample(llvm::map_range(shrunkenConeSample, ceil)); -Optional> -FlatAffineConstraints::findIntegerSample() const { - return Simplex(*this).findIntegerSample(); + // 6) Return transform * concat(boundedSample, coneSample). + SmallVector &sample = boundedSample.getValue(); + sample.append(coneSample.begin(), coneSample.end()); + return transform.preMultiplyColumn(sample); } /// Helper to evaluate an affine expression at a point. @@ -2204,15 +2277,22 @@ return -1; } -void FlatAffineConstraints::setAndEliminate(unsigned pos, int64_t constVal) { - assert(pos < getNumIds() && "invalid position"); - for (unsigned r = 0, e = getNumInequalities(); r < e; r++) { - atIneq(r, getNumCols() - 1) += atIneq(r, pos) * constVal; - } - for (unsigned r = 0, e = getNumEqualities(); r < e; r++) { - atEq(r, getNumCols() - 1) += atEq(r, pos) * constVal; - } - removeId(pos); +void FlatAffineConstraints::setAndEliminate(unsigned pos, + ArrayRef values) { + if (values.empty()) + return; + assert(pos + values.size() <= getNumIds() && + "invalid position or too many values"); + // Setting x_j = p in sum_i a_i x_i + c is equivalent to adding p*a_j to the + // constant term and removing the id x_j. We do this for all the ids + // pos, pos + 1, ... pos + values.size() - 1. + for (unsigned r = 0, e = getNumInequalities(); r < e; r++) + for (unsigned i = 0, numVals = values.size(); i < numVals; ++i) + atIneq(r, getNumCols() - 1) += atIneq(r, pos + i) * values[i]; + for (unsigned r = 0, e = getNumEqualities(); r < e; r++) + for (unsigned i = 0, numVals = values.size(); i < numVals; ++i) + atEq(r, getNumCols() - 1) += atEq(r, pos + i) * values[i]; + removeIdRange(pos, pos + values.size()); } LogicalResult FlatAffineConstraints::constantFoldId(unsigned pos) { diff --git a/mlir/lib/Analysis/LinearTransform.cpp b/mlir/lib/Analysis/LinearTransform.cpp --- a/mlir/lib/Analysis/LinearTransform.cpp +++ b/mlir/lib/Analysis/LinearTransform.cpp @@ -111,23 +111,32 @@ return {echelonCol, LinearTransform(std::move(resultMatrix))}; } -SmallVector LinearTransform::applyTo(ArrayRef v) { - assert(v.size() == matrix.getNumRows() && - "vector dimension should be matrix output dimension"); - - SmallVector result; - result.reserve(v.size()); - for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col) { - int64_t elem = 0; +SmallVector +LinearTransform::postMultiplyRow(ArrayRef rowVec) const { + assert(rowVec.size() == matrix.getNumRows() && + "row vector dimension should match transform output dimension"); + + SmallVector result(matrix.getNumColumns(), 0); + for (unsigned col = 0, e = matrix.getNumColumns(); col < e; ++col) for (unsigned i = 0, e = matrix.getNumRows(); i < e; ++i) - elem += v[i] * matrix(i, col); - result.push_back(elem); - } + result[col] += rowVec[i] * matrix(i, col); + return result; +} + +SmallVector +LinearTransform::preMultiplyColumn(ArrayRef colVec) const { + assert(matrix.getNumColumns() == colVec.size() && + "column vector dimension should match transform input dimension"); + + SmallVector result(matrix.getNumRows(), 0); + for (unsigned row = 0, e = matrix.getNumRows(); row < e; row++) + for (unsigned i = 0, e = matrix.getNumColumns(); i < e; i++) + result[row] += matrix(row, i) * colVec[i]; return result; } FlatAffineConstraints -LinearTransform::applyTo(const FlatAffineConstraints &fac) { +LinearTransform::applyTo(const FlatAffineConstraints &fac) const { FlatAffineConstraints result(fac.getNumDimIds()); for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i) { @@ -135,7 +144,7 @@ int64_t c = eq.back(); - SmallVector newEq = applyTo(eq.drop_back()); + SmallVector newEq = postMultiplyRow(eq.drop_back()); newEq.push_back(c); result.addEquality(newEq); } @@ -145,7 +154,7 @@ int64_t c = ineq.back(); - SmallVector newIneq = applyTo(ineq.drop_back()); + SmallVector newIneq = postMultiplyRow(ineq.drop_back()); newIneq.push_back(c); result.addInequality(newIneq); } diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp --- a/mlir/lib/Analysis/Presburger/Simplex.cpp +++ b/mlir/lib/Analysis/Presburger/Simplex.cpp @@ -682,30 +682,42 @@ return result; } -Optional> Simplex::getSamplePointIfIntegral() const { - // The tableau is empty, so no sample point exists. - if (empty) - return {}; +SmallVector Simplex::getRationalSample() const { + assert(!empty && "This should not be called when Simplex is empty."); - SmallVector sample; + SmallVector sample; + sample.reserve(var.size()); // Push the sample value for each variable into the vector. for (const Unknown &u : var) { if (u.orientation == Orientation::Column) { // If the variable is in column position, its sample value is zero. - sample.push_back(0); + sample.emplace_back(0, 1); } else { // If the variable is in row position, its sample value is the entry in // the constant column divided by the entry in the common denominator - // column. If this is not an integer, then the sample point is not - // integral so we return None. - if (tableau(u.pos, 1) % tableau(u.pos, 0) != 0) - return {}; - sample.push_back(tableau(u.pos, 1) / tableau(u.pos, 0)); + // column. + sample.emplace_back(tableau(u.pos, 1), tableau(u.pos, 0)); } } return sample; } +Optional> Simplex::getSamplePointIfIntegral() const { + // If the tableau is empty, no sample point exists. + if (empty) + return {}; + SmallVector rationalSample = getRationalSample(); + SmallVector integerSample; + integerSample.reserve(var.size()); + for (const Fraction &coord : rationalSample) { + // If the sample is non-integral, return None. + if (coord.num % coord.den != 0) + return {}; + integerSample.push_back(coord.num / coord.den); + } + return integerSample; +} + /// Given a simplex for a polytope, construct a new simplex whose variables are /// identified with a pair of points (x, y) in the original polytope. Supports /// some operations needed for generalized basis reduction. In what follows, 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 @@ -212,45 +212,12 @@ {0, -2, 99}, // 2y <= 99. }, {})); -} - -TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) { - // 1 <= 5x and 5x <= 4 (no solution). - EXPECT_TRUE( - makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty()); - // 1 <= 5x and 5x <= 9 (solution: x = 1). - EXPECT_FALSE( - makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty()); - - // Unbounded sets. - EXPECT_TRUE(makeFACFromConstraints(3, - { - {2, 0, 0, -1}, // 2x >= 1 - {-2, 0, 0, 1}, // 2x <= 1 - {0, 2, 0, -1}, // 2y >= 1 - {0, -2, 0, 1}, // 2y <= 1 - {0, 0, 2, -1}, // 2z >= 1 - }, - {}) - .isIntegerEmpty()); - - EXPECT_FALSE(makeFACFromConstraints(3, - { - {2, 0, 0, -1}, // 2x >= 1 - {-3, 0, 0, 3}, // 3x <= 3 - {0, 0, 5, -6}, // 5z >= 6 - {0, 0, -7, 17}, // 7z <= 17 - {0, 3, 0, -2}, // 3y >= 2 - }, - {}) - .isIntegerEmpty()); - // 2D cone with apex at (10000, 10000) and // edges passing through (1/3, 0) and (2/3, 0). - EXPECT_FALSE( + checkSample( + true, makeFACFromConstraints( - 2, {{300000, -299999, -100000}, {-300000, 299998, 200000}}, {}) - .isIntegerEmpty()); + 2, {{300000, -299999, -100000}, {-300000, 299998, 200000}}, {})); // Cartesian product of a tetrahedron and a 2D cone. // The tetrahedron has vertices at @@ -276,7 +243,7 @@ // -300000p + 299998q + 200000 >= 0 {0, 0, 0, -300000, 299998, 200000}, }, - {}, TestFunction::Empty); + {}); // Cartesian product of same tetrahedron as above and {(p, q) : 1/3 <= p <= // 2/3}. Since the second set is empty, the whole set is too. @@ -297,7 +264,7 @@ // 3p <= 2 {0, 0, 0, -3, 0, 2}, }, - {}, TestFunction::Empty); + {}); // Cartesian product of same tetrahedron as above and // {(p, q, r) : 1 <= p <= 2 and p = 3q + 3r}. @@ -321,8 +288,7 @@ }, { {0, 0, 0, 1, -3, -3, 0}, // p = 3q + 3r - }, - TestFunction::Empty); + }); // Cartesian product of a tetrahedron and a 2D cone. // The tetrahedron is empty and has vertices at @@ -344,7 +310,7 @@ // -300000p + 299998q + 200000 >= 0 {0, 0, 0, -300000, 299998, 200000}, }, - {}, TestFunction::Empty); + {}); // Cartesian product of same tetrahedron as above and // {(p, q) : 1/3 <= p <= 2/3}. @@ -362,7 +328,47 @@ // 3p <= 2 {0, 0, 0, -3, 0, 2}, }, - {}, TestFunction::Empty); + {}); + + checkSample(true, makeFACFromConstraints(3, + { + {2, 0, 0, -1}, // 2x >= 1 + }, + {{ + {1, -1, 0, -1}, // y = x - 1 + {0, 1, -1, 0}, // z = y + }})); +} + +TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) { + // 1 <= 5x and 5x <= 4 (no solution). + EXPECT_TRUE( + makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty()); + // 1 <= 5x and 5x <= 9 (solution: x = 1). + EXPECT_FALSE( + makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty()); + + // Unbounded sets. + EXPECT_TRUE(makeFACFromConstraints(3, + { + {0, 2, 0, -1}, // 2y >= 1 + {0, -2, 0, 1}, // 2y <= 1 + {0, 0, 2, -1}, // 2z >= 1 + }, + {{2, 0, 0, -1}} // 2x = 1 + ) + .isIntegerEmpty()); + + EXPECT_FALSE(makeFACFromConstraints(3, + { + {2, 0, 0, -1}, // 2x >= 1 + {-3, 0, 0, 3}, // 3x <= 3 + {0, 0, 5, -6}, // 5z >= 6 + {0, 0, -7, 17}, // 7z <= 17 + {0, 3, 0, -2}, // 3y >= 2 + }, + {}) + .isIntegerEmpty()); EXPECT_FALSE(makeFACFromConstraints(3, { diff --git a/mlir/unittests/Analysis/LinearTransformTest.cpp b/mlir/unittests/Analysis/LinearTransformTest.cpp --- a/mlir/unittests/Analysis/LinearTransformTest.cpp +++ b/mlir/unittests/Analysis/LinearTransformTest.cpp @@ -22,7 +22,7 @@ // In column echelon form, each row's last non-zero value can be at most one // column to the right of the last non-zero column among the previous rows. for (unsigned row = 0, nRows = m.getNumRows(); row < nRows; ++row) { - SmallVector rowVec = transform.applyTo(m.getRow(row)); + SmallVector rowVec = transform.postMultiplyRow(m.getRow(row)); for (unsigned col = lastAllowedNonZeroCol + 1, nCols = m.getNumColumns(); col < nCols; ++col) { EXPECT_EQ(rowVec[col], 0);