[Mlir-commits] [mlir] 8663915 - [MLIR][Presburger] Implement getDomainSet and getRangeSet for PresburgerRelation

[llvmlistbot at llvm.org]

Author: iambrj
Date: 2023-08-19T16:12:43+05:30
New Revision: 866391580b90531896dff7ffc28ecd961cb18084

URL: https://github.com/llvm/llvm-project/commit/866391580b90531896dff7ffc28ecd961cb18084
DIFF: https://github.com/llvm/llvm-project/commit/866391580b90531896dff7ffc28ecd961cb18084.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
    mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
    mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h b/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
index 28af83bdd1316d..f54878272d86d0 100644
--- a/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
+++ b/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

diff  --git a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
index abe9651019007d..634047d43464ee 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
+++ b/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();

diff  --git a/mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp
index 012a9bd5d084c1..edb41050f9dfa6 100644
--- a/mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp
+++ b/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));
+}