From 370292719aca34521b7cdc923d9f2f17b0a1346a Mon Sep 17 00:00:00 2001 From: Christopher Bate Date: Fri, 6 Dec 2024 21:40:15 +0000 Subject: [PATCH 1/3] [mlir][presburger] Preserve relative ordering of symbols and non-symbol vars when setting up SymbolicLexSimplex The Presburger library's SymbolicLexSimplex class organizes its tableau so that symbols occuply the columns in the range `[3, 3 + num_symbols)`. When creating the SymbolicLexOpt from an IntegerRelation, the API allows the user to specify what variables of the system should be considered as symbols. The constructor then reorders the coefficient when constructing the initial constraint rows. This reordering was also manually at certain users of SymbolicLexOpt, e.g. `IntegerRelation::computeReprWithOnlyDivLocals`. Previously, at least one of the reordering implementations was not stable (in `computeReprWithOnlyDivLocals`). If it considered the the second variable to be a symbol in a 4 variable system: ``` var0 sym1 var2 var3 ``` then the simplex would get arranged as ``` sym1 var3 var2 var0 ``` I discovered tha, in certain cases, this reordering can cause `SymbolicLexOpt::computeSymbolicLexMax` to become extremely inefficient due to a large number of branches through the context tableau as well as an explosion in the size of the coefficients. See the additional test added. So, this change updates the locations where the reordering logic is implemented in order to ensure the partitioning is stable. It is possible that this is just masking a more complicated bug (hard to say for a non-expert on this code). My debugging efforts lead me to believe that there is not a bug in the simplex solver, however. Rather, some pre-processing procedures that rely on heuristics/patterns (e.g. detection of division representations) can fail, which can result in a sub-optimal problem description being handed off to the simplex solver. In any case, a non-stable partitioning of the variables is unexpected and makes it difficult to proactively try to avoid problematic system descriptions by users of the Presburger library. --- .../Analysis/Presburger/IntegerRelation.cpp | 27 +++++++------------ mlir/lib/Analysis/Presburger/Simplex.cpp | 15 ++++++----- .../Analysis/Presburger/PresburgerSetTest.cpp | 14 ++++++++++ 3 files changed, 32 insertions(+), 24 deletions(-) diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp index 74cdf567c0e56..15d453896924a 100644 --- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp @@ -229,25 +229,17 @@ PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const { // SymbolicLexOpt requires these to form a contiguous range. // // Take a copy so we can perform mutations. - IntegerRelation copy = *this; std::vector reprs(getNumLocalVars()); - copy.getLocalReprs(&reprs); + this->getLocalReprs(&reprs); - // Iterate through all the locals. The last `numNonDivLocals` are the locals - // that have been scanned already and do not have division representations. unsigned numNonDivLocals = 0; - unsigned offset = copy.getVarKindOffset(VarKind::Local); - for (unsigned i = 0, e = copy.getNumLocalVars(); i < e - numNonDivLocals;) { - if (!reprs[i]) { - // Whenever we come across a local that does not have a division - // representation, we swap it to the `numNonDivLocals`-th last position - // and increment `numNonDivLocal`s. `reprs` also needs to be swapped. - copy.swapVar(offset + i, offset + e - numNonDivLocals - 1); - std::swap(reprs[i], reprs[e - numNonDivLocals - 1]); - ++numNonDivLocals; + llvm::SmallBitVector isSymbol(getNumVars(), true); + unsigned offset = getVarKindOffset(VarKind::Local); + for (unsigned i = 0, e = getNumLocalVars(); i < e; ++i) { + if (reprs[i]) continue; - } - ++i; + isSymbol[offset + i] = false; + ++numNonDivLocals; } // If there are no non-div locals, we're done. @@ -266,9 +258,10 @@ PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const { // and the returned set of assignments to the "symbols" that makes the lexmin // unbounded. SymbolicLexOpt lexminResult = - SymbolicLexSimplex(copy, /*symbolOffset*/ 0, + SymbolicLexSimplex(*this, IntegerPolyhedron(PresburgerSpace::getSetSpace( - /*numDims=*/copy.getNumVars() - numNonDivLocals))) + /*numDims=*/getNumVars() - numNonDivLocals)), + isSymbol) .computeSymbolicIntegerLexMin(); PresburgerRelation result = lexminResult.lexopt.getDomain().unionSet(lexminResult.unboundedDomain); diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp index 4ffa2d546af4d..384dc05cbcbc0 100644 --- a/mlir/lib/Analysis/Presburger/Simplex.cpp +++ b/mlir/lib/Analysis/Presburger/Simplex.cpp @@ -64,13 +64,14 @@ SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM, const llvm::SmallBitVector &isSymbol) : SimplexBase(nVar, mustUseBigM) { assert(isSymbol.size() == nVar && "invalid bitmask!"); - // Invariant: nSymbol is the number of symbols that have been marked - // already and these occupy the columns - // [getNumFixedCols(), getNumFixedCols() + nSymbol). - for (unsigned symbolIdx : isSymbol.set_bits()) { - var[symbolIdx].isSymbol = true; - swapColumns(var[symbolIdx].pos, getNumFixedCols() + nSymbol); - ++nSymbol; + // Iterate through all the variables. Move symbols to the left and non-symbols + // to the right while preserving relative ordering. + for (unsigned i = 0; i < nVar; ++i) { + if (isSymbol[i]) { + var[i].isSymbol = true; + swapColumns(var[i].pos, getNumFixedCols() + nSymbol); + nSymbol++; + } } } diff --git a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp index 8e31a8bb2030b..7561d2044bd77 100644 --- a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp +++ b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp @@ -855,6 +855,20 @@ TEST(SetTest, computeReprWithOnlyDivLocals) { PresburgerSet(parseIntegerPolyhedron( {"(x) : (x - 3*(x floordiv 3) == 0)"})), /*numToProject=*/2); + + testComputeRepr( + parseIntegerPolyhedron("(e, a, b, c)[] : (" + "a >= 0," + "b >= 0," + "c >= 0," + "e >= 0," + "15 - a >= 0," + "7 - b >= 0," + "5 - c >= 0," + "e - a * 192 - c * 32 - b * 4 >= 0," + "3 - e + a * 192 + c * 32 + b * 4 >= 0)"), + parsePresburgerSet({"(i) : (i >= 0, i <= 3071)"}), + /*numToProject=*/3); } TEST(SetTest, subtractOutputSizeRegression) { From 64b57d302f57f559a6d9775ac2f9b42404addb3e Mon Sep 17 00:00:00 2001 From: Christopher Bate Date: Mon, 6 Jan 2025 18:56:00 +0000 Subject: [PATCH 2/3] Fix SimplexBase constructor to use stable partitioning when reordering columns. When the SimplexBase constructor reorders columns into symbol and non-symbol segments, use a stable partitioning. --- mlir/lib/Analysis/Presburger/Simplex.cpp | 26 ++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp index 384dc05cbcbc0..3b8b06dce8f25 100644 --- a/mlir/lib/Analysis/Presburger/Simplex.cpp +++ b/mlir/lib/Analysis/Presburger/Simplex.cpp @@ -68,9 +68,31 @@ SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM, // to the right while preserving relative ordering. for (unsigned i = 0; i < nVar; ++i) { if (isSymbol[i]) { + + // Move the column from its current position to the end of + // the symbols segment and update the position metadata for each column + // unknown (which should all be vars at construction time). The + // segment of non-symbol up until the current variable's column are + // shifted to the right and the current column is then moved before the + // right-shifted segment. + tableau.moveColumns(var[i].pos, 1, getNumFixedCols() + nSymbol); + + // Update the column position metadata for the unknowns associated with + // the right-shifted columns. + for (unsigned col = getNumFixedCols() + nSymbol; col < var[i].pos; ++col) + unknownFromColumn(col).pos = col + 1; + + // Perform the equivalent rearrangement on the col-to-unknown + // mapping. + std::rotate(colUnknown.begin() + getNumFixedCols() + nSymbol, + colUnknown.begin() + var[i].pos, + colUnknown.begin() + var[i].pos + 1); + + // Update the column mapping for the current variable. + var[i].pos = getNumFixedCols() + nSymbol; var[i].isSymbol = true; - swapColumns(var[i].pos, getNumFixedCols() + nSymbol); - nSymbol++; + + ++nSymbol; } } } From 34ad76a6a89fa3c8bec87ebe8d6aa24a94f04211 Mon Sep 17 00:00:00 2001 From: Christopher Bate Date: Mon, 6 Jan 2025 19:14:14 +0000 Subject: [PATCH 3/3] Address additional review comments. --- mlir/include/mlir/Analysis/Presburger/Utils.h | 3 +- .../Analysis/Presburger/IntegerRelation.cpp | 43 +++++++++---------- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h index 0e6d18279d67e..6d06daa91d376 100644 --- a/mlir/include/mlir/Analysis/Presburger/Utils.h +++ b/mlir/include/mlir/Analysis/Presburger/Utils.h @@ -96,7 +96,8 @@ enum class ReprKind { Inequality, Equality, None }; /// set to None. struct MaybeLocalRepr { ReprKind kind = ReprKind::None; - explicit operator bool() const { return kind != ReprKind::None; } + explicit operator bool() const { return hasRepr(); } + bool hasRepr() const { return kind != ReprKind::None; } union { unsigned equalityIdx; struct { diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp index 15d453896924a..d15efaa1b86f9 100644 --- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp @@ -225,27 +225,6 @@ PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const { if (getNumLocalVars() == 0) return PresburgerRelation(*this); - // Move all the non-div locals to the end, as the current API to - // SymbolicLexOpt requires these to form a contiguous range. - // - // Take a copy so we can perform mutations. - std::vector reprs(getNumLocalVars()); - this->getLocalReprs(&reprs); - - unsigned numNonDivLocals = 0; - llvm::SmallBitVector isSymbol(getNumVars(), true); - unsigned offset = getVarKindOffset(VarKind::Local); - for (unsigned i = 0, e = getNumLocalVars(); i < e; ++i) { - if (reprs[i]) - continue; - isSymbol[offset + i] = false; - ++numNonDivLocals; - } - - // If there are no non-div locals, we're done. - if (numNonDivLocals == 0) - return PresburgerRelation(*this); - // We computeSymbolicIntegerLexMin by considering the non-div locals as // "non-symbols" and considering everything else as "symbols". This will // compute a function mapping assignments to "symbols" to the @@ -257,10 +236,28 @@ PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const { // exists, which is the union of the domain of the returned lexmin function // and the returned set of assignments to the "symbols" that makes the lexmin // unbounded. + + // Construct a BitVector that identifies which variables we would like to + // treat as symbols. We want to treat all variables as "symbols" except for + // the locals that don't have a division representation. + std::vector reprs(getNumLocalVars()); + this->getLocalReprs(&reprs); + + llvm::SmallBitVector isSymbol(getNumVars(), true); + unsigned offset = getVarKindOffset(VarKind::Local); + for (unsigned i = 0, e = getNumLocalVars(); i < e; ++i) { + isSymbol[offset + i] = reprs[i].hasRepr(); + } + + // If there are no non-div locals (non-symbols), we're done. + if (isSymbol.all()) + return PresburgerRelation(*this); + SymbolicLexOpt lexminResult = - SymbolicLexSimplex(*this, + SymbolicLexSimplex(/*constraints=*/*this, + /*symbolDomain=*/ IntegerPolyhedron(PresburgerSpace::getSetSpace( - /*numDims=*/getNumVars() - numNonDivLocals)), + /*numDims=*/isSymbol.count())), isSymbol) .computeSymbolicIntegerLexMin(); PresburgerRelation result =