[MLIR] FlatAffineConstraints::isIntegerEmpty: fix bug in computation of duals

[arjunp]

Summary

The method that was previously used for computing dual variables was incorrect.
This was used in the integer emptiness check algorithm, where this bug could lead to much longer running times. (Due to the way it is used, this never results in an incorrect emptiness check result.)

This patch fixes the dual computation and adds some additional asserts that catch this bug, along with regression test cases that trigger the asserts when the incorrect dual computation is used.

[Groverkss]

mlir/lib/Analysis/Presburger/Simplex.cpp

831–833

1019–1034

It's not obvious what this macro does. Could you write a comment about what it does?

[Groverkss]

mlir/unittests/Analysis/AffineStructuresTest.cpp

361–377

It would be nice to have comments for each equality/inequality as in other tests to keep consistency. Something like:

checkSample(true, makeFACFromConstraints(3,
                                         {
                                              {3, 3, 0, 3}, // 3x + 3y +3 >= 0
                                              {-4, -8, -1, 4}, // -4x -8y -z + 4 >= 0
                                              {-7, -4, 1, 1}, // ...
                                              {2, -7, -8, -7}, // ...
                                              {9, 8, -9, -7} // ...
                                          },
                                         {}));

[grosser]

Amazing. Congratulations that you identified this issue!

[bondhugula]

mlir/lib/Analysis/Presburger/Simplex.cpp

1019

Why is this under NDEBUG? The code comments completely indicate otherwise.

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

1019

This section is just for doing the assert on line 101. To do the assert it's necessary to set up the correct basis. The basis is reverted at the end of the NDEBUG section. In other words, the section under NDEBUG has no functional effect except for the assert. I'll update the documentation later today.

[grosser]

mlir/lib/Analysis/Presburger/Simplex.cpp

1019

One way to clarify this could be to move all computation in the NDEBUG into a separate function and then call this function from the assert()?

[arjunp]

Address comments.

[arjunp]

Rebase and write the regression tests using the new parser.

[arjunp]

@bondhugula gentle reminder :) can you take another look at the patch?

[Groverkss]

LGTM! I would just wait for @bondhugula to see if his concerns were addressed.

[bondhugula]

Did it previously just lead to longer running times or also incorrect emptiness results? It's unclear from the commit summary.

[arjunp]

In D113803#3167136, @bondhugula wrote:

Did it previously just lead to longer running times or also incorrect emptiness results? It's unclear from the commit summary.

It lead to computing wrong dual coefficients.

Due to the way these are used, this in turn only leads to longer running times rather than wrong emptiness check results. i.e., it can lead to a "bad" basis that leads to longer (possibly exponentially bad) running times, but the returned basis would still always be a valid basis.

[Groverkss]

LGTM

mlir/lib/Analysis/Presburger/Simplex.cpp

public:
   /// that this can be rolled back later.
   void addEqualityForDirection(ArrayRef<int64_t> dir) {
     assert(
         std::any_of(dir.begin(), dir.end(), [](int64_t x) { return x != 0; }) &&
         "Direction passed is the zero vector!");
     snapshotStack.push_back(simplex.getSnapshot());
     simplex.addEquality(getCoeffsForDirection(dir));
   }
+  /// Compute max(dotProduct(dir, x - y)).
+  Fraction computeWidth(ArrayRef<int64_t> dir) {
+    Optional<Fraction> maybeWidth =
+        simplex.computeOptimum(Direction::Up, getCoeffsForDirection(dir));
+    assert(maybeWidth.hasValue() && "Width should not be unbounded!");
+    return *maybeWidth;
+  }
 
   /// Compute max(dotProduct(dir, x - y)) and save the dual variables for only
   /// the direction equalities to `dual`.
   Fraction computeWidthAndDuals(ArrayRef<int64_t> dir,
                                 SmallVectorImpl<int64_t> &dual,
                                 int64_t &dualDenom) {
+    // We can't just call into computeWidth or computeOptimum since we need to
+    // access the state of the tableau after computing the optimum, and these
+    // functions rollback the insertion of the objective function into the
+    // tableau before returning. We instead add a row for the objective function
+    // ourselves, call into computeOptimum, compute the duals from the tableau
+    // state, and finally rollback the addition of the row before returning.
     unsigned snap = simplex.getSnapshot();
     unsigned conIndex = simplex.addRow(getCoeffsForDirection(dir));
     unsigned row = simplex.con[conIndex].pos;
     Optional<Fraction> maybeWidth =
         simplex.computeRowOptimum(Simplex::Direction::Up, row);
     assert(maybeWidth.hasValue() && "Width should not be unbounded!");
     dualDenom = simplex.tableau(row, 0);
     dual.clear();
+
     // The increment is i += 2 because equalities are added as two inequalities,
     // one positive and one negative. Each iteration processes one equality.
     for (unsigned i = simplexConstraintOffset; i < conIndex; i += 2) {
-      // The dual variable is the negative of the coefficient of the new row
-      // in the column of the constraint, if the constraint is in a column.
-      // Note that the second inequality for the equality is negated.
+      // The dual variable for an inequality in column orientation is the
+      // negative of its coefficient at the objective row. If the inequality is
+      // in row orientation, the corresponding dual variable is zero.
       //
-      // We want the dual for the original equality. If the positive inequality
-      // is in column position, the negative of its row coefficient is the
-      // desired dual. If the negative inequality is in column position, its row
-      // coefficient is the desired dual. (its coefficients are already the
-      // negated coefficients of the original equality, so we don't need to
-      // negate it now.)
+      // We want the dual for the original equality, which corresponds to two
+      // inequalities: a positive inequality, which has the same coefficients as
+      // the equality, and a negative equality, which has negated coefficients.
       //
-      // If neither are in column position, we move the negated inequality to
-      // column position. Since the inequality must have sample value zero
-      // (since it corresponds to an equality), we are free to pivot with
-      // any column. Since both the unknowns have sample value before and after
-      // pivoting, no other sample values will change and the tableau will
-      // remain consistent. To pivot, we just need to find a column that has a
-      // non-zero coefficient in this row. There must be one since otherwise the
-      // equality would be 0 == 0, which should never be passed to
-      // addEqualityForDirection.
+      // Note that at most one of these inequalities can be in column
+      // orientation because the column unknowns should form a basis and hence
+      // must be linearly independent. If the positive inequality is in column
+      // position, its dual is the dual corresponding to the equality. If the
+      // negative inequality is in column position, the negation of its dual is
+      // the dual corresponding to the equality. If neither is in column
+      // position, then that means that this equality is redundant, and its dual
+      // is zero.
       //
-      // After finding a column, we pivot with the column, after which we can
-      // get the dual from the inequality in column position as explained above.
-      if (simplex.con[i].orientation == Orientation::Column) {
+      // Note that it is NOT valid to perform pivots during the computation of
+      // the duals. This entire dual computation must be performed on the same
+      // tableau configuration.
+      assert(!(simplex.con[i].orientation == Orientation::Column &&
+               simplex.con[i + 1].orientation == Orientation::Column) &&
+             "Both inequalities for the equality cannot be in column "
+             "orientation!");
+      if (simplex.con[i].orientation == Orientation::Column)
         dual.push_back(-simplex.tableau(row, simplex.con[i].pos));
-      } else {
-        if (simplex.con[i + 1].orientation == Orientation::Row) {
-          unsigned ineqRow = simplex.con[i + 1].pos;
-          // Since it is an equality, the sample value must be zero.
-          assert(simplex.tableau(ineqRow, 1) == 0 &&
-                 "Equality's sample value must be zero.");
-          for (unsigned col = 2; col < simplex.nCol; ++col) {
-            if (simplex.tableau(ineqRow, col) != 0) {
-              simplex.pivot(ineqRow, col);
-              break;
-            }
-          }
-          assert(simplex.con[i + 1].orientation == Orientation::Column &&
-                 "No pivot found. Equality has all-zeros row in tableau!");
-        }
+      else if (simplex.con[i + 1].orientation == Orientation::Column)
         dual.push_back(simplex.tableau(row, simplex.con[i + 1].pos));
-      }
+      else
+        dual.push_back(0);
     }
     simplex.rollback(snap);
     return *maybeWidth;
   }
 
   /// Remove the last equality that was added through addEqualityForDirection.
   ///
   /// We do this by rolling back to the snapshot at the top of the stack, which
private:
   Simplex simplex;
   /// The first index of the equality constraints, the index immediately after
   /// the last constraint in the initial product simplex.
   unsigned simplexConstraintOffset;
   /// A stack of snapshots, used for rolling back.
   SmallVector<unsigned, 8> snapshotStack;
 };
 
+// Return a + scale*b;
+static SmallVector<int64_t, 8> scaleAndAdd(ArrayRef<int64_t> a, int64_t scale,
+                                           ArrayRef<int64_t> b) {
+  assert(a.size() == b.size());
+  SmallVector<int64_t, 8> res;
+  res.reserve(a.size());
+  for (unsigned i = 0, e = a.size(); i < e; ++i)
+    res.push_back(a[i] + scale * b[i]);
+  return res;
+}
+
 /// Reduce the basis to try and find a direction in which the polytope is
 /// "thin". This only works for bounded polytopes.
 ///
 /// This is an implementation of the algorithm described in the paper
 /// "An Implementation of Generalized Basis Reduction for Integer Programming"
 /// by W. Cook, T. Rutherford, H. E. Scarf, D. Shallcross.
 ///
 /// Let b_{level}, b_{level + 1}, ... b_n be the current basis.
if (dual[i - level] % dualDenom != 0) {
       basis.addToRow(i, i + 1, 1);
       widthI[1] = gbrSimplex.computeWidthAndDuals(
           basis.getRow(i + 1), candidateDual[1], candidateDualDenom[1]);
 
       unsigned j = widthI[0] < widthI[1] ? 0 : 1;
       if (j == 0)
         // Subtract 1 to go from u = ceil(dual) back to floor(dual).
         basis.addToRow(i, i + 1, -1);
+
+      // width_i(b{i+1} + u*b_i) should be minimized at our value of u.
+      // We assert that this holds by checking that the values of width_i at
+      // u - 1 and u + 1 are greater than or equal to the value at u. If the
+      // width is lesser at either of the adjacent values, then our computed
+      // value of u is clearly not the minimizer. Otherwise by convexity the
+      // computed value of u is really the minimizer.
+
+      // Check the value at u - 1.
+      assert(gbrSimplex.computeWidth(scaleAndAdd(
+                 basis.getRow(i + 1), -1, basis.getRow(i))) >= widthI[j] &&
+             "Computed u value does not minimize the width!");
+      // Check the value at u + 1.
+      assert(gbrSimplex.computeWidth(scaleAndAdd(
+                 basis.getRow(i + 1), +1, basis.getRow(i))) >= widthI[j] &&
+             "Computed u value does not minimize the width!");
+
       dual = std::move(candidateDual[j]);
       dualDenom = candidateDualDenom[j];
       return widthI[j];
     }
+
     assert(i + 1 - level < width.size() && "width_{i+1} wasn't saved");
-    // When dual minimizes f_i(b_{i+1} + dual*b_i), this is equal to
-    // width_{i+1}(b_{i+1}).
+    // f_i(b_{i+1} + dual*b_i) == width_{i+1}(b_{i+1}) when `dual` minimizes the
+    // LHS. (note: the basis has already been updated, so b_{i+1} + dual*b_i in
+    // the above expression is equal to basis.getRow(i+1) below.)
+    assert(gbrSimplex.computeWidth(basis.getRow(i + 1)) ==
+           width[i + 1 - level]);
     return width[i + 1 - level];
   };
 
   // In the ith iteration of the loop, gbrSimplex has constraints for directions
   // from `level` to i - 1.
   unsigned i = level;
   while (i < basis.getNumRows() - 1) {
     if (i >= level + width.size()) {

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, FindSampleTest) {
   checkSample(true, makeFACFromConstraints(3,
                                            {
                                                {2, 0, 0, -1}, // 2x >= 1
                                            },
                                            {{
                                                {1, -1, 0, -1}, // y = x - 1
                                                {0, 1, -1, 0},  // z = y
                                            }}));
+
+  // Regression tests for the computation of dual coefficients.
+  checkSample(false, parseFAC("(x, y, z) : ("
+                              "6*x - 4*y + 9*z + 2 >= 0,"
+                              "x + 5*y + z + 5 >= 0,"
+                              "-4*x + y + 2*z - 1 >= 0,"
+                              "-3*x - 2*y - 7*z - 1 >= 0,"
+                              "-7*x - 5*y - 9*z - 1 >= 0)",
+                              &context));
+  checkSample(true, parseFAC("(x, y, z) : ("
+                             "3*x + 3*y + 3 >= 0,"
+                             "-4*x - 8*y - z + 4 >= 0,"
+                             "-7*x - 4*y + z + 1 >= 0,"
+                             "2*x - 7*y - 8*z - 7 >= 0,"
+                             "9*x + 8*y - 9*z - 7 >= 0)",
+                             &context));
 }
 
 TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) {
   // 1 <= 5x and 5x <= 4 (no solution).
   EXPECT_TRUE(
       makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty());
   // 1 <= 5x and 5x <= 9 (solution: x = 1).
   EXPECT_FALSE(
       makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty());