[MLIR][Presburger] add Simplex:addDivisionVariable

[arjunp]

Summary

This is a convenience function for adding new divisions to the Simplex given the numerator and denominator.

This will be needed for symbolic integer lexmin support.

[Groverkss]

mlir/lib/Analysis/Presburger/Simplex.cpp

821–825

826

An assert to check if denom == 0 would be useful.

[arjunp]

Address Kunwar's comments

[Groverkss]

LGTM.

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

public:
   /// 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);
 
+  /// 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);
+
   /// Mark the tableau as being empty.
   void markEmpty();
 
   /// Get a snapshot of the current state. This is used for rolling back.
   /// The same basis will not necessarily be restored on rolling back.
   /// The snapshot only captures the set of variables and constraints present
   /// in the Simplex.
   unsigned getSnapshot() const;

mlir/lib/Analysis/Presburger/Simplex.cpp

 /// is reached.
 void SimplexBase::rollback(unsigned snapshot) {
   while (undoLog.size() > snapshot) {
     undo(undoLog.back());
     undoLog.pop_back();
   }
 }
 
+/// We add the usual floor division constraints:
+/// `0 <= coeffs - denom*q <= denom - 1`, where `q` is the new division
+/// variable.
+///
+/// This constrains the remainder `coeffs - denom*q` to be in the
+/// range `[0, denom - 1]`, which fixes the integer value of the quotient `q`.
+void SimplexBase::addDivisionVariable(ArrayRef<int64_t> coeffs, int64_t denom) {
+  assert(denom != 0 && "Cannot divide by zero!\n");
+  appendVariable();
+
+  SmallVector<int64_t, 8> ineq(coeffs.begin(), coeffs.end());
+  int64_t constTerm = ineq.back();
+  ineq.back() = -denom;
+  ineq.push_back(constTerm);
+  addInequality(ineq);
+
+  for (int64_t &coeff : ineq)
+    coeff = -coeff;
+  ineq.back() += denom - 1;
+  addInequality(ineq);
+}
+
 void SimplexBase::appendVariable(unsigned count) {
   if (count == 0)
     return;
   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,

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

TEST(SimplexTest, IsRationalSubsetOf) {
   EXPECT_FALSE(sim1.isRationalSubsetOf(s2));
   EXPECT_FALSE(sim1.isRationalSubsetOf(empty));
 
   EXPECT_TRUE(sim2.isRationalSubsetOf(univ));
   EXPECT_TRUE(sim2.isRationalSubsetOf(s1));
   EXPECT_TRUE(sim2.isRationalSubsetOf(s2));
   EXPECT_FALSE(sim2.isRationalSubsetOf(empty));
 }
+
+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]);
+}