[MLIR] PresburgerSet subtraction: fix bug where the set `b` was not restored properly on return

[arjunp]

Summary

When subtracting b \ c, when there are divisions in c, these division
constraints get added to b. b must be restored to its original state
when returning, but these added divisions constraints were not removed in
one of the return paths. This patch fixes this and deduplicates the
restoration logic by encapuslating it in a lambda restoreState. The patch
also includes a regression test for the bug fix.

[Groverkss]

mlir/lib/Analysis/PresburgerSet.cpp

208–209

Could you maybe remove the last x added divisions instead? The problem with this is, in case mergeLocalIds removes any duplicate inequalities (I plan to do that sometime), it will cause problems here. Removing the last x seems better to me.

mlir/unittests/Analysis/PresburgerSetTest.cpp

664

Could you maybe add a test case where division merging also takes place? Something that has maybe more divisions:

A: [x/5] [y/2] [x/5]
B: [y/2] [x/7] [y/2]

I think this will be useful in the future.

Not exactly sure what kind of test, but something that tests merging and has some duplicate divisions. (It's fine if they are two test cases checking duplicate and merging separately).

[arjunp]

mlir/lib/Analysis/PresburgerSet.cpp

208–209

Well, then we have to manually keep track of how many constraints we have added which is error-prone. What exactly do you mean by removing duplicate inequalities?

mlir/unittests/Analysis/PresburgerSetTest.cpp

664

Sure, we should definitely add such tests. But I suppose that would be a separate patch, since it's not really related to the bug of forgetting to rollback added constraints?

[Groverkss]

LGTM.

mlir/lib/Analysis/PresburgerSet.cpp

208–209

I think explaining what I mean would be better when I actually push the changes. This looks fine to me for now.

mlir/unittests/Analysis/PresburgerSetTest.cpp

664

Sure, we can do this as a separate patch.

[nicolasvasilache]

mlir/lib/Analysis/PresburgerSet.cpp

204

this feels like you want an RAII llvm::scope_exit to handle all the branches for you automatically.

[arjunp]

mlir/lib/Analysis/PresburgerSet.cpp

204

Whoops, nice catch. This is an outdated comment. I eventually decided against it because only one of the return paths actually needs to restore at scope exit, the other one already restores earlier.

mlir/lib/Analysis/PresburgerSet.cpp

 static void subtractRecursively(FlatAffineConstraints &b, Simplex &simplex,
                                 const PresburgerSet &s, unsigned i,
                                 PresburgerSet &result) {
   if (i == s.getNumFACs()) {
     result.unionFACInPlace(b);
     return;
   }
   FlatAffineConstraints sI = s.getFlatAffineConstraints(i);
-  unsigned bInitNumLocals = b.getNumLocalIds();
+
+  // Below, we append some additional constraints and ids to b. We want to
+  // rollback b to its initial state before returning, which we will do by
+  // removing all constraints beyond the original number of inequalities
+  // and equalities, so we store these counts first.
+  const unsigned bInitNumIneqs = b.getNumInequalities();
+  const unsigned bInitNumEqs = b.getNumEqualities();
+  const unsigned bInitNumLocals = b.getNumLocalIds();
+  // Similarly, we also want to rollback simplex to its original state.
+  const unsigned initialSnapshot = simplex.getSnapshot();
+
+  // Automatically restore the original state when we return.
+  auto restoreState = [&]() {
+    b.removeIdRange(FlatAffineConstraints::IdKind::Local, bInitNumLocals,
+                    b.getNumLocalIds());
+    b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
+    b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
+    simplex.rollback(initialSnapshot);
+  };
 
   // Find out which inequalities of sI correspond to division inequalities for
   // the local variables of sI.
   std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
       sI.getNumLocalIds());
   sI.getLocalReprs(repr);
 
   // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
for (Optional<std::pair<unsigned, unsigned>> &maybePair : repr) {
     b.addInequality(sI.getInequality(maybePair->second));
 
     assert(maybePair->first != maybePair->second &&
            "Upper and lower bounds must be different inequalities!");
     isDivInequality[maybePair->first] = true;
     isDivInequality[maybePair->second] = true;
   }
 
-  unsigned initialSnapshot = simplex.getSnapshot();
   unsigned offset = simplex.getNumConstraints();
   unsigned numLocalsAdded = b.getNumLocalIds() - bInitNumLocals;
   simplex.appendVariable(numLocalsAdded);
 
   unsigned snapshotBeforeIntersect = simplex.getSnapshot();
   simplex.intersectFlatAffineConstraints(sI);
 
   if (simplex.isEmpty()) {
     /// b ^ s_i is empty, so b \ s_i = b. We move directly to i + 1.
-    simplex.rollback(initialSnapshot);
-    b.removeIdRange(FlatAffineConstraints::IdKind::Local, bInitNumLocals,
-                    b.getNumLocalIds());
+    /// We are ignoring level i completely, so we restore the state
+    /// *before* going to level i + 1.
+    restoreState();
     subtractRecursively(b, simplex, s, i + 1, result);
     return;
   }
 
   simplex.detectRedundant();
 
   // Equalities are added to simplex as a pair of inequalities.
   unsigned totalNewSimplexInequalities =
static void subtractRecursively(FlatAffineConstraints &b, Simplex &simplex,
   // is not satisfied, and then add the ineq to b and simplex because
   // ineq must be satisfied by all later parts.
   auto processInequality = [&](ArrayRef<int64_t> ineq) {
     recurseWithInequality(getComplementIneq(ineq));
     b.addInequality(ineq);
     simplex.addInequality(ineq);
   };
 
-  // processInequality appends some additional constraints to b. We want to
-  // rollback b to its initial state before returning, which we will do by
-  // removing all constraints beyond the original number of inequalities
-  // and equalities, so we store these counts first.
-  unsigned bInitNumIneqs = b.getNumInequalities();
-  unsigned bInitNumEqs = b.getNumEqualities();
-
   // Process all the inequalities, ignoring redundant inequalities and division
   // inequalities. The result is correct whether or not we ignore these, but
   // ignoring them makes the result simpler.
   for (unsigned j = 0, e = sI.getNumInequalities(); j < e; j++) {
     if (isMarkedRedundant[j])
       continue;
     if (isDivInequality[j])
       continue;
for (unsigned j = 0, e = sI.getNumEqualities(); j < e; ++j) {
     // represented in terms of inequalities and not equalities, so we do not
     // check for division inequalities here.
     if (!isMarkedRedundant[offset + 2 * j])
       processInequality(coeffs);
     if (!isMarkedRedundant[offset + 2 * j + 1])
       processInequality(getNegatedCoeffs(coeffs));
   }
 
-  // Rollback b and simplex to their initial states.
-  b.removeIdRange(FlatAffineConstraints::IdKind::Local, bInitNumLocals,
-                  b.getNumLocalIds());
-  b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
-  b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
-
-  simplex.rollback(initialSnapshot);
+  restoreState();
 }
 
 /// Return the set difference fac \ set.
 ///
 /// The FAC here is modified in subtractRecursively, so it cannot be a const
 /// reference even though it is restored to its original state before returning
 /// from that function.
 PresburgerSet PresburgerSet::getSetDifference(FlatAffineConstraints fac,

mlir/unittests/Analysis/PresburgerSetTest.cpp

 #include "mlir/Analysis/PresburgerSet.h"
 #include "./AffineStructuresParser.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace mlir {
 
+/// Parses a FlatAffineConstraints from a StringRef. It is expected that the
+/// string represents a valid IntegerSet, otherwise it will violate a gtest
+/// assertion.
+static FlatAffineConstraints parseFAC(StringRef str, MLIRContext *context) {
+  FailureOr<FlatAffineConstraints> fac = parseIntegerSetToFAC(str, context);
+
+  EXPECT_TRUE(succeeded(fac));
+
+  return *fac;
+}
+
 /// Compute the union of s and t, and check that each of the given points
 /// belongs to the union iff it belongs to at least one of s and t.
 static void testUnionAtPoints(PresburgerSet s, PresburgerSet t,
                               ArrayRef<SmallVector<int64_t, 4>> points) {
   PresburgerSet unionSet = s.unionSet(t);
   for (const SmallVector<int64_t, 4> &point : points) {
     bool inS = s.containsPoint(point);
     bool inT = t.containsPoint(point);

 
 void expectEqual(const PresburgerSet &s, const PresburgerSet &t) {
   EXPECT_TRUE(s.isEqual(t));
 }
 
 void expectEmpty(PresburgerSet s) { EXPECT_TRUE(s.isIntegerEmpty()); }
 
 TEST(SetTest, divisions) {
+  MLIRContext context;
   // Note: we currently need to add the equalities as inequalities to the FAC
   // since detecting divisions based on equalities is not yet supported.
 
   // evens = {x : exists q, x = 2q}.
   PresburgerSet evens{
       makeFACFromConstraints(2, {{1, -2, 0}, {-1, 2, 1}}, {{1, -2, 0}}, 1)};
   // odds = {x : exists q, x = 2q + 1}.
   PresburgerSet odds{
       makeFACFromConstraints(2, {{1, -2, 0}, {-1, 2, 1}}, {{1, -2, -1}}, 1)};
   // multiples6 = {x : exists q, x = 6q}.
   PresburgerSet multiples3{
       makeFACFromConstraints(2, {{1, -3, 0}, {-1, 3, 2}}, {{1, -3, 0}}, 1)};
   // multiples6 = {x : exists q, x = 6q}.
   PresburgerSet multiples6{
       makeFACFromConstraints(2, {{1, -6, 0}, {-1, 6, 5}}, {{1, -6, 0}}, 1)};
 
   // evens /\ odds = empty.
   expectEmpty(PresburgerSet(evens).intersect(PresburgerSet(odds)));
   // evens U odds = universe.
   expectEqual(evens.unionSet(odds), PresburgerSet::getUniverse(1));
   expectEqual(evens.complement(), odds);
   expectEqual(odds.complement(), evens);
   // even multiples of 3 = multiples of 6.
   expectEqual(multiples3.intersect(evens), multiples6);
-}
-
-/// Parses a FlatAffineConstraints from a StringRef. It is expected that the
-/// string represents a valid IntegerSet, otherwise it will violate a gtest
-/// assertion.
-static FlatAffineConstraints parseFAC(StringRef str, MLIRContext *context) {
-  FailureOr<FlatAffineConstraints> fac = parseIntegerSetToFAC(str, context);
-
-  EXPECT_TRUE(succeeded(fac));
 
-  return *fac;
+  PresburgerSet setA =
+      makeSetFromFACs(1, {parseFAC("(x) : (-x >= 0)", &context)});
+  PresburgerSet setB =
+      makeSetFromFACs(1, {parseFAC("(x) : (x floordiv 2 - 4 >= 0)", &context)});
+  EXPECT_TRUE(setA.subtract(setB).isEqual(setA));
 }
 
 /// Coalesce `set` and check that the `newSet` is equal to `set and that
 /// `expectedNumFACs` matches the number of FACs in the coalesced set.
 /// If one of the two
 void expectCoalesce(size_t expectedNumFACs, const PresburgerSet set) {
   PresburgerSet newSet = set.coalesce();
   EXPECT_TRUE(set.isEqual(newSet));