[MLIR][Presburger] add Simplex:addDivisionVariable

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.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D122159

3 files changed, 36 insertions, 0 deletions

diff --git a/mlir/include/mlir/Analysis/Presburger/Simplex.h b/mlir/include/mlir/Analysis/Presburger/Simplex.h
index 07d7318..8a4ea2a 100644
--- a/mlir/include/mlir/Analysis/Presburger/Simplex.h
+++ b/mlir/include/mlir/Analysis/Presburger/Simplex.h
@@ -185,6 +185,10 @@ public:
   /// 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();
 
diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp
index 8c9f017..a79bbb0 100644
--- a/mlir/lib/Analysis/Presburger/Simplex.cpp
+++ b/mlir/lib/Analysis/Presburger/Simplex.cpp
@@ -818,6 +818,28 @@ void SimplexBase::rollback(unsigned snapshot) {
   }
 }
 
+/// 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;
diff --git a/mlir/unittests/Analysis/Presburger/SimplexTest.cpp b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
index a66f1fe..e605c89 100644
--- a/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
@@ -538,3 +538,13 @@ TEST(SimplexTest, IsRationalSubsetOf) {
   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]);
+}