[MLIR] Support checking if two FlatAffineConstraints are equal

[arjunp]

Summary

This patch adds support for checking if two PresburgerSets are equal. In particular, one can check if two FlatAffineConstraints are equal by constructing PrebsurgerSets from them and comparing these.

[ftynse]

mlir/lib/Analysis/PresburgerSet.cpp

297–299

return setMinusThis.isInegerEmpty();

[arjunp]

Addressed Alex's comment.

mlir/include/mlir/Analysis/PresburgerSet.h

public:
   void unionSetInPlace(const PresburgerSet &set);
 
   /// Return the union of this set and the given set.
   PresburgerSet unionSet(const PresburgerSet &set) const;
 
   /// Return the intersection of this set and the given set.
   PresburgerSet intersect(const PresburgerSet &set) const;
 
-  /// Return true if the set contains the given point, or false otherwise.
+  /// Return true if the set contains the given point, and false otherwise.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
   /// Print the set's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
   /// Return the complement of this set.
   PresburgerSet complement() const;
 
   /// Return the set difference of this set and the given set, i.e.,
   /// return `this \ set`.
   PresburgerSet subtract(const PresburgerSet &set) const;
 
+  /// Return true if this set is equal to the given set, and false otherwise.
+  bool isEqual(const PresburgerSet &set) const;
+
   /// Return a universe set of the specified type that contains all points.
   static PresburgerSet getUniverse(unsigned nDim = 0, unsigned nSym = 0);
   /// Return an empty set of the specified type that contains no points.
   static PresburgerSet getEmptySet(unsigned nDim = 0, unsigned nSym = 0);
 
   /// Return true if all the sets in the union are known to be integer empty
   /// false otherwise.
   bool isIntegerEmpty() const;

mlir/lib/Analysis/PresburgerSet.cpp

PresburgerSet PresburgerSet::subtract(const PresburgerSet &set) const {
   assertDimensionsCompatible(set, *this);
   PresburgerSet result(nDim, nSym);
   // We compute (V_i t_i) \ (V_i set_i) as V_i (t_i \ V_i set_i).
   for (const FlatAffineConstraints &fac : flatAffineConstraints)
     result.unionSetInPlace(getSetDifference(fac, set));
   return result;
 }
 
+/// Two sets S and T are equal iff S contains T and T contains S.
+/// By "S contains T", we mean that S is a superset of or equal to T.
+///
+/// S contains T iff T \ S is empty, since if T \ S contains a
+/// point then this is a point that is contained in T but not S.
+///
+/// Therefore, S is equal to T iff S \ T and T \ S are both empty.
+bool PresburgerSet::isEqual(const PresburgerSet &set) const {
+  assertDimensionsCompatible(set, *this);
+  PresburgerSet thisMinusSet = subtract(set);
+  if (!thisMinusSet.isIntegerEmpty())
+    return false;
+  PresburgerSet setMinusThis = set.subtract(*this);
+  return setMinusThis.isIntegerEmpty();
+}
+
 /// Return true if all the sets in the union are known to be integer empty,
 /// false otherwise.
 bool PresburgerSet::isIntegerEmpty() const {
   assert(nSym == 0 && "isIntegerEmpty is intended for non-symbolic sets");
   // The set is empty iff all of the disjuncts are empty.
   for (const FlatAffineConstraints &fac : flatAffineConstraints) {
     if (!fac.isIntegerEmpty())
       return false;

mlir/unittests/Analysis/PresburgerSetTest.cpp

 //===- SetTest.cpp - Tests for PresburgerSet ------------------------------===//
 //
 // 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 file contains tests for PresburgerSet. Each test works by computing
-// an operation (union, intersection, subtract, or complement) 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.
+// This file contains tests for PresburgerSet. The tests for union,
+// 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 <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.
 static void testUnionAtPoints(PresburgerSet s, PresburgerSet t,
                               ArrayRef<SmallVector<int64_t, 4>> points) {
   PresburgerSet unionSet = s.unionSet(t);
   for (const SmallVector<int64_t, 4> &point : points) {
     bool inS = s.containsPoint(point);
     bool inT = t.containsPoint(point);
     bool inUnion = unionSet.containsPoint(point);
     EXPECT_EQ(inUnion, inS || inT);
   }
 }
 
 /// Compute the intersection of s and t, and check that each of the given points
-/// belongs to the intersection iff it belongs to both of s and t.
+/// belongs to the intersection iff it belongs to both s and t.
 static void testIntersectAtPoints(PresburgerSet s, PresburgerSet t,
                                   ArrayRef<SmallVector<int64_t, 4>> points) {
   PresburgerSet intersection = s.intersect(t);
   for (const SmallVector<int64_t, 4> &point : points) {
     bool inS = s.containsPoint(point);
     bool inT = t.containsPoint(point);
     bool inIntersection = intersection.containsPoint(point);
     EXPECT_EQ(inIntersection, inS && inT);
testComplementAtPoints(
        {10, 10},
        {10, 11},
        {11, 10},
        {2, 10},
        {2, 11},
        {1, 10}});
 }
 
+TEST(SetTest, isEqual) {
+  // set = [2, 8] U [10, 20].
+  PresburgerSet universe = PresburgerSet::getUniverse(1);
+  PresburgerSet emptySet = PresburgerSet::getEmptySet(1);
+  PresburgerSet set =
+      makeSetFromFACs(1, {
+                             makeFACFromIneqs(1, {{1, -2},    // x >= 2.
+                                                  {-1, 8}}),  // x <= 8.
+                             makeFACFromIneqs(1, {{1, -10},   // x >= 10.
+                                                  {-1, 20}}), // x <= 20.
+                         });
+
+  // universe != emptySet.
+  EXPECT_FALSE(universe.isEqual(emptySet));
+  // emptySet != universe.
+  EXPECT_FALSE(emptySet.isEqual(universe));
+  // emptySet == emptySet.
+  EXPECT_TRUE(emptySet.isEqual(emptySet));
+  // universe == universe.
+  EXPECT_TRUE(universe.isEqual(universe));
+
+  // universe U emptySet == universe.
+  EXPECT_TRUE(universe.unionSet(emptySet).isEqual(universe));
+  // universe U universe == universe.
+  EXPECT_TRUE(universe.unionSet(universe).isEqual(universe));
+  // emptySet U emptySet == emptySet.
+  EXPECT_TRUE(emptySet.unionSet(emptySet).isEqual(emptySet));
+  // universe U emptySet != emptySet.
+  EXPECT_FALSE(universe.unionSet(emptySet).isEqual(emptySet));
+  // universe U universe != emptySet.
+  EXPECT_FALSE(universe.unionSet(universe).isEqual(emptySet));
+  // emptySet U emptySet != universe.
+  EXPECT_FALSE(emptySet.unionSet(emptySet).isEqual(universe));
+
+  // set \ set == emptySet.
+  EXPECT_TRUE(set.subtract(set).isEqual(emptySet));
+  // set == set.
+  EXPECT_TRUE(set.isEqual(set));
+  // set U (universe \ set) == universe.
+  EXPECT_TRUE(set.unionSet(set.complement()).isEqual(universe));
+  // set U (universe \ set) != set.
+  EXPECT_FALSE(set.unionSet(set.complement()).isEqual(set));
+  // set != set U (universe \ set).
+  EXPECT_FALSE(set.isEqual(set.unionSet(set.complement())));
+
+  // square is one unit taller than rect.
+  PresburgerSet square =
+      makeSetFromFACs(2, {makeFACFromIneqs(2, {
+                                                  {1, 0, -2}, // x >= 2.
+                                                  {0, 1, -2}, // y >= 2.
+                                                  {-1, 0, 9}, // x <= 9.
+                                                  {0, -1, 9}  // y <= 9.
+                                              })});
+  PresburgerSet rect =
+      makeSetFromFACs(2, {makeFACFromIneqs(2, {
+                                                  {1, 0, -2}, // x >= 2.
+                                                  {0, 1, -2}, // y >= 2.
+                                                  {-1, 0, 9}, // x <= 9.
+                                                  {0, -1, 8}  // y <= 8.
+                                              })});
+  EXPECT_FALSE(square.isEqual(rect));
+  PresburgerSet universeRect = square.unionSet(square.complement());
+  PresburgerSet universeSquare = rect.unionSet(rect.complement());
+  EXPECT_TRUE(universeRect.isEqual(universeSquare));
+  EXPECT_FALSE(universeRect.isEqual(rect));
+  EXPECT_FALSE(universeSquare.isEqual(square));
+  EXPECT_FALSE(rect.complement().isEqual(square.complement()));
+}
+
 } // namespace mlir