[MLIR] Simplex: support adding new variables dynamically

[arjunp]

Fix a bug.

[Groverkss]

Added some comments.

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

165

Please use triple / doc comments here.

mlir/lib/Analysis/Presburger/Simplex.cpp

348

An assert for i. j not exceeding number of columns would be useful here.

453–454

nit: Did you mean to write "Therefore, no constraint currently..." ?

[arjunp]

Address Kunwar's comments.

[arjunp]

Thanks for the review!

[Groverkss]

LGTM

[bondhugula]

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

166

This isn't appropriately named I think because you are allowing one or more variables. I'd prefer not having a default argument here and just naming it appendVariables. It addresses both the naming and makes it explicit for the single var case as well and is as compact.

mlir/lib/Analysis/Presburger/Simplex.cpp

483

You are missing an early:

if (count == 0)
  return

[arjunp]

Thanks for the review!

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

166

I agree with you. I only named it this way for consistency with the insert*Id/append*Id functions that were already present in FlatAffineConstraints (https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Analysis/AffineStructures.h#L258-L269). Shall I rename all of these to plural?

mlir/lib/Analysis/Presburger/Simplex.cpp

483

Thanks, added.

[bondhugula]

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

166

Oh, I see. We can leave it as is then - was a minor comment.

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

public:
   /// Returns the number of constraints in the tableau.
   unsigned numConstraints() const;
 
   /// 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);
 
+  // Add new variables to the end of the list of variables.
+  void appendVariable(unsigned count = 1);
+
   /// Mark the tableau as being empty.
   void markEmpty();
 
   /// Get a snapshot of the current state. This is used for rolling back.
   unsigned getSnapshot() const;
 
   /// Rollback to a snapshot. This invalidates all later snapshots.
   void rollback(unsigned snapshot);
private:
 
   /// Add a new row to the tableau and the associated data structures.
   unsigned addRow(ArrayRef<int64_t> coeffs);
 
   /// Normalize the given row by removing common factors between the numerator
   /// and the denominator.
   void normalizeRow(unsigned row);
 
-  /// Swap the two rows in the tableau and associated data structures.
+  /// Swap the two rows/columns in the tableau and associated data structures.
   void swapRows(unsigned i, unsigned j);
+  void swapColumns(unsigned i, unsigned j);
 
   /// Restore the unknown to a non-negative sample value.
   ///
   /// Returns true if the unknown was successfully restored to a non-negative
   /// sample value, false otherwise.
   LogicalResult restoreRow(Unknown &u);
 
   /// Compute the maximum or minimum of the specified Unknown, depending on
   /// direction. The specified unknown may be pivoted. If the unknown is
   /// restricted, it will have a non-negative sample value on return.
   /// Should not be called if the Simplex is empty.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
   Optional<Fraction> computeOptimum(Direction direction, Unknown &u);
 
   /// Mark the specified unknown redundant. This operation is added to the undo
   /// log and will be undone by rollbacks. The specified unknown must be in row
   /// orientation.
   void markRowRedundant(Unknown &u);
 
   /// Enum to denote operations that need to be undone during rollback.
   enum class UndoLogEntry {
     RemoveLastConstraint,
+    RemoveLastVariable,
     UnmarkEmpty,
     UnmarkLastRedundant
   };
 
   /// Undo the operation represented by the log entry.
   void undo(UndoLogEntry entry);
 
   /// Find a row that can be used to pivot the column in the specified

mlir/lib/Analysis/Presburger/Simplex.cpp

void Simplex::swapRows(unsigned i, unsigned j) {
   if (i == j)
     return;
   tableau.swapRows(i, j);
   std::swap(rowUnknown[i], rowUnknown[j]);
   unknownFromRow(i).pos = i;
   unknownFromRow(j).pos = j;
 }
 
+void Simplex::swapColumns(unsigned i, unsigned j) {
+  if (i == j)
+    return;
+  tableau.swapColumns(i, j);
+  std::swap(colUnknown[i], colUnknown[j]);
+  unknownFromColumn(i).pos = i;
+  unknownFromColumn(j).pos = j;
+}
+
 /// Mark this tableau empty and push an entry to the undo stack.
 void Simplex::markEmpty() {
   undoLog.push_back(UndoLogEntry::UnmarkEmpty);
   empty = true;
 }
 
 /// 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
if (entry == UndoLogEntry::RemoveLastConstraint) {
     // tableau.
     swapRows(constraint.pos, nRow - 1);
     // It is not strictly necessary to shrink the tableau, but for now we
     // maintain the invariant that the tableau has exactly nRow rows.
     tableau.resizeVertically(nRow - 1);
     nRow--;
     rowUnknown.pop_back();
     con.pop_back();
+  } else if (entry == UndoLogEntry::RemoveLastVariable) {
+    // Whenever we are rolling back the addition of a variable, it is guaranteed
+    // that the variable will be in column position.
+    //
+    // We can see this as follows: any constraint that depends on this variable
+    // was added after this variable was added, so the addition of such
+    // constraints should already have been rolled back by the time we get to
+    // rolling back the addition of the variable. Therefore, currently no
+    // constraint currently has a component along the variable, so the variable
+    // itself must be part of the basis.
+    assert(var.back().orientation == Orientation::Column &&
+           "Variable to be removed must be in column orientation!");
+    // Move this variable to the last column and remove the column from the
+    // tableau.
+    swapColumns(var.back().pos, nCol - 1);
+    tableau.resizeHorizontally(nCol - 1);
+    var.pop_back();
+    colUnknown.pop_back();
+    nCol--;
   } else if (entry == UndoLogEntry::UnmarkEmpty) {
     empty = false;
   } else if (entry == UndoLogEntry::UnmarkLastRedundant) {
     nRedundant--;
   }
 }
 
 /// Rollback to the specified snapshot.
 ///
 /// We undo all the log entries until the log size when the snapshot was taken
 /// is reached.
 void Simplex::rollback(unsigned snapshot) {
   while (undoLog.size() > snapshot) {
     undo(undoLog.back());
     undoLog.pop_back();
   }
 }
 
+void Simplex::appendVariable(unsigned count) {
+  var.reserve(var.size() + count);
+  colUnknown.reserve(colUnknown.size() + count);
+  for (unsigned i = 0; i < count; ++i) {
+    nCol++;
+    var.emplace_back(Orientation::Column, /*restricted=*/false,
+                     /*pos=*/nCol - 1);
+    colUnknown.push_back(var.size() - 1);
+  }
+  tableau.resizeHorizontally(nCol);
+  undoLog.insert(undoLog.begin(), count, UndoLogEntry::RemoveLastVariable);
+}
+
 /// Add all the constraints from the given FlatAffineConstraints.
 void Simplex::intersectFlatAffineConstraints(const FlatAffineConstraints &fac) {
   assert(fac.getNumIds() == numVariables() &&
          "FlatAffineConstraints must have same dimensionality as simplex");
   for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
     addInequality(fac.getInequality(i));
   for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
     addEquality(fac.getEquality(i));

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

TEST(SimplexTest, addInequality_already_redundant) {
   simplex.addInequality({1, -1}); // x >= 1.
   simplex.addInequality({1, 0});  // x >= 0.
   simplex.detectRedundant();
   ASSERT_FALSE(simplex.isEmpty());
   EXPECT_FALSE(simplex.isMarkedRedundant(0));
   EXPECT_TRUE(simplex.isMarkedRedundant(1));
 }
 
+TEST(SimplexTest, appendVariable) {
+  Simplex simplex(1);
+
+  unsigned snapshot1 = simplex.getSnapshot();
+  simplex.appendVariable();
+  EXPECT_EQ(simplex.numVariables(), 2u);
+
+  int64_t yMin = 2, yMax = 5;
+  simplex.addInequality({0, 1, -yMin}); // y >= 2.
+  simplex.addInequality({0, -1, yMax}); // y <= 5.
+
+  unsigned snapshot2 = simplex.getSnapshot();
+  simplex.appendVariable(2);
+  EXPECT_EQ(simplex.numVariables(), 4u);
+  simplex.rollback(snapshot2);
+
+  EXPECT_EQ(simplex.numVariables(), 2u);
+  EXPECT_EQ(simplex.numConstraints(), 2u);
+  EXPECT_EQ(simplex.computeIntegerBounds({0, 1, 0}),
+            std::make_pair(yMin, yMax));
+
+  simplex.rollback(snapshot1);
+  EXPECT_EQ(simplex.numVariables(), 1u);
+  EXPECT_EQ(simplex.numConstraints(), 0u);
+}
+
 } // namespace mlir