[Mlir-commits] [mlir] 31cb999 - [MLIR][Presburger] subtract: fix bug when an input set has duplicate divisions

[Arjun P]

Author: Arjun P
Date: 2022-04-01T12:38:45+01:00
New Revision: 31cb99959f1766a23005325b6e93fe69ddafdffb

URL: https://github.com/llvm/llvm-project/commit/31cb99959f1766a23005325b6e93fe69ddafdffb
DIFF: https://github.com/llvm/llvm-project/commit/31cb99959f1766a23005325b6e93fe69ddafdffb.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
    mlir/lib/Analysis/Presburger/IntegerRelation.cpp
    mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
    mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 4a32dc90c4731..47a1ae2ae656a 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -409,6 +409,8 @@ class IntegerRelation : public PresburgerSpace {
   /// 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
index 2ae384d2bbcf6..bfa9a6539077d 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/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

diff  --git a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
index c99fc81bd3c8c..21d242b90765a 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
+++ b/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);

diff  --git a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp
index 548ba133c2854..f63870259d9ec 100644
--- a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp
+++ b/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) {