[MLIR][Presburger] subtract: fix bug when an input set has duplicate divisions

[arjunp]

Summary

Previously, when an input set had a duplicate division, the duplicates might
be removed by a call to mergeLocalIds due to being detected as being duplicate
for the first time. The subtraction implementation cannot handle existing
locals being removed, so this would lead to unexpected behaviour. Resolve this
by removing all the duplicates up front.

[Groverkss]

LGTM.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1099–1103

We should later make it so that the merge function returns false if the divs belong to the same relation, such that it does not remove duplicate divisions. But for now, it's fine.

[arjunp]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1099–1103

We'll have to think what that even means. Let's say the same div occurs twice in each set. What is the guarantee on the number of occurrences in the final set? Will it be the max of the number of occurrences in each set? Seems weird. I think it's better to just have the users remove duplicates upfront if they don't want this to happen. Anyway, we should eventually just make addFloorDiv not add duplicates maybe, at least by default.

1099–1103

(Not a perfect solution but that might avoid this problem entirely in many cases.)

[arjunp]

rebase

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

public:
 
   /// Removes redundant constraints using Simplex. Although the algorithm can
   /// theoretically take exponential time in the worst case (rare), it is known
   /// to perform much better in the average case. If V is the number of vertices
   /// in the polytope and C is the number of constraints, the algorithm takes
   /// O(VC) time.
   void removeRedundantConstraints();
 
+  void removeDuplicateDivs();
+
   /// Converts identifiers of kind srcKind in the range [idStart, idLimit) to
   /// variables of kind dstKind and placed after all the other variables of kind
   /// dstKind. The internal ordering among the moved variables is preserved.
   void convertIdKind(IdKind srcKind, unsigned idStart, unsigned idLimit,
                      IdKind dstKind);
   void convertToLocal(IdKind kind, unsigned idStart, unsigned idLimit) {
     convertIdKind(kind, idStart, idLimit, IdKind::Local);
   }

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void IntegerRelation::mergeLocalIds(IntegerRelation &other) {
   // have any division representation.
   std::copy(divsB.begin() + initLocals, divsB.end(),
             divsA.begin() + initLocals);
   std::copy(denomsB.begin() + initLocals, denomsB.end(),
             denomsA.begin() + initLocals);
 
   // Merge function that merges the local variables in both sets by treating
   // them as the same identifier.
   auto merge = [&relA, &relB](unsigned i, unsigned j) -> bool {
     relA.eliminateRedundantLocalId(i, j);
     relB.eliminateRedundantLocalId(i, j);
     return true;
   };
 
   // Merge all divisions by removing duplicate divisions.
   unsigned localOffset = getIdKindOffset(IdKind::Local);
-  removeDuplicateDivs(divsA, denomsA, localOffset, merge);
+  presburger::removeDuplicateDivs(divsA, denomsA, localOffset, merge);
+}
+
+void IntegerRelation::removeDuplicateDivs() {
+  std::vector<SmallVector<int64_t, 8>> divs;
+  SmallVector<unsigned, 4> denoms;
+
+  getLocalReprs(divs, denoms);
+  auto merge = [this](unsigned i, unsigned j) -> bool {
+    eliminateRedundantLocalId(i, j);
+    return true;
+  };
+  presburger::removeDuplicateDivs(divs, denoms, getIdKindOffset(IdKind::Local),
+                                  merge);
 }
 
 /// Removes local variables using equalities. Each equality is checked if it
 /// can be reduced to the form: `e = affine-expr`, where `e` is a local
 /// variable and `affine-expr` is an affine expression not containing `e`.
 /// If an equality satisfies this form, the local variable is replaced in
 /// each constraint and then removed. The equality used to replace this local
 /// variable is also removed.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 /// As a heuristic, we try adding all the constraints and check if simplex
 /// says that the intersection is empty. If it is, then subtracting this
 /// disjuncts is a no-op and we just skip it. Also, in the process we find out
 /// that some constraints are redundant. These redundant constraints are
 /// ignored.
 ///
 /// b and simplex are callee saved, i.e., their values on return are
 /// semantically equivalent to their values when the function is called.
+///
+/// b should not have duplicate divs because this might lead to existing
+/// divs disappearing in the call to mergeLocalIds below, which cannot be
+/// handled.
 static void subtractRecursively(IntegerRelation &b, Simplex &simplex,
                                 const PresburgerRelation &s, unsigned i,
                                 PresburgerRelation &result) {
   if (i == s.getNumDisjuncts()) {
     result.unionInPlace(b);
     return;
   }
+
   IntegerRelation sI = s.getDisjunct(i);
+  // Remove the duplicate divs up front to avoid them possibly disappearing
+  // in the call to mergeLocalIds below.
+  sI.removeDuplicateDivs();
 
   // 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 IntegerRelation::CountsSnapshot bCounts = b.getCounts();
   // Similarly, we also want to rollback simplex to its original state.
   const unsigned initialSnapshot = simplex.getSnapshot();

 /// const reference even though it is restored to its original state before
 /// returning from that function.
 static PresburgerRelation getSetDifference(IntegerRelation disjunct,
                                            const PresburgerRelation &set) {
   assert(disjunct.isSpaceCompatible(set) && "Spaces should match");
   if (disjunct.isEmptyByGCDTest())
     return PresburgerRelation::getEmpty(disjunct.getCompatibleSpace());
 
+  // Remove duplicate divs up front here as subtractRecursively does not support
+  // this set having duplicate divs.
+  disjunct.removeDuplicateDivs();
+
   PresburgerRelation result =
       PresburgerRelation::getEmpty(disjunct.getCompatibleSpace());
   Simplex simplex(disjunct);
   subtractRecursively(disjunct, simplex, set, 0, result);
   return result;
 }
 
 /// Return the complement of this set.

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

TEST(SetTest, divisions) {
   // 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));
 }
 
+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.
+  IntegerRelation setA(1);
+  setA.addLocalFloorDiv({1, 0}, 2);
+  setA.addLocalFloorDiv({1, 0, 0}, 2);
+  EXPECT_TRUE(setA.isEqual(setA));
+}
+
 /// Coalesce `set` and check that the `newSet` is equal to `set` and that
 /// `expectedNumPoly` matches the number of Poly in the coalesced set.
 void expectCoalesce(size_t expectedNumPoly, const PresburgerSet &set) {
   PresburgerSet newSet = set.coalesce();
   EXPECT_TRUE(set.isEqual(newSet));
   EXPECT_TRUE(expectedNumPoly == newSet.getNumDisjuncts());
 }