[MLIR] Introduce coalesce for PresburgerSet

[webmiche]

Summary

This patch provides functionality for simplifying PresburgerSets by checking if any FlatAffineConstraints in the set is contained in another, and removing such redundant FACs.

This is part of a series of patches to provide functionality for integer set coalescing in MLIR.

[webmiche]

• update FACUtils license

[webmiche]

update the revision to the proper HEAD state

[webmiche]

• add header guard to FACUtils.h

[arjunp]

Hi, thanks for the patch!

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

241

Maybe clarify the name as isRedundantInequality.

Maybe even add a function isRedundantEquality for equalities as well as some tests for it.

mlir/include/mlir/Analysis/PresburgerSet.h

99

Nit: inparticular -> in particular

subset-> subsets

mlir/lib/Analysis/Presburger/Simplex.cpp

1244

nit: maybe remove this blank line

1253

Don't you have to check both inequalities that make up the equality?

1259

A bit of documentation might be useful here.

mlir/lib/Analysis/PresburgerSet.cpp

386

nit: I think you can drop the this-> everywhere. We're only operating on one PresburgerSet.

387

You can use SmallBitVector.

410

I think you can drop these braces.

mlir/unittests/Analysis/FACUtils.h

1 ↗

(On Diff #375517)

IMO you can just name it Utils.h.

9 ↗

(On Diff #375517)

I believe it wants you to remove the trailing underscore.

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

10

Not sure about this. Let's see what the others think.

416–418

nit: The numbering [0] [1] [2] doesn't seem to be used, so you can probably remove it.

429–439

nit: I feel these blank lines aren't needed.

mlir/unittests/Analysis/PresburgerSetTest.cpp

614

Please add some tests involving equalities.

614

Please add a set with zero FACs in the set.

[arjunp]

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

244

all inequalities -> all constraints

mlir/lib/Analysis/PresburgerSet.cpp

398

Please write a complete sentence; capitalize the first letter and put a full stop at the end. Also, you can use backticks to refer to variables like bs1.

[Groverkss]

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

243–245

Could you mention what happens in the case either of the sets is empty?

mlir/lib/Analysis/PresburgerSet.cpp

389

Use preincrement: ++i instead of i++.

mlir/unittests/Analysis/FACUtils.h

19–22 ↗

(On Diff #375517)

The same function is used in AffineStructuresTest. If you decide to use this utils file, you can also use it there.

[webmiche]

• address review comments

[webmiche]

mlir/lib/Analysis/Presburger/Simplex.cpp

1268

the both -> both the

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

1259–1268

What are subplanes? do you mean hyperplanes?

mlir/lib/Analysis/PresburgerSet.cpp

387

Maybe we can call it something more descriptive, e.g., isRedundant or something.

401

I think there is nothing called bs1 in the code now.

[webmiche]

• update codumentation

mlir/lib/Analysis/Presburger/Simplex.cpp

1259–1268

I was thinking in 2D 😅 . I think in 3D, one would describe the thing I am thinking of as a half-space.

[bondhugula]

Great to see this. Separately, I think it'll be good to create in the future some high-level documentation of PresburgerSet and the underlying machinery being used (GBR, simplex, etc.) as part of a larger document perhaps on Affine Analysis. This can go into the doc tree and appear at mlir.llvm.org.

[bondhugula]

mlir/unittests/Analysis/Utils.h

16–17

FACUTILS_H -> UTILS_H

26–47

We shouldn't be having these definitions in the header file but in the .cpp file. Is this some sort of an exception for unit tests includes?

[bondhugula]

mlir/unittests/Analysis/PresburgerSetTest.cpp

614

Doc comment here.

614

const ref for set?

[webmiche]

split Utils.h, add possible const refs

Utils.h was split into Utils.h and Utils.cpp. Furthermore, addressed the remaining comments, in particular the PresburgerSet in coalesce is const now.

[rriddle]

mlir/unittests/Analysis/Utils.h

15

Can you rename this to something that highlights that this is for affine analyses? Utils is fairly generic of a name.

[webmiche]

mlir/unittests/Analysis/Utils.h

15

I agree. What do you suggest? FACUtils? FACGenerationUtils?

[rriddle]

mlir/unittests/Analysis/Utils.h

15

Either of those is fine, I don't have a strong preference.

[webmiche]

rename Utils -> FACGenerationUtils

[arjunp]

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

248

248–249

Personally I feel that these lines make this sound more complex than it really is. Maybe just call the function isSubsetOf and say that it returns true if this Simplex's polytope is a subset of the given FAC; then hopefully everyone knows that empty set is subset of empty set, etc.

mlir/lib/Analysis/Presburger/Simplex.cpp

1268

It is not clear which two half spaces "both half-spaces defined by coeffs" is referring to.

[bondhugula]

mlir/unittests/Analysis/FACGenerationUtils.cpp

24 ↗

(On Diff #378916)

Use ArrayRef instead of const SmallVector ref.

26 ↗

(On Diff #378916)

Likewise.

mlir/unittests/Analysis/Presburger/CMakeLists.txt

8–9

Please place the libraries each on a separate line.

[webmiche]

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

248–249

As discussed, I will rename this function to isRationalSubsetOf. Notice that there can be cases where an FAC is a subset in terms of integer points, but has some additional rational points. In this case, this function returns false.

A full subset check at this point would be too slow to justify for a heuristic mainly targeted at runtime optimization.

mlir/lib/Analysis/Presburger/Simplex.cpp

1268

I am thinking here about the one for which coeffs is >= 0 and the one where coeffs is <= 0. If both hold, then we know that coeffs== 0, so it holds as an equality.

I am unsure how to write this better, while still justifying why we need both the minimum and the maximum.

/// Computes the minimum and maximum value `coeffs` can take. If the
/// minimum is greater than 0 and the maximum is less than 0, `coeffs`>= 0
/// and `coeffs` <= 0 respectively hold, hence `coeffs`== 0 and
/// the equality defined by `coeffs` holds in the entire polytope.

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

1268

since 0 <= minimum <= maximum <= 0, the checks be changed to just minimum && maximum && *minimum == 0 && *maximum == 0 instead right? Then you can write the doc:

Check whether the equality given by `coeffs == 0` is redundant given the existing constraints. This is redundant when `coeffs` is already always zero under the existing constraints. `coeffs` is always zero when the minimum and maximum value that `coeffs` can take are both zero.

[webmiche]

• address review comments, update unittests to use AffineParser

I was able to delete all the FACGenerationUtils.{h,cpp}, as I am now using the AffineStructuresParser. Feedback on that is very welcome! :)

[arjunp]

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

246–249

You can simplify the documentation now as I mentioned -- you can just say it returns true if this Simplex's polytope is a subset of the given FAC

[webmiche]

• update isRationalSubsetOf comment

[webmiche]

• fix typo

[bondhugula]

This revision appears to be pending for quite some time. Is there a plan of action here? Is this waiting on reviewers? @arjunp @Groverkss

[arjunp]

LGTM. Please fix the nits below though.

mlir/lib/Analysis/Presburger/Simplex.cpp

1258

1260

[webmiche]

updated isRedundantInequality comment

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

public:
   /// Returns an integer sample point if one exists, or None
   /// otherwise. This should only be called for bounded sets.
   Optional<SmallVector<int64_t, 8>> findIntegerSample();
 
   /// Print the tableau's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
+  /// Check if the specified inequality already holds in the polytope.
+  bool isRedundantInequality(ArrayRef<int64_t> coeffs);
+
+  /// Check if the specified equality already holds in the polytope.
+  bool isRedundantEquality(ArrayRef<int64_t> coeffs);
+
+  /// Check if the specified FAC contains the polytope,
+  /// i.e. whether all constraints of the specified FAC hold for the polytope.
+  /// If polytope is empty, returns always true.
+  /// If `fac` is empty, returns false unless the polytope is empty.
+  bool containedIn(const FlatAffineConstraints &fac);
+
 private:
   friend class GBRSimplex;
 
   enum class Orientation { Row, Column };
 
   /// An Unknown is either a variable or a constraint. It is always associated
   /// with either a row or column. Whether it's a row or a column is specified
   /// by the orientation and pos identifies the specific row or column it is

mlir/include/mlir/Analysis/PresburgerSet.h

public:
   /// Return true if all the sets in the union are known to be integer empty
   /// false otherwise.
   bool isIntegerEmpty() const;
 
   /// Find an integer sample from the given set. This should not be called if
   /// any of the FACs in the union are unbounded.
   bool findIntegerSample(SmallVectorImpl<int64_t> &sample);
 
+  /// Simplifies the representation of a PresburgerSet.
+  ///
+  /// In particular, removes all FACs which are subsets of other FACs in the
+  /// union.
+  PresburgerSet coalesce() const;
+
 private:
   /// Construct an empty PresburgerSet.
   PresburgerSet(unsigned nDim = 0, unsigned nSym = 0)
       : nDim(nDim), nSym(nSym) {}
 
   /// Return the set difference fac \ set.
   static PresburgerSet getSetDifference(FlatAffineConstraints fac,
                                         const PresburgerSet &set);

mlir/lib/Analysis/Presburger/Simplex.cpp

for (unsigned col = 0; col < nCol; ++col)
       os << tableau(row, col) << '\t';
     os << '\n';
   }
   os << '\n';
 }
 
 void Simplex::dump() const { print(llvm::errs()); }
 
+bool Simplex::containedIn(const FlatAffineConstraints &fac) {
+  if (isEmpty())
+    return true;
+
+  for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
+    if (!isRedundantInequality(fac.getInequality(i)))
+      return false;
+
+  for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
+    if (!isRedundantEquality(fac.getEquality(i)))
+      return false;
+
+  return true;
+}
+
+/// Computes the minimum value `coeffs` can take. If the value is greater
+/// than zero, the polytope entirely lies in the half-space defined
+/// by `coeffs`.
+bool Simplex::isRedundantInequality(ArrayRef<int64_t> coeffs) {
+  Optional<Fraction> minimum = computeOptimum(Direction::Down, coeffs);
+  return minimum && *minimum >= Fraction(0, 1);
+}
+
+/// Computes the minimum and maximum value `coeffs` can take. If the
+/// minimum is greater than 0 and the maximum is less than 0, the
+/// polytope is entirely contained in both half-spaces defined by `coeffs`,
+/// hence it is contained in the equality defined by `coeffs`.
+bool Simplex::isRedundantEquality(ArrayRef<int64_t> coeffs) {
+  Optional<Fraction> minimum = computeOptimum(Direction::Down, coeffs);
+  Optional<Fraction> maximum = computeOptimum(Direction::Up, coeffs);
+  return minimum && maximum && *maximum <= Fraction(0, 1) &&
+         *minimum >= Fraction(0, 1);
+}
+
 } // namespace mlir

mlir/lib/Analysis/PresburgerSet.cpp

for (const FlatAffineConstraints &fac : flatAffineConstraints) {
     if (Optional<SmallVector<int64_t, 8>> opt = fac.findIntegerSample()) {
       sample = std::move(*opt);
       return true;
     }
   }
   return false;
 }
 
+PresburgerSet PresburgerSet::coalesce() const {
+  PresburgerSet newSet = PresburgerSet::getEmptySet(getNumDims(), getNumSyms());
+  llvm::SmallBitVector isRedundant(getNumFACs());
+
+  for (unsigned i = 0, e = flatAffineConstraints.size(); i < e; ++i) {
+    if (isRedundant[i])
+      continue;
+    Simplex simplex(flatAffineConstraints[i]);
+
+    // Check whether the polytope of `simplex` is empty. If so, it is trivially
+    // redundant.
+    if (simplex.isEmpty()) {
+      isRedundant[i] = true;
+      continue;
+    }
+
+    // Check whether `FlatAffineConstraints[i]` is contained in any FAC, that is
+    // different from itself and not yet marked as redundant.
+    for (unsigned j = 0, e = flatAffineConstraints.size(); j < e; ++j) {
+      if (j == i || isRedundant[j])
+        continue;
+
+      if (simplex.containedIn(flatAffineConstraints[j])) {
+        isRedundant[i] = true;
+        break;
+      }
+    }
+  }
+
+  for (unsigned i = 0, e = flatAffineConstraints.size(); i < e; ++i)
+    if (!isRedundant[i])
+      newSet.unionFACInPlace(flatAffineConstraints[i]);
+
+  return newSet;
+}
+
 void PresburgerSet::print(raw_ostream &os) const {
   os << getNumFACs() << " FlatAffineConstraints:\n";
   for (const FlatAffineConstraints &fac : flatAffineConstraints) {
     fac.print(os);
     os << '\n';
   }
 }
 
 void PresburgerSet::dump() const { print(llvm::errs()); }

mlir/unittests/Analysis/AffineStructuresTest.cpp

 //===- AffineStructuresTest.cpp - Tests for AffineStructures ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
+#include "./Utils.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
 
case TestFunction::Sample:
     }
     break;
   case TestFunction::Empty:
     EXPECT_EQ(!hasSample, fac.isIntegerEmpty());
     break;
   }
 }
 
-/// Construct a FlatAffineConstraints from a set of inequality and
-/// equality constraints.
-static FlatAffineConstraints
-makeFACFromConstraints(unsigned ids, ArrayRef<SmallVector<int64_t, 4>> ineqs,
-                       ArrayRef<SmallVector<int64_t, 4>> eqs,
-                       unsigned syms = 0) {
-  FlatAffineConstraints fac(ineqs.size(), eqs.size(), ids + 1, ids - syms, syms,
-                            /*numLocals=*/0);
-  for (const auto &eq : eqs)
-    fac.addEquality(eq);
-  for (const auto &ineq : ineqs)
-    fac.addInequality(ineq);
-  return fac;
-}
-
 /// Check sampling for all the permutations of the dimensions for the given
 /// constraint set. Since the GBR algorithm progresses dimension-wise, different
 /// orderings may cause the algorithm to proceed differently. At least some of
 ///.these permutations should make it past the heuristics and test the
 /// implementation of the GBR algorithm itself.
 /// Use TestFunction fn to test.
 static void checkPermutationsSample(bool hasSample, unsigned nDim,
                                     ArrayRef<SmallVector<int64_t, 4>> ineqs,
TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) {
   // Set with symbols.
   FlatAffineConstraints fac6 = makeFACFromConstraints(2,
                                                       {
                                                           {1, 1, 0},
                                                       },
                                                       {
                                                           {1, -1, 0},
                                                       },
-                                                      1);
+                                                      0, 1);
   EXPECT_FALSE(fac6.isIntegerEmpty());
 }
 
 TEST(FlatAffineConstraintsTest, removeRedundantConstraintsTest) {
   FlatAffineConstraints fac = makeFACFromConstraints(1,
                                                      {
                                                          {1, -2}, // x >= 2.
                                                          {-1, 2}  // x <= 2.

mlir/unittests/Analysis/CMakeLists.txt

 add_mlir_unittest(MLIRAnalysisTests
   AffineStructuresTest.cpp
   LinearTransformTest.cpp
   PresburgerSetTest.cpp
+  Utils.cpp
 )
 
 target_link_libraries(MLIRAnalysisTests
   PRIVATE MLIRLoopAnalysis)
 
 add_subdirectory(Presburger)

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   MatrixTest.cpp
   SimplexTest.cpp
+  ../Utils.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
-  PRIVATE MLIRPresburger)
-
+  PRIVATE MLIRPresburger
+  MLIRLoopAnalysis)

mlir/unittests/Analysis/Presburger/SimplexTest.cpp

 //===- SimplexTest.cpp - Tests for Simplex --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Simplex.h"
+#include "../Utils.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace mlir {
 
 /// Take a snapshot, add constraints making the set empty, and rollback.
 /// The set should not be empty after rolling back.
TEST(SimplexTest, appendVariable) {
   EXPECT_EQ(simplex.computeIntegerBounds({0, 1, 0}),
             std::make_pair(yMin, yMax));
 
   simplex.rollback(snapshot1);
   EXPECT_EQ(simplex.getNumVariables(), 1u);
   EXPECT_EQ(simplex.getNumConstraints(), 0u);
 }
 
+TEST(SimplexTest, isRedundantInequality) {
+  Simplex simplex(2);
+  simplex.addInequality({0, -1, 2}); // y <= 2.
+  simplex.addInequality({1, 0, 0});  // x >= 0.
+  simplex.addEquality({-1, 1, 0});   // y = x.
+
+  EXPECT_TRUE(simplex.isRedundantInequality({-1, 0, 2})); // x <= 2.
+  EXPECT_TRUE(simplex.isRedundantInequality({0, 1, 0}));  // y >= 0.
+
+  EXPECT_FALSE(simplex.isRedundantInequality({-1, 0, -1})); // x <= -1.
+  EXPECT_FALSE(simplex.isRedundantInequality({0, 1, -2}));  // y >= 2.
+  EXPECT_FALSE(simplex.isRedundantInequality({0, 1, -1}));  // y >= 1.
+}
+
+TEST(SimplexTest, isRedundantEquality) {
+  Simplex simplex(2);
+  simplex.addInequality({0, -1, 2}); // y <= 2.
+  simplex.addInequality({1, 0, 0});  // x >= 0.
+  simplex.addEquality({-1, 1, 0});   // y = x.
+
+  EXPECT_TRUE(simplex.isRedundantEquality({-1, 1, 0})); // y = x.
+  EXPECT_TRUE(simplex.isRedundantEquality({1, -1, 0})); // x = y.
+
+  EXPECT_FALSE(simplex.isRedundantEquality({0, 1, -1})); // y = 1.
+
+  simplex.addEquality({0, -1, 2}); // y = 2.
+
+  EXPECT_TRUE(simplex.isRedundantEquality({-1, 0, 2})); // x = 2.
+}
+
+TEST(SimplexTest, ContainedIn) {
+
+  FlatAffineConstraints univ = FlatAffineConstraints::getUniverse(1, 0);
+  FlatAffineConstraints empty = makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                     {-1, -1}}); // x <= -1.
+  FlatAffineConstraints s1 = makeFACFromIneqs(1, {{1, 0},        // x >= 0.
+                                                  {-1, 4}});     // x <= 4.
+  FlatAffineConstraints s2 = makeFACFromIneqs(1, {{1, -1},       // x >= 1.
+                                                  {-1, 3}});     // x <= 3.
+
+  Simplex simUniv(univ);
+  Simplex simEmpty(empty);
+  Simplex sim1(s1);
+  Simplex sim2(s2);
+
+  EXPECT_TRUE(simUniv.containedIn(univ));
+  EXPECT_TRUE(simEmpty.containedIn(empty));
+  EXPECT_TRUE(sim1.containedIn(s1));
+  EXPECT_TRUE(sim2.containedIn(s2));
+
+  EXPECT_TRUE(simEmpty.containedIn(univ));
+  EXPECT_TRUE(simEmpty.containedIn(s1));
+  EXPECT_TRUE(simEmpty.containedIn(s2));
+  EXPECT_TRUE(simEmpty.containedIn(empty));
+
+  EXPECT_TRUE(simUniv.containedIn(univ));
+  EXPECT_FALSE(simUniv.containedIn(s1));
+  EXPECT_FALSE(simUniv.containedIn(s2));
+  EXPECT_FALSE(simUniv.containedIn(empty));
+
+  EXPECT_TRUE(sim1.containedIn(univ));
+  EXPECT_TRUE(sim1.containedIn(s1));
+  EXPECT_FALSE(sim1.containedIn(s2));
+  EXPECT_FALSE(sim1.containedIn(empty));
+
+  EXPECT_TRUE(sim2.containedIn(univ));
+  EXPECT_TRUE(sim2.containedIn(s1));
+  EXPECT_TRUE(sim2.containedIn(s2));
+  EXPECT_FALSE(sim2.containedIn(empty));
+}
+
 } // namespace mlir

mlir/unittests/Analysis/PresburgerSetTest.cpp

 // intersection, subtract, and complement work by computing the operation on
 // two sets and checking, for a set of points, that the resulting set contains
 // the point iff the result is supposed to contain it. The test for isEqual just
 // checks if the result for two sets matches the expected result.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/PresburgerSet.h"
+#include "./Utils.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace mlir {
 
 /// Compute the union of s and t, and check that each of the given points
 /// belongs to the union iff it belongs to at least one of s and t.
for (const SmallVector<int64_t, 4> &point : points) {
     bool inComplement = complement.containsPoint(point);
     if (inS)
       EXPECT_FALSE(inComplement);
     else
       EXPECT_TRUE(inComplement);
   }
 }
 
-/// Construct a FlatAffineConstraints from a set of inequality and
-/// equality constraints. `numIds` is the total number of ids, of which
-/// `numLocals` is the number of local ids.
-static FlatAffineConstraints
-makeFACFromConstraints(unsigned numIds, ArrayRef<SmallVector<int64_t, 4>> ineqs,
-                       ArrayRef<SmallVector<int64_t, 4>> eqs,
-                       unsigned numLocals = 0) {
-  FlatAffineConstraints fac(/*numReservedInequalities=*/ineqs.size(),
-                            /*numReservedEqualities=*/eqs.size(),
-                            /*numReservedCols=*/numIds + 1,
-                            /*numDims=*/numIds - numLocals,
-                            /*numSymbols=*/0, numLocals);
-  for (const SmallVector<int64_t, 4> &eq : eqs)
-    fac.addEquality(eq);
-  for (const SmallVector<int64_t, 4> &ineq : ineqs)
-    fac.addInequality(ineq);
-  return fac;
-}
-
-/// Construct a FlatAffineConstraints having `numDims` dimensions from the given
-/// set of inequality constraints. This is a convenience function to be used
-/// when the FAC to be constructed does not have any local ids and does not have
-/// equalties.
-static FlatAffineConstraints
-makeFACFromIneqs(unsigned numDims, ArrayRef<SmallVector<int64_t, 4>> ineqs) {
-  return makeFACFromConstraints(numDims, ineqs, /*eqs=*/{});
-}
-
 /// Construct a PresburgerSet having `numDims` dimensions and no symbols from
 /// the given list of FlatAffineConstraints. Each FAC in `facs` should also have
 /// `numDims` dimensions and no symbols, although it can have any number of
 /// local ids.
 static PresburgerSet makeSetFromFACs(unsigned numDims,
                                      ArrayRef<FlatAffineConstraints> facs) {
   PresburgerSet set = PresburgerSet::getEmptySet(numDims);
   for (const FlatAffineConstraints &fac : facs)
TEST(SetTest, divisions) {
   // evens U odds = universe.
   expectEqual(evens.unionSet(odds), PresburgerSet::getUniverse(1));
   expectEqual(evens.complement(), odds);
   expectEqual(odds.complement(), evens);
   // even multiples of 3 = multiples of 6.
   expectEqual(multiples3.intersect(evens), multiples6);
 }
 
+/// Coalesce `set` and check that the `newSet` is equal to `set and that
+/// `expectedNumFACs` matches the number of FACs in the coalesced set.
+/// If one of the two
+void expectCoalesce(size_t expectedNumFACs, const PresburgerSet set) {
+  PresburgerSet newSet = set.coalesce();
+  EXPECT_TRUE(set.isEqual(newSet));
+  EXPECT_TRUE(expectedNumFACs == newSet.getNumFACs());
+}
+
+TEST(SetTest, coalesceNoFAC) {
+  PresburgerSet set = makeSetFromFACs(0, {});
+  expectCoalesce(0, set);
+}
+
+TEST(SetTest, coalesceContainedOneDim) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, 0},    // x >= 0.
+                                                  {-1, 4}}), // x <= 4.
+                             makeFACFromIneqs(1, {{1, -1},   // x >= 1.
+                                                  {-1, 2}}), // x <= 2.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceFirstEmpty) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                  {-1, -1}}), // x <= -1.
+                             makeFACFromIneqs(1, {{1, -1},    // x >= 1.
+                                                  {-1, 2}}),  // x <= 2.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceSecondEmpty) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, -1},    // x >= 1.
+                                                  {-1, 2}}),  // x <= 2.
+                             makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                  {-1, -1}}), // x <= -1.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceBothEmpty) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, -3},    // x >= 3.
+                                                  {-1, -1}}), // x <= -1.
+                             makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                  {-1, -1}}), // x <= -1.
+                         });
+  expectCoalesce(0, set);
+}
+
+TEST(SetTest, coalesceFirstUniv) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {}),
+                             makeFACFromIneqs(1, {{1, 0},    // x >= 0.
+                                                  {-1, 1}}), // x <= 1.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceSecondUniv) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, 0},    // x >= 0.
+                                                  {-1, 1}}), // x <= 1.
+                             makeFACFromIneqs(1, {}),
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceBothUniv) {
+  PresburgerSet set = makeSetFromFACs(1, {
+                                             makeFACFromIneqs(1, {}),
+                                             makeFACFromIneqs(1, {}),
+                                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceFirstUnivSecondEmpty) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {}),
+                             makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                  {-1, -1}}), // x <= -1.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceFirstEmptySecondUniv) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {}),
+                             makeFACFromIneqs(1, {{1, 0},     // x >= 0.
+                                                  {-1, -1}}), // x <= -1.
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceCutOneDim) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, 0},    // x >= 0.
+                                                  {-1, 3}}), // x <= 3.
+                             makeFACFromIneqs(1, {{1, -2},   // x >= 2.
+                                                  {-1, 4}}), // x <= 4.
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceSeparateOneDim) {
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, 0},    // x >= 0.
+                                                  {-1, 2}}), // x <= 2.
+                             makeFACFromIneqs(1, {{1, -3},   // x >= 3.
+                                                  {-1, 4}}), // x <= 4.
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceContainedTwoDim) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, 0},    // y >= 0
+                                                  {0, -1, 3}}), // y <= 3.
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, -2},   // y >= 2.
+                                                  {0, -1, 3}}), // y <= 3
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceCutTwoDim) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, 0},    // y >= 0
+                                                  {0, -1, 2}}), // y <= 2.
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, -1},   // y >= 1.
+                                                  {0, -1, 3}}), // y <= 3
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceSeparateTwoDim) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, 0},    // y >= 0
+                                                  {0, -1, 1}}), // y <= 1.
+                             makeFACFromIneqs(2, {{1, 0, 0},    // x >= 0.
+                                                  {-1, 0, 3},   // x <= 3.
+                                                  {0, 1, -2},   // y >= 2.
+                                                  {0, -1, 3}}), // y <= 3
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceContainedEq) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, 0},    // x >= 0.
+                                                     {-1, 0, 3}},  // x <= 3.
+                                                    {{1, -1, 0}}), // x = y
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, -1},  // x >= 1.
+                                                     {-1, 0, 2}}, // x <= 2.
+                                                    {{1, -1, 0}}) // x = y
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceCuttingEq) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, -1},   // x >= 1.
+                                                     {-1, 0, 3}},  // x <= 3.
+                                                    {{1, -1, 0}}), // x = y
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, 0},   // x >= 0.
+                                                     {-1, 0, 2}}, // x <= 2.
+                                                    {{1, -1, 0}}) // x = y
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceSeparateEq) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, -3},   // x >= 3.
+                                                     {-1, 0, 4}},  // x <= 4.
+                                                    {{1, -1, 0}}), // x = y
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, 0},   // x >= 0.
+                                                     {-1, 0, 1}}, // x <= 1.
+                                                    {{1, -1, 0}}) // x = y
+                         });
+  expectCoalesce(2, set);
+}
+
+TEST(SetTest, coalesceContainedEqAsIneq) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, 0},   // x >= 0.
+                                                     {-1, 0, 3},  // x <= 3.
+                                                     {1, -1, 0},  // x >= y
+                                                     {-1, 1, 0}}, // x <= y
+                                                    {}),
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, -1},  // x >= 1.
+                                                     {-1, 0, 2}}, // x <= 2.
+                                                    {{1, -1, 0}}) // x = y
+                         });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceContainedEqComplex) {
+  PresburgerSet set =
+      makeSetFromFACs(2, {
+                             makeFACFromConstraints(2, {},
+                                                    {{1, -1, 0},   // x = y
+                                                     {1, 0, -2}}), // x = 2
+                             makeFACFromConstraints(2,
+                                                    {{1, 0, -1},  // x >= 1.
+                                                     {-1, 0, 2}}, // x <= 2.
+                                                    {{1, -1, 0}}) // x = y
+                         });
+  expectCoalesce(1, set);
+}
+
 } // namespace mlir

mlir/unittests/Analysis/Utils.h

+//===- Utils.h - utility functions for testing ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines several utility functions for creating
+// FlatAffineConstraints from sets of equalites and inequalities used in
+// testing several features of PresburgerSet, Simplex and AffineStructures.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_UNITTESTS_ANALYSIS_UTILS_H
+#define MLIR_UNITTESTS_ANALYSIS_UTILS_H
+
+#include "mlir/Analysis/AffineStructures.h"
+
+namespace mlir {
+
+/// Construct a FlatAffineConstraints from a set of inequality and
+/// equality constraints. `numIds` is the total number of ids, of which
+/// `numLocals` is the number of local ids.
+FlatAffineConstraints
+makeFACFromConstraints(unsigned numIds, ArrayRef<SmallVector<int64_t, 4>> ineqs,
+                       ArrayRef<SmallVector<int64_t, 4>> eqs,
+                       unsigned numLocals = 0, unsigned numSymbols = 0);
+
+/// Construct a FlatAffineConstraints having `numDims` dimensions from the given
+/// set of inequality constraints. This is a convenience function to be used
+/// when the FAC to be constructed does not have any local ids and does not have
+/// equalties.
+FlatAffineConstraints makeFACFromIneqs(unsigned dims,
+                                       ArrayRef<SmallVector<int64_t, 4>> ineqs);
+
+} // namespace mlir
+
+#endif // MLIR_UNITTESTS_ANALYSIS_FACUTILS_H

mlir/unittests/Analysis/Utils.cpp

+//===- Utils.cpp - Testing Utility Functions Implementation ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "./Utils.h"
+#include "mlir/Analysis/AffineStructures.h"
+
+using namespace mlir;
+
+FlatAffineConstraints mlir::makeFACFromConstraints(
+    unsigned numIds, ArrayRef<SmallVector<int64_t, 4>> ineqs,
+    ArrayRef<SmallVector<int64_t, 4>> eqs, unsigned numLocals /* = 0*/,
+    unsigned numSymbols /* = 0*/) {
+  FlatAffineConstraints fac(/*numReservedInequalities=*/ineqs.size(),
+                            /*numReservedEqualities=*/eqs.size(),
+                            /*numReservedCols=*/numIds + 1,
+                            /*numDims=*/numIds - numLocals - numSymbols,
+                            /*numSymbols=*/numSymbols,
+                            /*numLocals=*/numLocals);
+  for (const SmallVector<int64_t, 4> &eq : eqs)
+    fac.addEquality(eq);
+  for (const SmallVector<int64_t, 4> &ineq : ineqs)
+    fac.addInequality(ineq);
+  return fac;
+}
+
+FlatAffineConstraints
+mlir::makeFACFromIneqs(unsigned dims, ArrayRef<SmallVector<int64_t, 4>> ineqs) {
+  return makeFACFromConstraints(dims, ineqs, {});
+}