[MLIR][Presburger] LexSimplex::addEquality: add equalities as fixed columns

[arjunp]

Summary

In LexSimplex, instead of adding equalities as a pair of inequalities,
add them as a single row, move them into the basis, and keep them there.

There will always be a valid basis involving all non-redundant equalities. Such
equalities will then be ignored in some other operations, such as when looking
for pivot columns. This speeds them up a little bit.

More importantly, this is an important precursor patch to adding support for
symbolic integer lexmin, as this heuristic can sometimes make a big difference there.

[Groverkss]

mlir/include/mlir/Analysis/Presburger/Simplex.h

311–314

I don't think you need to delete this. You can return numFixedCols here instead and keep using this function.

[arjunp]

• address comments
• clang-format

[Groverkss]

mlir/include/mlir/Analysis/Presburger/Simplex.h

455

Can you please write a doc comment specifying how this is different from SimplexBase::addEquality?

[arjunp]

Fix a bug by making addEquality virtual. Moved the implementation of addEquality from SimplexBase to Simplex.

[arjunp]

mlir/include/mlir/Analysis/Presburger/Simplex.h

455

Moved SimplexBase::addEquality to Simplex::addEquality. Added a doc comment to explain what LexSimplex::addEquality is doing. Simplex and LexSimplex just do things differently. I think these implementation details needn't be mentioned in the header.

[Groverkss]

LGTM.

[arjunp]

rebase

mlir/include/mlir/Analysis/Presburger/Simplex.h

public:
   /// so-called big M parameter. For more information see the documentation for
   /// LexSimplex.
   SimplexBase(unsigned nVar, bool mustUseBigM);
 
   /// Returns true if the tableau is empty (has conflicting constraints),
   /// false otherwise.
   bool isEmpty() const;
 
-  /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
-  /// is the current number of variables, then the corresponding inequality is
-  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
-  virtual void addInequality(ArrayRef<int64_t> coeffs) = 0;
-
   /// Returns the number of variables in the tableau.
   unsigned getNumVariables() const;
 
   /// Returns the number of constraints in the tableau.
   unsigned getNumConstraints() const;
 
+  /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+  /// is the current number of variables, then the corresponding inequality is
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
+  virtual void addInequality(ArrayRef<int64_t> coeffs) = 0;
+
   /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding equality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
-  void addEquality(ArrayRef<int64_t> coeffs);
+  virtual void addEquality(ArrayRef<int64_t> coeffs) = 0;
 
   /// Add new variables to the end of the list of variables.
   void appendVariable(unsigned count = 1);
 
   /// Append a new variable to the simplex and constrain it such that its only
   /// integer value is the floor div of `coeffs` and `denom`.
   void addDivisionVariable(ArrayRef<int64_t> coeffs, int64_t denom);
 
protected:
   /// Return any row that this column can be pivoted with, ignoring tableau
   /// consistency.
   ///
   /// Returns an empty optional if no pivot is possible, which happens only when
   /// the column unknown is a variable and no constraint has a non-zero
   /// coefficient for it.
   Optional<unsigned> findAnyPivotRow(unsigned col);
 
+  /// Return any column that this row can be pivoted with, ignoring tableau
+  /// consistency. Equality rows are not considered.
+  ///
+  /// Returns an empty optional if no pivot is possible, which happens only when
+  /// the column unknown is a variable and no constraint has a non-zero
+  /// coefficient for it.
+  Optional<unsigned> findAnyPivotCol(unsigned row);
+
   /// Swap the row with the column in the tableau's data structures but not the
   /// tableau itself. This is used by pivot.
   void swapRowWithCol(unsigned row, unsigned col);
 
   /// Pivot the row with the column.
   void pivot(unsigned row, unsigned col);
   void pivot(Pivot pair);
 
protected:
   void swapColumns(unsigned i, unsigned j);
 
   /// Enum to denote operations that need to be undone during rollback.
   enum class UndoLogEntry {
     RemoveLastConstraint,
     RemoveLastVariable,
     UnmarkEmpty,
     UnmarkLastRedundant,
+    UnmarkLastEquality,
     RestoreBasis
   };
 
   /// Undo the addition of the last constraint. This will only be called from
   /// undo, when rolling back.
   virtual void undoLastConstraint() = 0;
 
   /// Remove the last constraint, which must be in row orientation.
   void removeLastConstraintRowOrientation();
 
   /// Undo the operation represented by the log entry.
   void undo(UndoLogEntry entry);
 
-  /// Return the number of fixed columns, as described in the constructor above,
-  /// this is the number of columns beyond those for the variables in var.
-  unsigned getNumFixedCols() const { return usingBigM ? 3u : 2u; }
+  unsigned getNumFixedCols() const { return numFixedCols; }
 
   /// Stores whether or not a big M column is present in the tableau.
   bool usingBigM;
 
+  /// denom + const + maybe M + equality columns
+  unsigned numFixedCols;
+
   /// The number of rows in the tableau.
   unsigned nRow;
 
   /// The number of columns in the tableau, including the common denominator
   /// and the constant column.
   unsigned nCol;
 
   /// The number of redundant rows in the tableau. These are the first
public:
   ~LexSimplex() override = default;
 
   /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding inequality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
   ///
   /// This just adds the inequality to the tableau and does not try to create a
   /// consistent tableau configuration.
-  void addInequality(ArrayRef<int64_t> coeffs) final {
-    addRow(coeffs, /*makeRestricted=*/true);
-  }
+  void addInequality(ArrayRef<int64_t> coeffs) final;
+
+  /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+  /// is the current number of variables, then the corresponding equality is
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
+  void addEquality(ArrayRef<int64_t> coeffs) final;
 
   /// Get a snapshot of the current state. This is used for rolling back.
   unsigned getSnapshot() { return SimplexBase::getSnapshotBasis(); }
 
   /// Return the lexicographically minimum rational solution to the constraints.
   MaybeOptimum<SmallVector<Fraction, 8>> findRationalLexMin();
 
   /// Return the lexicographically minimum integer solution to the constraints.
   ///
   /// Note: this should be used only when the lexmin is really needed. To obtain
public:
   /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding inequality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
   ///
   /// This also tries to restore the tableau configuration to a consistent
   /// state and marks the Simplex empty if this is not possible.
   void addInequality(ArrayRef<int64_t> coeffs) final;
 
+  /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
+  /// is the current number of variables, then the corresponding equality is
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
+  void addEquality(ArrayRef<int64_t> coeffs) final;
+
   /// Compute the maximum or minimum value of the given row, depending on
   /// direction. The specified row is never pivoted. On return, the row may
   /// have a negative sample value if the direction is down.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
   MaybeOptimum<Fraction> computeRowOptimum(Direction direction, unsigned row);
 

mlir/lib/Analysis/Presburger/Simplex.cpp

 using namespace mlir;
 using namespace presburger;
 
 using Direction = Simplex::Direction;
 
 const int nullIndex = std::numeric_limits<int>::max();
 
 SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM)
-    : usingBigM(mustUseBigM), nRow(0), nCol(getNumFixedCols() + nVar),
-      nRedundant(0), tableau(0, nCol), empty(false) {
-  colUnknown.insert(colUnknown.begin(), getNumFixedCols(), nullIndex);
+    : usingBigM(mustUseBigM), numFixedCols(mustUseBigM ? 3 : 2), nRow(0),
+      nCol(numFixedCols + nVar), nRedundant(0), tableau(0, nCol), empty(false) {
+  colUnknown.insert(colUnknown.begin(), numFixedCols, nullIndex);
   for (unsigned i = 0; i < nVar; ++i) {
     var.emplace_back(Orientation::Column, /*restricted=*/false,
-                     /*pos=*/getNumFixedCols() + i);
+                     /*pos=*/numFixedCols + i);
     colUnknown.push_back(i);
   }
 }
 
 const Simplex::Unknown &SimplexBase::unknownFromIndex(int index) const {
   assert(index != nullIndex && "nullIndex passed to unknownFromIndex");
   return index >= 0 ? var[index] : con[~index];
 }

 // This means that for some p, for all t < p,
 // C(t,k) = 0 => B(t,k) = B(t,j) * B(i,k) / B(i,j) and
 // C(t,k) < 0 => B(p,k) < B(t,j) * B(i,k) / B(i,j),
 // which is in contradiction to the fact that B.col(j) / B(i,j) must be
 // lexicographically smaller than B.col(k) / B(i,k), since it lexicographically
 // minimizes the change in sample value.
 LogicalResult LexSimplex::moveRowUnknownToColumn(unsigned row) {
   Optional<unsigned> maybeColumn;
-  for (unsigned col = 3; col < nCol; ++col) {
+  for (unsigned col = getNumFixedCols(); col < nCol; ++col) {
     if (tableau(row, col) <= 0)
       continue;
     maybeColumn =
         !maybeColumn ? col : getLexMinPivotColumn(row, *maybeColumn, col);
   }
 
   if (!maybeColumn) {
     markEmpty();

 }
 
 /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
 /// is the current number of variables, then the corresponding equality is
 /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
 ///
 /// We simply add two opposing inequalities, which force the expression to
 /// be zero.
-void SimplexBase::addEquality(ArrayRef<int64_t> coeffs) {
+void Simplex::addEquality(ArrayRef<int64_t> coeffs) {
   addInequality(coeffs);
   SmallVector<int64_t, 8> negatedCoeffs;
   for (int64_t coeff : coeffs)
     negatedCoeffs.emplace_back(-coeff);
   addInequality(negatedCoeffs);
 }
 
 unsigned SimplexBase::getNumVariables() const { return var.size(); }

 // iff no constraints have been added with a non-zero coefficient for this row.
 Optional<unsigned> SimplexBase::findAnyPivotRow(unsigned col) {
   for (unsigned row = nRedundant; row < nRow; ++row)
     if (tableau(row, col) != 0)
       return row;
   return {};
 }
 
+// This doesn't find a pivot column only if the row has zero coefficients for
+// every column not marked as an equality.
+Optional<unsigned> SimplexBase::findAnyPivotCol(unsigned row) {
+  for (unsigned col = getNumFixedCols(); col < nCol; ++col)
+    if (tableau(row, col) != 0)
+      return col;
+  return {};
+}
+
 // It's not valid to remove the constraint by deleting the column since this
 // would result in an invalid basis.
 void Simplex::undoLastConstraint() {
   if (con.back().orientation == Orientation::Column) {
     // We try to find any pivot row for this column that preserves tableau
     // consistency (except possibly the column itself, which is going to be
     // deallocated anyway).
     //
if (entry == UndoLogEntry::RemoveLastConstraint) {
     tableau.resizeHorizontally(nCol - 1);
     var.pop_back();
     colUnknown.pop_back();
     nCol--;
   } else if (entry == UndoLogEntry::UnmarkEmpty) {
     empty = false;
   } else if (entry == UndoLogEntry::UnmarkLastRedundant) {
     nRedundant--;
+  } else if (entry == UndoLogEntry::UnmarkLastEquality) {
+    numFixedCols--;
+    assert(getNumFixedCols() >= 2 + usingBigM &&
+           "The denominator, constant, big M and symbols are always fixed!");
   } else if (entry == UndoLogEntry::RestoreBasis) {
     assert(!savedBases.empty() && "No bases saved!");
 
     SmallVector<int, 8> basis = std::move(savedBases.back());
     savedBases.pop_back();
 
     for (int index : basis) {
       Unknown &u = unknownFromIndex(index);
if (u.orientation == Orientation::Column) {
       // entry in the constant column divided by the denominator.
       int64_t denom = tableau(u.pos, 0);
       sample.emplace_back(tableau(u.pos, 1), denom);
     }
   }
   return sample;
 }
 
+void LexSimplex::addInequality(ArrayRef<int64_t> coeffs) {
+  addRow(coeffs, /*makeRestricted=*/true);
+}
+
+/// Try to make the equality a fixed column by finding any pivot and performing
+/// it. The only time this is not possible is when the given equality's
+/// direction is already in the span of the existing fixed column equalities. In
+/// that case, we just leave it in row position.
+void LexSimplex::addEquality(ArrayRef<int64_t> coeffs) {
+  const Unknown &u = con[addRow(coeffs, /*makeRestricted=*/true)];
+  Optional<unsigned> pivotCol = findAnyPivotCol(u.pos);
+  if (!pivotCol)
+    return;
+
+  pivot(u.pos, *pivotCol);
+  swapColumns(*pivotCol, getNumFixedCols());
+  numFixedCols++;
+  undoLog.push_back(UndoLogEntry::UnmarkLastEquality);
+}
+
 MaybeOptimum<SmallVector<Fraction, 8>> LexSimplex::getRationalSample() const {
   if (empty)
     return OptimumKind::Empty;
 
   SmallVector<Fraction, 8> sample;
   sample.reserve(var.size());
   // Push the sample value for each variable into the vector.
   for (const Unknown &u : var) {

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

TEST(SimplexTest, addDivisionVariable) {
   Simplex simplex(/*nVar=*/1);
   simplex.addDivisionVariable({1, 0}, 2);
   simplex.addInequality({1, 0, -3}); // x >= 3.
   simplex.addInequality({-1, 0, 9}); // x <= 9.
   Optional<SmallVector<int64_t, 8>> sample = simplex.findIntegerSample();
   ASSERT_TRUE(sample.hasValue());
   EXPECT_EQ((*sample)[0] / 2, (*sample)[1]);
 }
+
+TEST(LexSimplexTest, addEquality) {
+  IntegerRelation rel(/*numDomain=*/0, /*numRange=*/1);
+  rel.addEquality({1, 0});
+  LexSimplex simplex(rel);
+  EXPECT_EQ(simplex.getNumConstraints(), 1u);
+}