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

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.

Reviewed By: Groverkss

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

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

@@ -409,6 +409,8 @@ public:
   /// 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.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

@@ -1104,7 +1104,20 @@ void IntegerRelation::mergeLocalIds(IntegerRelation &other) {
 
   // 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

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

@@ -135,6 +135,10 @@ PresburgerRelation::intersect(const PresburgerRelation &set) const {
 ///
 /// 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) {
@@ -142,7 +146,11 @@ static void subtractRecursively(IntegerRelation &b, Simplex &simplex,
     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
@@ -280,6 +288,10 @@ static PresburgerRelation getSetDifference(IntegerRelation disjunct,
   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);

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

@@ -457,6 +457,16 @@ TEST(SetTest, divisions) {
   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.
+  IntegerPolyhedron setA(PresburgerSpace::getSetSpace(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) {