[MLIR][Presburger] introduce SetCoalescer

[webmiche]

Summary

This patch refactors the current coalesce implementation. It introduces
the SetCoalescer, a class in which all coalescing functionality lives.
The main advantage over the old design is the fact that the vectors of
constraints do not have to be passed around, but are implemented as
private fields of the SetCoalescer. This will become especially
important once more inequality types are introduced.

[webmiche]

Rebased.

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerSet.h

114–122 ↗

(On Diff #414362)

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

383 ↗

(On Diff #414362)

Please add some documentation about what the class does and it's usage.

383 ↗

(On Diff #414362)

Can you have the implementations and definitions for this class separate? It makes the class more readable.

386 ↗

(On Diff #414362)

const PresburgerSet &s would be better.

[webmiche]

separate declaration and implementation and address comments.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

389 ↗

(On Diff #414827)

I don't think namespace matters, nobody else can access this anyway.

396–398 ↗

(On Diff #414827)

399 ↗

(On Diff #414827)

Is this needed? Why can't we just use set.integerPolyhedrons everywhere and return set in the end? (note that it looks a little long now but it will become set.disjuncts everywhere after the other patch lands, which is ok IMO)

415–427 ↗

(On Diff #414827)

Please write docs for these members.

442–457 ↗

(On Diff #414827)

I feel this looks more like an implementation doc, move it above the impl and just write here that it tries to coalesce pairs of overlapping polys.

443–445 ↗

(On Diff #414827)

Don't repeat that redundantIneqsA are the redundant ineqs of a in every declaration doc. Also, some docs still seem to write as though these are given as arguments to the function. Please update all these, and also write what these arrays are just once, at their actual declaration in the class.

[webmiche]

Rebased and addressed comments.

[webmiche]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

399 ↗

(On Diff #414827)

AFAIU, this makes coalesce still return a newSet, rather than changing the given set in-place. I felt like changing to in-place is a bit of a subtle change for such a refactoring and it should probably get an own patch.

[webmiche]

Rebase on a (hopefully) working commit and fix a small comment

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

399 ↗

(On Diff #414827)

What I was suggesting was a little different, to operate on the copy of s, i.e., set directly. We don't need another copy newSet of the copy set. But sure, we can do that in another patch.

By the way, do we need to store set here or only the PresburgerSpace of set? AFAICT, set is only used at the end of coalesce to know the number of dims?

[webmiche]

Removed the explicit set form the SetCoalescer. It only stores the dimensions and a copy of the disjunct vector now. Updated the documentation to reflect the change and rebased.

[mehdi_amini]

Sorry I had to revert, seems like this broke some builds: https://buildkite.com/mlir/mlir-core/builds/21122#964182ab-9f3b-4c5e-a5e9-8dfcc0bf0099

[webmiche]

Fixed the build failure. It seems that gcc really wants me to declare the SetCoalescer in the header file, whereas clang does not care.

[webmiche]

move forward declaration to the top of the namespace

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

 #ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERRELATION_H
 #define MLIR_ANALYSIS_PRESBURGER_PRESBURGERRELATION_H
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 
 namespace mlir {
 namespace presburger {
 
+/// The SetCoalescer class contains all functionality concerning the coalesce
+/// heuristic. It is built from a `PresburgerRelation` and has the `coalesce()`
+/// function as its main API.
+class SetCoalescer;
+
 /// A PresburgerRelation represents a union of IntegerRelations that live in
 /// the same PresburgerSpace with support for union, intersection, subtraction,
 /// and complement operations, as well as sampling.
 ///
 /// The IntegerRelations (relations) are stored in a vector, and the set
 /// represents the union of these relations. An empty list corresponds to
 /// the empty set.
 ///
protected:
   /// Construct an empty PresburgerRelation with the specified number of
   /// dimension and symbols.
   PresburgerRelation(unsigned numDomain = 0, unsigned numRange = 0,
                      unsigned numSymbols = 0)
       : PresburgerSpace(numDomain, numRange, numSymbols) {}
 
   /// The list of disjuncts that this set is the union of.
   SmallVector<IntegerRelation, 2> integerRelations;
+
+  friend class SetCoalescer;
 };
 
 class PresburgerSet : public PresburgerRelation {
 public:
   /// Return a universe set of the specified type that contains all points.
   static PresburgerSet getUniverse(unsigned numDims = 0,
                                    unsigned numSymbols = 0);
 

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

for (const IntegerRelation &disjunct : integerRelations) {
     Optional<uint64_t> volume = disjunct.computeVolume();
     if (!volume)
       return {};
     result += *volume;
   }
   return result;
 }
 
-/// Given an IntegerRelation `p` and one of its inequalities `ineq`, check
-/// that all inequalities of `cuttingIneqs` are redundant for the facet of `p`
-/// where `ineq` holds as an equality. `simp` must be the Simplex constructed
-/// from `p`.
-static bool isFacetContained(ArrayRef<int64_t> ineq, Simplex &simp,
-                             IntegerRelation &p,
-                             ArrayRef<ArrayRef<int64_t>> cuttingIneqs) {
+/// The SetCoalescer class contains all functionality concerning the coalesce
+/// heuristic. It is built from a `PresburgerRelation` and has the `coalesce()`
+/// function as its main API. The coalesce heuristic simplifies the
+/// representation of a PresburgerRelation. In particular, it removes all
+/// disjuncts which are subsets of other disjuncts in the union and it combines
+/// sets that overlap and can be combined in a convex way.
+class presburger::SetCoalescer {
+
+public:
+  /// Simplifies the representation of a PresburgerSet.
+  PresburgerRelation coalesce();
+
+  /// Construct a SetCoalescer from a PresburgerSet.
+  SetCoalescer(const PresburgerRelation &s);
+
+private:
+  /// The dimensionality of the set the SetCoalescer is coalescing.
+  unsigned numDomainIds;
+  unsigned numRangeIds;
+  unsigned numSymbolIds;
+
+  /// The current list of `IntegerRelation`s that the currently coalesced set is
+  /// the union of.
+  SmallVector<IntegerRelation, 2> disjuncts;
+  /// The list of `Simplex`s constructed from the elements of `disjuncts`.
+  SmallVector<Simplex, 2> simplices;
+
+  /// The list of all inversed equalities during typing. This ensures that
+  /// the constraints exist even after the typing function has concluded.
+  SmallVector<SmallVector<int64_t, 2>, 2> negEqs;
+
+  /// `redundantIneqsA` is the inequalities of `a` that are redundant for `b`
+  /// (similarly for `cuttingIneqsA`, `redundantIneqsB`, and `cuttingIneqsB`).
+  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsA;
+  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsA;
+
+  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsB;
+  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsB;
+
+  /// Given a Simplex `simp` and one of its inequalities `ineq`, check
+  /// that the facet of `simp` where `ineq` holds as an equality is contained
+  /// within `a`.
+  bool isFacetContained(ArrayRef<int64_t> ineq, Simplex &simp);
+
+  /// Adds `disjunct` to `disjuncts` and removes the disjuncts at position `i`
+  /// and `j`. Updates `simplices` to reflect the changes. `i` and `j` cannot
+  /// be equal.
+  void addCoalescedDisjunct(unsigned i, unsigned j,
+                            const IntegerRelation &disjunct);
+
+  /// Checks whether `a` and `b` can be combined in a convex sense, if there
+  /// exist cutting inequalities.
+  ///
+  /// An example of this case:
+  ///    ___________        ___________
+  ///   /   /  |   /       /          /
+  ///   \   \  |  /   ==>  \         /
+  ///    \   \ | /          \       /
+  ///     \___\|/            \_____/
+  ///
+  ///
+  LogicalResult coalescePairCutCase(unsigned i, unsigned j);
+
+  /// Types the inequality `ineq` according to its `IneqType` for `simp` into
+  /// `redundantIneqsB` and `cuttingIneqsB`. Returns success, if no separate
+  /// inequalities were encountered. Otherwise, returns failure.
+  LogicalResult typeInequality(ArrayRef<int64_t> ineq, Simplex &simp);
+
+  /// Types the equality `eq`, i.e. for `eq` == 0, types both `eq` >= 0 and
+  /// -`eq` >= 0 according to their `IneqType` for `simp` into
+  /// `redundantIneqsB` and `cuttingIneqsB`. Returns success, if no separate
+  /// inequalities were encountered. Otherwise, returns failure.
+  LogicalResult typeEquality(ArrayRef<int64_t> eq, Simplex &simp);
+
+  /// Replaces the element at position `i` with the last element and erases
+  /// the last element for both `disjuncts` and `simplices`.
+  void eraseDisjunct(unsigned i);
+
+  /// Attempts to coalesce the two IntegerRelations at position `i` and `j`
+  /// in `disjuncts` in-place. Returns whether the disjuncts were
+  /// successfully coalesced. The simplices in `simplices` need to be the ones
+  /// constructed from `disjuncts`. At this point, there are no empty
+  /// disjuncts in `disjuncts` left.
+  LogicalResult coalescePair(unsigned i, unsigned j);
+};
+
+/// Constructs a `SetCoalescer` from a `PresburgerRelation`. Only adds non-empty
+/// `IntegerRelation`s to the `disjuncts` vector.
+SetCoalescer::SetCoalescer(const PresburgerRelation &s) {
+
+  disjuncts = s.integerRelations;
+
+  simplices.reserve(s.getNumDisjuncts());
+  // Note that disjuncts.size() changes during the loop.
+  for (unsigned i = 0; i < disjuncts.size();) {
+    Simplex simp(disjuncts[i]);
+    if (simp.isEmpty()) {
+      disjuncts[i] = disjuncts[disjuncts.size() - 1];
+      disjuncts.pop_back();
+      continue;
+    }
+    ++i;
+    simplices.push_back(simp);
+  }
+  numDomainIds = s.getNumDomainIds();
+  numRangeIds = s.getNumRangeIds();
+  numSymbolIds = s.getNumSymbolIds();
+}
+
+/// Simplifies the representation of a PresburgerSet.
+PresburgerRelation SetCoalescer::coalesce() {
+  // For all tuples of IntegerRelations, check whether they can be
+  // coalesced. When coalescing is successful, the contained IntegerRelation
+  // is swapped with the last element of `disjuncts` and subsequently erased
+  // and similarly for simplices.
+  for (unsigned i = 0; i < disjuncts.size();) {
+
+    // TODO: This does some comparisons two times (index 0 with 1 and index 1
+    // with 0).
+    bool broken = false;
+    for (unsigned j = 0, e = disjuncts.size(); j < e; ++j) {
+      negEqs.clear();
+      redundantIneqsA.clear();
+      redundantIneqsB.clear();
+      cuttingIneqsA.clear();
+      cuttingIneqsB.clear();
+      if (i == j)
+        continue;
+      if (coalescePair(i, j).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 IntegerRelation
+    // now at position i is not compared.
+    if (!broken)
+      ++i;
+  }
+
+  PresburgerRelation newSet =
+      PresburgerRelation::getEmpty(numDomainIds, numRangeIds, numSymbolIds);
+  for (unsigned i = 0, e = disjuncts.size(); i < e; ++i)
+    newSet.unionInPlace(disjuncts[i]);
+
+  return newSet;
+}
+
+/// Given a Simplex `simp` and one of its inequalities `ineq`, check
+/// that all inequalities of `cuttingIneqsB` are redundant for the facet of
+/// `simp` where `ineq` holds as an equality is contained within `a`.
+bool SetCoalescer::isFacetContained(ArrayRef<int64_t> ineq, Simplex &simp) {
   unsigned snapshot = simp.getSnapshot();
   simp.addEquality(ineq);
-  if (llvm::any_of(cuttingIneqs, [&simp](ArrayRef<int64_t> curr) {
+  if (llvm::any_of(cuttingIneqsB, [&simp](ArrayRef<int64_t> curr) {
         return !simp.isRedundantInequality(curr);
       })) {
     simp.rollback(snapshot);
     return false;
   }
   simp.rollback(snapshot);
   return true;
 }
 
-/// Adds `disjunct` to `disjuncts` and removes the disjuncts at position `i` and
-/// `j`. Updates `simplices` to reflect the changes. `i` and `j` cannot be
-/// equal.
-static void addCoalescedDisjunct(SmallVectorImpl<IntegerRelation> &disjuncts,
-                                 unsigned i, unsigned j,
-                                 const IntegerRelation &disjunct,
-                                 SmallVectorImpl<Simplex> &simplices) {
+void SetCoalescer::addCoalescedDisjunct(unsigned i, unsigned j,
+                                        const IntegerRelation &disjunct) {
   assert(i != j && "The indices must refer to different disjuncts");
-
   unsigned n = disjuncts.size();
   if (j == n - 1) {
-    // This case needs special handling since position `n` - 1 is removed from
-    // the vector, hence the `IntegerRelation` at position `n` - 2 is lost
-    // otherwise.
+    // This case needs special handling since position `n` - 1 is removed
+    // from the vector, hence the `IntegerRelation` at position `n` - 2 is
+    // lost otherwise.
     disjuncts[i] = disjuncts[n - 2];
     disjuncts.pop_back();
     disjuncts[n - 2] = disjunct;
 
     simplices[i] = simplices[n - 2];
     simplices.pop_back();
     simplices[n - 2] = Simplex(disjunct);
 
   } else {
-    // Other possible edge cases are correct since for `j` or `i` == `n` - 2,
-    // the `IntegerRelation` at position `n` - 2 should be lost. The case
-    // `i` == `n` - 1 makes the first following statement a noop. Hence, in this
-    // case the same thing is done as above, but with `j` rather than `i`.
+    // Other possible edge cases are correct since for `j` or `i` == `n` -
+    // 2, the `IntegerRelation` at position `n` - 2 should be lost. The
+    // case `i` == `n` - 1 makes the first following statement a noop.
+    // Hence, in this case the same thing is done as above, but with `j`
+    // rather than `i`.
     disjuncts[i] = disjuncts[n - 1];
     disjuncts[j] = disjuncts[n - 2];
     disjuncts.pop_back();
     disjuncts[n - 2] = disjunct;
 
     simplices[i] = simplices[n - 1];
     simplices[j] = simplices[n - 2];
     simplices.pop_back();
     simplices[n - 2] = Simplex(disjunct);
   }
 }
 
-/// Given two disjuncts `a` and `b` at positions `i` and `j` in `disjuncts`
-/// and `redundantIneqsA` being the inequalities of `a` that are redundant for
-/// `b` (similarly for `cuttingIneqsA`, `redundantIneqsB`, and `cuttingIneqsB`),
-/// checks whether the facets of all cutting inequalites of `a` are contained in
-/// `b`. If so, a new disjunct consisting of all redundant inequalites of `a`
-/// and `b` and all equalities of both is created.
+/// Given two polyhedra `a` and `b` at positions `i` and `j` in
+/// `disjuncts` and `redundantIneqsA` being the inequalities of `a` that
+/// are redundant for `b` (similarly for `cuttingIneqsA`, `redundantIneqsB`,
+/// and `cuttingIneqsB`), Checks whether the facets of all cutting
+/// inequalites of `a` are contained in `b`. If so, a new polyhedron
+/// consisting of all redundant inequalites of `a` and `b` and all
+/// equalities of both is created.
 ///
 /// An example of this case:
 ///    ___________        ___________
 ///   /   /  |   /       /          /
 ///   \   \  |  /   ==>  \         /
 ///    \   \ | /          \       /
 ///     \___\|/            \_____/
 ///
 ///
-static LogicalResult
-coalescePairCutCase(SmallVectorImpl<IntegerRelation> &disjuncts,
-                    SmallVectorImpl<Simplex> &simplices, unsigned i, unsigned j,
-                    ArrayRef<ArrayRef<int64_t>> redundantIneqsA,
-                    ArrayRef<ArrayRef<int64_t>> cuttingIneqsA,
-                    ArrayRef<ArrayRef<int64_t>> redundantIneqsB,
-                    ArrayRef<ArrayRef<int64_t>> cuttingIneqsB) {
+LogicalResult SetCoalescer::coalescePairCutCase(unsigned i, unsigned j) {
   /// All inequalities of `b` need to be redundant. We already know that the
   /// redundant ones are, so only the cutting ones remain to be checked.
   Simplex &simp = simplices[i];
   IntegerRelation &disjunct = disjuncts[i];
-  if (llvm::any_of(cuttingIneqsA, [&simp, &disjunct,
-                                   &cuttingIneqsB](ArrayRef<int64_t> curr) {
-        return !isFacetContained(curr, simp, disjunct, cuttingIneqsB);
+  if (llvm::any_of(cuttingIneqsA, [this, &simp](ArrayRef<int64_t> curr) {
+        return !isFacetContained(curr, simp);
       }))
     return failure();
   IntegerRelation newSet(disjunct.getNumDomainIds(), disjunct.getNumRangeIds(),
                          disjunct.getNumSymbolIds(), disjunct.getNumLocalIds());
 
   for (ArrayRef<int64_t> curr : redundantIneqsA)
     newSet.addInequality(curr);
 
   for (ArrayRef<int64_t> curr : redundantIneqsB)
     newSet.addInequality(curr);
 
-  addCoalescedDisjunct(disjuncts, i, j, newSet, simplices);
+  addCoalescedDisjunct(i, j, newSet);
   return success();
 }
 
-/// Types the inequality `ineq` according to its `IneqType` for `simp` into
-/// `redundantIneqs` and `cuttingIneqs`. Returns success, if no separate
-/// inequalities were encountered. Otherwise, returns failure.
-static LogicalResult
-typeInequality(ArrayRef<int64_t> ineq, Simplex &simp,
-               SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
-               SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs) {
+LogicalResult SetCoalescer::typeInequality(ArrayRef<int64_t> ineq,
+                                           Simplex &simp) {
   Simplex::IneqType type = simp.findIneqType(ineq);
   if (type == Simplex::IneqType::Redundant)
-    redundantIneqs.push_back(ineq);
+    redundantIneqsB.push_back(ineq);
   else if (type == Simplex::IneqType::Cut)
-    cuttingIneqs.push_back(ineq);
+    cuttingIneqsB.push_back(ineq);
   else
     return failure();
   return success();
 }
 
-/// Types the equality `eq`, i.e. for `eq` == 0, types both `eq` >= 0 and -`eq`
-/// >= 0 according to their `IneqType` for `simp` into `redundantIneqs` and
-/// `cuttingIneqs`. Returns success, if no separate inequalities were
-/// encountered. Otherwise, returns failure.
-static LogicalResult
-typeEquality(ArrayRef<int64_t> eq, Simplex &simp,
-             SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
-             SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs,
-             SmallVectorImpl<SmallVector<int64_t, 2>> &negEqs) {
-  if (typeInequality(eq, simp, redundantIneqs, cuttingIneqs).failed())
+LogicalResult SetCoalescer::typeEquality(ArrayRef<int64_t> eq, Simplex &simp) {
+  if (typeInequality(eq, simp).failed())
     return failure();
   negEqs.push_back(getNegatedCoeffs(eq));
   ArrayRef<int64_t> inv(negEqs.back());
-  if (typeInequality(inv, simp, redundantIneqs, cuttingIneqs).failed())
+  if (typeInequality(inv, simp).failed())
     return failure();
   return success();
 }
 
-/// Replaces the element at position `i` with the last element and erases the
-/// last element for both `disjuncts` and `simplices`.
-static void eraseDisjunct(unsigned i,
-                          SmallVectorImpl<IntegerRelation> &disjuncts,
-                          SmallVectorImpl<Simplex> &simplices) {
+void SetCoalescer::eraseDisjunct(unsigned i) {
   assert(simplices.size() == disjuncts.size() &&
          "simplices and disjuncts must be equally as long");
   disjuncts[i] = disjuncts.back();
   disjuncts.pop_back();
   simplices[i] = simplices.back();
   simplices.pop_back();
 }
 
-/// Attempts to coalesce the two IntegerRelations at position `i` and `j` in
-/// `disjuncts` in-place. Returns whether the disjuncts were successfully
-/// coalesced. The simplices in `simplices` need to be the ones constructed from
-/// `disjuncts`. At this point, there are no empty disjuncts in
-/// `disjuncts` left.
-static LogicalResult coalescePair(unsigned i, unsigned j,
-                                  SmallVectorImpl<IntegerRelation> &disjuncts,
-                                  SmallVectorImpl<Simplex> &simplices) {
+LogicalResult SetCoalescer::coalescePair(unsigned i, unsigned j) {
 
   IntegerRelation &a = disjuncts[i];
   IntegerRelation &b = disjuncts[j];
-  /// Handling of local ids is not yet implemented, so these cases are skipped.
+  /// Handling of local ids is not yet implemented, so these cases are
+  /// skipped.
   /// TODO: implement local id support.
   if (a.getNumLocalIds() != 0 || b.getNumLocalIds() != 0)
     return failure();
   Simplex &simpA = simplices[i];
   Simplex &simpB = simplices[j];
 
-  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsA;
-  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsA;
-  SmallVector<SmallVector<int64_t, 2>, 2> negEqs;
-
-  // Organize all inequalities and equalities 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 IntegerRelations cannot
-  // be coalesced.
+  // Organize all inequalities and equalities 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 IntegerRelations
+  // cannot be coalesced.
   for (int k = 0, e = a.getNumInequalities(); k < e; ++k)
-    if (typeInequality(a.getInequality(k), simpB, redundantIneqsA,
-                       cuttingIneqsA)
-            .failed())
+    if (typeInequality(a.getInequality(k), simpB).failed())
       return failure();
 
   for (int k = 0, e = a.getNumEqualities(); k < e; ++k)
-    if (typeEquality(a.getEquality(k), simpB, redundantIneqsA, cuttingIneqsA,
-                     negEqs)
-            .failed())
+    if (typeEquality(a.getEquality(k), simpB).failed())
       return failure();
 
-  SmallVector<ArrayRef<int64_t>, 2> redundantIneqsB;
-  SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsB;
+  std::swap(redundantIneqsA, redundantIneqsB);
+  std::swap(cuttingIneqsA, cuttingIneqsB);
+
   for (int k = 0, e = b.getNumInequalities(); k < e; ++k)
-    if (typeInequality(b.getInequality(k), simpA, redundantIneqsB,
-                       cuttingIneqsB)
-            .failed())
+    if (typeInequality(b.getInequality(k), simpA).failed())
       return failure();
 
   for (int k = 0, e = b.getNumEqualities(); k < e; ++k)
-    if (typeEquality(b.getEquality(k), simpA, redundantIneqsB, cuttingIneqsB,
-                     negEqs)
-            .failed())
+    if (typeEquality(b.getEquality(k), simpA).failed())
       return failure();
 
   // If there are no cutting inequalities of `a`, `b` is contained
-  // within `a` (and vice versa for `b`).
+  // within `a`.
   if (cuttingIneqsA.empty()) {
-    eraseDisjunct(j, disjuncts, simplices);
+    eraseDisjunct(j);
     return success();
   }
 
-  if (cuttingIneqsB.empty()) {
-    eraseDisjunct(i, disjuncts, simplices);
+  // Try to apply the cut case
+  if (coalescePairCutCase(i, j).succeeded())
     return success();
-  }
 
-  // Try to apply the cut case
-  if (coalescePairCutCase(disjuncts, simplices, i, j, redundantIneqsA,
-                          cuttingIneqsA, redundantIneqsB, cuttingIneqsB)
-          .succeeded())
+  // Swap the vectors to compare the pair (j,i) instead of (i,j).
+  std::swap(redundantIneqsA, redundantIneqsB);
+  std::swap(cuttingIneqsA, cuttingIneqsB);
+
+  // If there are no cutting inequalities of `a`, `b` is contained
+  // within `a`.
+  if (cuttingIneqsA.empty()) {
+    eraseDisjunct(i);
     return success();
+  }
 
-  if (coalescePairCutCase(disjuncts, simplices, j, i, redundantIneqsB,
-                          cuttingIneqsB, redundantIneqsA, cuttingIneqsA)
-          .succeeded())
+  // Try to apply the cut case
+  if (coalescePairCutCase(j, i).succeeded())
     return success();
 
   return failure();
 }
 
 PresburgerRelation PresburgerRelation::coalesce() const {
-  PresburgerRelation newSet = PresburgerRelation::getEmpty(
-      getNumDomainIds(), getNumRangeIds(), getNumSymbolIds());
-  SmallVector<IntegerRelation, 2> disjuncts = integerRelations;
-  SmallVector<Simplex, 2> simplices;
-
-  simplices.reserve(getNumDisjuncts());
-  // Note that disjuncts.size() changes during the loop.
-  for (unsigned i = 0; i < disjuncts.size();) {
-    Simplex simp(disjuncts[i]);
-    if (simp.isEmpty()) {
-      disjuncts[i] = disjuncts[disjuncts.size() - 1];
-      disjuncts.pop_back();
-      continue;
-    }
-    ++i;
-    simplices.push_back(simp);
-  }
-
-  // For all tuples of IntegerRelations, check whether they can be coalesced.
-  // When coalescing is successful, the contained IntegerRelation is swapped
-  // with the last element of `disjuncts` and subsequently erased and
-  // similarly for simplices.
-  for (unsigned i = 0; i < disjuncts.size();) {
-
-    // TODO: This does some comparisons two times (index 0 with 1 and index 1
-    // with 0).
-    bool broken = false;
-    for (unsigned j = 0, e = disjuncts.size(); j < e; ++j) {
-      if (i == j)
-        continue;
-      if (coalescePair(i, j, disjuncts, 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 IntegerRelation
-    // now at position i is not compared.
-    if (!broken)
-      ++i;
-  }
-
-  for (unsigned i = 0, e = disjuncts.size(); i < e; ++i)
-    newSet.unionInPlace(disjuncts[i]);
-
-  return newSet;
+  return SetCoalescer(*this).coalesce();
 }
 
 void PresburgerRelation::print(raw_ostream &os) const {
   os << "Number of Disjuncts: " << getNumDisjuncts() << "\n";
   for (const IntegerRelation &disjunct : integerRelations) {
     disjunct.print(os);
     os << '\n';
   }