[MLIR][Presburger] Refactor looping strategy in coalesce

This patch refactors the looping strategy of coalesce for future patches. The new strategy works in-place and uses IneqType to organize inequalities into vectors of the same type. Future coalesce cases will pattern match on this organization. E.g. the contained case needs all inequalities and equalities to be redundant, so this case becomes checking whether the respective vectors are empty. For other cases, the patterns consider the types of all inequalities of both sets making it wasteful to only consider whether a can be coalesced with b in one step, as inequalities would need to be typed again for the opposite case. Therefore, the new strategy tries to coalesce a with b and b with a in a single step.

Reviewed By: Groverkss, arjunp

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

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

@@ -401,39 +401,148 @@ Optional<uint64_t> PresburgerSet::computeVolume() const {
   return result;
 }
 
+/// Types the inequalities of `p` according to their `IneqType` for `simp` into
+/// `redundantIneqs` and `cuttingIneqs`. Returns success, if no separate
+/// inequalities were encountered. Otherwise, returns failure.
+LogicalResult
+typeInequalities(const IntegerPolyhedron &p, Simplex &simp,
+                 SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
+                 SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs) {
+  for (unsigned i = 0, e = p.getNumInequalities(); i < e; ++i) {
+    Simplex::IneqType type = simp.findIneqType(p.getInequality(i));
+    if (type == Simplex::IneqType::Redundant)
+      redundantIneqs.push_back(p.getInequality(i));
+    else if (type == Simplex::IneqType::Cut)
+      cuttingIneqs.push_back(p.getInequality(i));
+    else
+      return failure();
+  }
+  return success();
+}
+
+/// Replaces the element at position `i` with the last element and erases the
+/// last element for both `polyhedrons` and `simplices`.
+void erasePolyhedron(unsigned i,
+                     SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
+                     SmallVectorImpl<Simplex> &simplices) {
+  assert(simplices.size() == polyhedrons.size() &&
+         "simplices and polyhedrons must be equally as long");
+  polyhedrons[i] = polyhedrons.back();
+  polyhedrons.pop_back();
+  simplices[i] = simplices.back();
+  simplices.pop_back();
+}
+
+/// Attempts to coalesce the two IntegerPolyhedrons at position `i` and `j` in
+/// `polyhedrons` in-place. Returns whether the polyhedrons were successfully
+/// coalesced. The simplices in `simplices` need to be the ones constructed from
+/// `polyhedrons`. At this point, there are no empty polyhedrons in
+/// `polyhedrons` left.
+LogicalResult coalescePair(unsigned i, unsigned j,
+                           SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
+                           SmallVectorImpl<Simplex> &simplices) {
+
+  IntegerPolyhedron &a = polyhedrons[i];
+  IntegerPolyhedron &b = polyhedrons[j];
+  Simplex &simpA = simplices[i];
+  Simplex &simpB = simplices[j];
+
+  /// Check that all equalities are redundant in a (and in b).
+  bool onlyRedundantEqsA = true;
+  for (unsigned k = 0, e = a.getNumEqualities(); k < e; ++k)
+    if (!simpB.isRedundantEquality(a.getEquality(k))) {
+      onlyRedundantEqsA = false;
+      break;
+    }
+
+  bool onlyRedundantEqsB = true;
+  for (unsigned k = 0, e = b.getNumEqualities(); k < e; ++k)
+    if (!simpA.isRedundantEquality(b.getEquality(k))) {
+      onlyRedundantEqsB = false;
+      break;
+    }
+
+  /// If there are non-redundant equalities for both, exit early.
+  if (!onlyRedundantEqsB && !onlyRedundantEqsA)
+    return failure();
+
+  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsA;
+  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsA;
+
+  /// Organize all inequalities of `a` according to their type for `b` into
+  /// `redundantIneqsA` and `cuttingIneqsA` (and vice versa for all inequalities
+  /// of `b` according to their type in `a`). If a separate inequality is
+  /// encountered during typing, the two IntegerPolyhedrons cannot be coalesced.
+  if (typeInequalities(a, simpB, redundantIneqsA, cuttingIneqsA).failed())
+    return failure();
+
+  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsB;
+  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsB;
+
+  if (typeInequalities(b, simpA, redundantIneqsB, cuttingIneqsB).failed())
+    return failure();
+
+  /// If there are no cutting inequalities of `a` and all equalities of `a` are
+  /// redundant, then all constraints of `a` are redundant making `b` contained
+  /// within a (and vice versa for `b`).
+  if (cuttingIneqsA.size() == 0 && onlyRedundantEqsA) {
+    erasePolyhedron(j, polyhedrons, simplices);
+    return success();
+  }
+
+  if (cuttingIneqsB.size() == 0 && onlyRedundantEqsB) {
+    erasePolyhedron(i, polyhedrons, simplices);
+    return success();
+  }
+
+  return failure();
+}
+
 PresburgerSet PresburgerSet::coalesce() const {
   PresburgerSet newSet =
       PresburgerSet::getEmptySet(getNumDimIds(), getNumSymbolIds());
-  llvm::SmallBitVector isRedundant(getNumPolys());
-
-  for (unsigned i = 0, e = integerPolyhedrons.size(); i < e; ++i) {
-    if (isRedundant[i])
-      continue;
-    Simplex simplex(integerPolyhedrons[i]);
-
-    // Check whether the polytope of `simplex` is empty. If so, it is trivially
-    // redundant.
-    if (simplex.isEmpty()) {
-      isRedundant[i] = true;
+  SmallVector<IntegerPolyhedron, 2> polyhedrons = integerPolyhedrons;
+  SmallVector<Simplex, 2> simplices;
+
+  simplices.reserve(getNumPolys());
+  for (unsigned i = 0; i < polyhedrons.size();) {
+    Simplex simp(polyhedrons[i]);
+    if (simp.isEmpty()) {
+      polyhedrons[i] = polyhedrons[polyhedrons.size() - 1];
+      polyhedrons.pop_back();
       continue;
     }
+    i++;
+    simplices.push_back(simp);
+  }
 
-    // Check whether `IntegerPolyhedron[i]` is contained in any Poly, that is
-    // different from itself and not yet marked as redundant.
-    for (unsigned j = 0, e = integerPolyhedrons.size(); j < e; ++j) {
-      if (j == i || isRedundant[j])
+  /// For all tuples of IntegerPolyhedrons, check whether they can be coalesced.
+  /// When coalescing is successful, the contained IntegerPolyhedron is swapped
+  /// with the last element of `polyhedrons` and subsequently erased and
+  /// similarly for simplices.
+  for (unsigned i = 0; i < polyhedrons.size(); ++i) {
+
+    /// TODO: This does some comparisons two times (index 0 with 1 and index 1
+    /// with 0).
+    bool broken = false;
+    for (unsigned j = 0, e = polyhedrons.size(); j < e; ++j) {
+      if (i == j)
         continue;
-
-      if (simplex.isRationalSubsetOf(integerPolyhedrons[j])) {
-        isRedundant[i] = true;
+      if (coalescePair(i, j, polyhedrons, simplices).succeeded()) {
+        broken = true;
         break;
       }
     }
+
+    /// Only if the inner loop was not broken, i is incremented. This is
+    /// required as otherwise, if a coalescing occurs, the IntegerPolyhedron
+    /// now at position i is not compared.
+    if (!broken)
+      ++i;
   }
 
-  for (unsigned i = 0, e = integerPolyhedrons.size(); i < e; ++i)
-    if (!isRedundant[i])
-      newSet.unionPolyInPlace(integerPolyhedrons[i]);
+  for (unsigned i = 0, e = polyhedrons.size(); i < e; ++i)
+    newSet.unionPolyInPlace(polyhedrons[i]);
 
   return newSet;
 }

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

@@ -609,6 +609,16 @@ TEST(SetTest, coalesceContainedEqComplex) {
   expectCoalesce(1, set);
 }
 
+TEST(SetTest, coalesceThreeContained) {
+  PresburgerSet set =
+      parsePresburgerSetFromPolyStrings(1, {
+                                               "(x) : (x >= 0, -x + 1 >= 0)",
+                                               "(x) : (x >= 0, -x + 2 >= 0)",
+                                               "(x) : (x >= 0, -x + 3 >= 0)",
+                                           });
+  expectCoalesce(1, set);
+}
+
 static void
 expectComputedVolumeIsValidOverapprox(const PresburgerSet &set,
                                       Optional<uint64_t> trueVolume,