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

This patch implements getDomainSet and getRangeSet for PresburgerRelation

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D158263

Author: iambrj

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

@@ -97,6 +97,11 @@ class PresburgerRelation {
   /// operation returns (A intersection C) -> B.
   PresburgerRelation intersectDomain(const PresburgerSet &set) const;
 
+  /// Return a set corresponding to the domain of the relation.
+  PresburgerSet getDomainSet() const;
+  /// Return a set corresponding to 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

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

@@ -164,6 +164,20 @@ PresburgerRelation::intersectDomain(const PresburgerSet &set) const {
   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();

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

@@ -191,7 +191,7 @@ TEST(PresburgerRelationTest, inverse) {
   }
 }
 
-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)",
@@ -296,3 +296,29 @@ TEST(IntegerRelationTest, symbolicLexOpt) {
   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));
+}