[Mlir-commits] [mlir] [mlir][presburger] Preserve relative ordering of symbols and non-symbol vars when setting up SymbolicLexSimplex (PR #119036)

[Christopher Bate]

https://github.com/christopherbate updated https://github.com/llvm/llvm-project/pull/119036

>From 370292719aca34521b7cdc923d9f2f17b0a1346a Mon Sep 17 00:00:00 2001
From: Christopher Bate <cbate at nvidia.com>
Date: Fri, 6 Dec 2024 21:40:15 +0000
Subject: [PATCH] [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 74cdf567c0e569..15d453896924af 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<MaybeLocalRepr> 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 4ffa2d546af4dd..384dc05cbcbc01 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 8e31a8bb2030b6..7561d2044bd775 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) {