diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h --- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h @@ -409,6 +409,8 @@ /// 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. diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp --- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp @@ -1104,7 +1104,20 @@ // 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> divs; + SmallVector 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 diff --git a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp --- a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp +++ b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp @@ -135,6 +135,10 @@ /// /// 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 @@ 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 @@ 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); diff --git a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp --- a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp +++ b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp @@ -457,6 +457,16 @@ 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) {