[mlir][presburger] Preserve relative ordering of symbols and non-symbol vars when setting up SymbolicLexSimplex

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);

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++;
+    }
   }
 }
 

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) {