[MLIR][Presburger] Implement getDomainSet and getRangeSet for PresburgerRelation

[iambrj]

Summary

This patch implements getDomainSet and getRangeSet for PresburgerRelation

[Groverkss]

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

100

I think you can drop "all points". Just "Return a set corresponding to the domain of the relation" should be fine.

102

I generally remove newlines between related methods.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

169–171

You can drop braces here.

[iambrj]

Address comments

[iambrj]

Fix missing base commit

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

public:
   PresburgerRelation intersectRange(const PresburgerSet &set) const;
 
   /// Return the domain intersection of the given `set` with `this` relation.
   ///
   /// Formally, let the relation `this` be R: A -> B and `set` is C, then this
   /// operation returns (A intersection C) -> B.
   PresburgerRelation intersectDomain(const PresburgerSet &set) const;
 
+  /// Return a set corresponding to all points in the domain of the relation.
+  PresburgerSet getDomainSet() const;
+
+  /// Return a set corresponding to all points in the range of the relation.
+  PresburgerSet getRangeSet() const;
+
   /// Invert the relation, i.e. swap its domain and range.
   ///
   /// Formally, if `this`: A -> B then `inverse` updates `this` in-place to
   /// `this`: B -> A.
   void inverse();
 
   /// Compose `this` relation with the given relation `rel` in-place.
   ///

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

assert(space.getDomainSpace().isCompatible(set.getSpace()) &&
          "Domain of `this` must be compatible with range of `set`");
 
   PresburgerRelation other = set;
   other.insertVarInPlace(VarKind::Domain, 0, getNumDomainVars());
   other.inverse();
   return intersect(other);
 }
 
+PresburgerSet PresburgerRelation::getDomainSet() const {
+  PresburgerSet result = PresburgerSet::getEmpty(space.getDomainSpace());
+  for (const IntegerRelation &cs : disjuncts) {
+    result.unionInPlace(cs.getDomainSet());
+  }
+  return result;
+}
+
+PresburgerSet PresburgerRelation::getRangeSet() const {
+  PresburgerSet result = PresburgerSet::getEmpty(space.getRangeSpace());
+  for (const IntegerRelation &cs : disjuncts) {
+    result.unionInPlace(cs.getRangeSet());
+  }
+  return result;
+}
+
 void PresburgerRelation::inverse() {
   for (IntegerRelation &cs : disjuncts)
     cs.inverse();
 
   if (getNumDisjuncts())
     setSpace(getDisjunct(0).getSpaceWithoutLocals());
 }
 

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

PresburgerRelation inverseRel = parsePresburgerRelationFromPresburgerSet(
          // (x, y) -> (x - N, y + N)
          "(x, y, a, b)[N] : (x - N - a == 0, y + N - b == 0)"},
         2);
 
     EXPECT_TRUE(rel.isEqual(inverseRel));
   }
 }
 
-TEST(IntegerRelationTest, symbolicLexOpt) {
+TEST(PresburgerRelationTest, symbolicLexOpt) {
   PresburgerRelation rel1 = parsePresburgerRelationFromPresburgerSet(
       {"(x, y)[N, M] : (x >= 0, y >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1>= 0, "
        "2 * N - x >= 0, 2 * N - y >= 0)",
        "(x, y)[N, M] : (x >= 0, y >= 0, N - 1 >= 0, M >= 0, M - 2 * N - 1>= 0, "
        "x - N >= 0, M - x >= 0, y - 2 * N >= 0, M - y >= 0)"},
       1);
 
   SymbolicLexOpt lexmin1 = rel1.findSymbolicIntegerLexMin();
TEST(PresburgerRelationTest, symbolicLexOpt) {
   PWMAFunction expectedLexMax3 =
       parsePWMAF({{"(x) : (x >= 0, 1 - x >= 0)", "(x) -> (1, 1, x)"}});
 
   EXPECT_TRUE(lexmin3.unboundedDomain.isIntegerEmpty());
   EXPECT_TRUE(lexmin3.lexopt.isEqual(expectedLexMin3));
   EXPECT_TRUE(lexmax3.unboundedDomain.isIntegerEmpty());
   EXPECT_TRUE(lexmax3.lexopt.isEqual(expectedLexMax3));
 }
+
+TEST(PresburgerRelationTest, getDomainAndRangeSet) {
+  PresburgerRelation rel = parsePresburgerRelationFromPresburgerSet(
+      {// (x, y) -> (x + N, y - N)
+       "(x, y, a, b)[N] : (a >= 0, b >= 0, N - a >= 0, N - b >= 0, x - a + N "
+       "== 0, y - b - N == 0)",
+       // (x, y) -> (- y, - x)
+       "(x, y, a, b)[N] : (a >= 0, b >= 0, 2 * N - a >= 0, 2 * N - b >= 0, a + "
+       "y == 0, b + x == 0)"},
+      2);
+
+  PresburgerSet domainSet = rel.getDomainSet();
+
+  PresburgerSet expectedDomainSet = parsePresburgerSet(
+      {"(x, y)[N] : (x + N >= 0, -x >= 0, y - N >= 0, 2 * N - y >= 0)",
+       "(x, y)[N] : (x + 2 * N >= 0, -x >= 0, y + 2 * N >= 0, -y >= 0)"});
+
+  EXPECT_TRUE(domainSet.isEqual(expectedDomainSet));
+
+  PresburgerSet rangeSet = rel.getRangeSet();
+
+  PresburgerSet expectedRangeSet = parsePresburgerSet(
+      {"(x, y)[N] : (x >= 0, 2 * N - x >= 0, y >= 0, 2 * N - y >= 0)"});
+
+  EXPECT_TRUE(rangeSet.isEqual(expectedRangeSet));
+}