[MLIR][Presburger] Refactor looping strategy in coalesce

[webmiche]

Summary

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.

[Groverkss]

Instead of creating a new Simplex object each time you do a coalescePair call, can you instead precompute them and only recompute if two polyhedrons get coalesced? This should reduce the number of times you compute Simplex from n^2 to at max 2n. AFAIU creating a Simplex object is the most expensive operation in coalescing so this is useful to do.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

557

Could you write a comment that since a polyhedron may have been coalesced, polyhedrons.size() can change after an iteration.

558

For this loop, polyhedrons.size() doesn't change since you break if it changes. You can use:

for (unsigned j = i + 1, e = polyhedrons.size(); j < e; ++j)

instead.

560

This will overflow if i = 0 right?

[arjunp]

Thanks for this patch.

Please add some tests.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

503–504

I think "Erases polyhedrons[i] and polyhedrons[j] and appends poly to the vector`" would be more understandable

508–514

Why can't you append poly directly? Also, local ids are missing.

516

Please write a comment for this

516

const &p?

519

nit: can just use i here

519

521

You can return LogicalResult

531–536

You can just erase the empty poly here.

541–549

isRedundant* perform pivots, so are expensive. Early exit the loop as soon as onlyRedundantEqs becomes false.

545–550

If using the simplex array, you can use that here. Also, note in the comment that this only happens when there is exactly one empty poly left, since if there were two empty left one of them will get removed as a subset of the other.

551–554

I can't see a usecase for not coalescing in-place. I would just call it coalesceInPlace and do all the operations directly on integerPolyhedrons. We can write a function coalesce() { newSet = *this; newSet.coalesceInPlace(); return newSet; } if someone really wants a copy returned.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

551–554

But let's leave that for a different patch. Non-in-place might be useful when this is const I guess. The above coalesce would be able to be const.

[webmiche]

I implemented a slightly different looping strategy that now does not shorten the polyhedron vector, but rather replaces elements and keeps track of which ones are actually part of the polytope. Additionally, I implemented the vector of simplices and addressed comments.

[webmiche]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

541–549

I do early exit here: k < e && onlyRedundantEqsA. If you want, I can make this part of the body rather than part of the for.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

541–549

Oops, I missed that. Yeah, putting it in the body might make it easier to spot.

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

518–524

Please write a comment on what exactly this does. Also, a comment that simplices and polyhedrons are guaranteed to be at least of size 2.

590

You can check for rational emptiness here only. After this, you don't need to check that since coalesce of two non empty polyhedrons should not be empty.

594

596

606

For now, can you just iterate from j=0 to n - coalescing? I'm not convinced that this covers all cases and may cause problems later. A TODO comment for this would be good and we can discuss this in a later patch.

609

n - coalescing does not change here so you can do e = n-coalescing, j < e.

[webmiche]

Removed the addCoalescedPolyhedron function as currently only erasing is needed. Furthermore, addressed remaining comments.

[grosser]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

539

newly ???

It seems there is a very missing in the last subclause of this sentence.

[webmiche]

The IntegerPolyhedra at the end of the vector are not needed anymore, so they are erased now.

[webmiche]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

539

Is "the IntegerPolyhedron now at position i is not compared" more understandable?

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

429

assert that polyhedrons and simplices have the same length

430

You can use polyhedrons.back()

532

And similarly for simplices

[webmiche]

addressed remaining comments

[webmiche]

include the commit that actually adresses the comments ;)

[webmiche]

added a testcase with multiple coalescings happening

[arjunp]

LGTM, please address minor comments below.

mlir/include/mlir/Analysis/Presburger/Simplex.h

311 ↗

(On Diff #411115)

Can you send this as a separate patch? :)

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

413–419

Drop braces for single-line if statements

484–487

drop braces

[webmiche]

Addressed minor comments and moved the change in Simplex.h to another patch (that already landed).

[Groverkss]

LGTM.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

for (const IntegerPolyhedron &poly : integerPolyhedrons) {
     Optional<uint64_t> volume = poly.computeVolume();
     if (!volume)
       return {};
     result += *volume;
   }
   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());
+  SmallVector<IntegerPolyhedron, 2> polyhedrons = integerPolyhedrons;
+  SmallVector<Simplex, 2> simplices;
 
-  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;
+  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;
 }
 
 void PresburgerSet::print(raw_ostream &os) const {
   os << getNumPolys() << " IntegerPolyhedron:\n";
   for (const IntegerPolyhedron &poly : integerPolyhedrons) {
     poly.print(os);
     os << '\n';
   }
 }
 
 void PresburgerSet::dump() const { print(llvm::errs()); }

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

TEST(SetTest, coalesceContainedEqComplex) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x - 2 == 0, x - y == 0)",
              "(x,y) : (x - 1 >= 0, -x + 2 >= 0, x - y == 0)",
          });
   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,
                                       Optional<uint64_t> resultBound) {
   expectComputedVolumeIsValidOverapprox(set.computeVolume(), trueVolume,
                                         resultBound);
 }