[MLIR][Presburger] subtraction: add support for divs defined by equalties

[Groverkss]

LGTM.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

if (level > frames.size()) {
       // will not be added.) Also add b's locals to sI, in such a way that both
       // have the same locals in the same order in the end.
       b.mergeLocalIds(sI);
 
       // Mark which inequalities of sI are division inequalities and add all
       // such inequalities to b.
       llvm::SmallBitVector canIgnoreIneq(sI.getNumInequalities() +
                                          2 * sI.getNumEqualities());
-      for (MaybeLocalRepr &maybeInequality : repr) {
+      for (MaybeLocalRepr &maybeRepr : repr) {
         assert(
-            maybeInequality.kind == ReprKind::Inequality &&
+            maybeRepr &&
             "Subtraction is not supported when a representation of the local "
             "variables of the subtrahend cannot be found!");
-        unsigned lb = maybeInequality.repr.inequalityPair.lowerBoundIdx;
-        unsigned ub = maybeInequality.repr.inequalityPair.upperBoundIdx;
+
+        if (maybeRepr.kind == ReprKind::Inequality) {
+          unsigned lb = maybeRepr.repr.inequalityPair.lowerBoundIdx;
+          unsigned ub = maybeRepr.repr.inequalityPair.upperBoundIdx;
 
           b.addInequality(sI.getInequality(lb));
           b.addInequality(sI.getInequality(ub));
 
           assert(lb != ub &&
                  "Upper and lower bounds must be different inequalities!");
           canIgnoreIneq[lb] = true;
           canIgnoreIneq[ub] = true;
+        } else {
+          assert(maybeRepr.kind == ReprKind::Equality &&
+                 "ReprKind isn't inequality so should be equality");
+          unsigned idx = maybeRepr.repr.equalityIdx;
+          b.addEquality(sI.getEquality(idx));
+          // We can ignore both inequalities corresponding to this equality.
+          unsigned offset = sI.getNumInequalities();
+          canIgnoreIneq[offset + 2 * idx] = true;
+          canIgnoreIneq[offset + 2 * idx + 1] = true;
+        }
       }
 
       unsigned offset = simplex.getNumConstraints();
       unsigned numLocalsAdded =
           b.getNumLocalIds() - initBCounts.getSpace().getNumLocalIds();
       simplex.appendVariable(numLocalsAdded);
 
       unsigned snapshotBeforeIntersect = simplex.getSnapshot();

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

TEST(SetTest, divisions) {
   expectEqual(evens.complement(), odds);
   expectEqual(odds.complement(), evens);
   // even multiples of 3 = multiples of 6.
   expectEqual(multiples3.intersect(evens), multiples6);
 
   PresburgerSet setA{parsePoly("(x) : (-x >= 0)")};
   PresburgerSet setB{parsePoly("(x) : (x floordiv 2 - 4 >= 0)")};
   EXPECT_TRUE(setA.subtract(setB).isEqual(setA));
+
+  IntegerPolyhedron evensDefByEquality(PresburgerSpace::getSetSpace(
+      /*numDims=*/1, /*numSymbols=*/0, /*numLocals=*/1));
+  evensDefByEquality.addEquality({1, -2, 0});
+  expectEqual(evens, PresburgerSet(evensDefByEquality));
 }
 
 TEST(SetTest, subtractDuplicateDivsRegression) {
   // Previously, subtracting sets with duplicate divs might result in crashes
   // due to existing divs being removed when merging local ids, due to being
   // identified as being duplicates for the first time.
   IntegerPolyhedron setA(PresburgerSpace::getSetSpace(1));
   setA.addLocalFloorDiv({1, 0}, 2);