[MLIR][Presburger] Add intersectDomain/Range to IntegerRelation

This patch adds support for intersection a set with a relation.

Reviewed By: arjunp

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

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

@@ -505,7 +505,19 @@ class IntegerRelation {
   /// Return a set corresponding to all points in the range of the relation.
   IntegerPolyhedron getRangeSet() const;
 
-  /// Invert the relation i.e., swap it's domain and range.
+  /// Intersect the given `poly` with the domain in-place.
+  ///
+  /// Formally, let the relation `this` be R: A -> B and poly is C, then this
+  /// operation modifies R to be (A intersection C) -> B.
+  void intersectDomain(const IntegerPolyhedron &poly);
+
+  /// Intersect the given `poly` with the range in-place.
+  ///
+  /// Formally, let the relation `this` be R: A -> B and poly is C, then this
+  /// operation modifies R to be A -> (B intersection C).
+  void intersectRange(const IntegerPolyhedron &poly);
+
+  /// Invert the relation i.e., swap its domain and range.
   ///
   /// Formally, let the relation `this` be R: A -> B, then this operation
   /// modifies R to be B -> A.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

@@ -2102,6 +2102,35 @@ IntegerPolyhedron IntegerRelation::getRangeSet() const {
   return IntegerPolyhedron(std::move(copyRel));
 }
 
+void IntegerRelation::intersectDomain(const IntegerPolyhedron &poly) {
+  assert(getDomainSet().getSpace().isCompatible(poly.getSpace()) &&
+         "Domain set is not compatible with poly");
+
+  // Treating the poly as a relation, convert it from `0 -> R` to `R -> 0`.
+  IntegerRelation rel = poly;
+  rel.inverse();
+
+  // Append dummy range variables to make the spaces compatible.
+  rel.appendId(IdKind::Range, getNumRangeIds());
+
+  // Intersect in place.
+  mergeLocalIds(rel);
+  append(rel);
+}
+
+void IntegerRelation::intersectRange(const IntegerPolyhedron &poly) {
+  assert(getRangeSet().getSpace().isCompatible(poly.getSpace()) &&
+         "Range set is not compatible with poly");
+
+  IntegerRelation rel = poly;
+
+  // Append dummy domain variables to make the spaces compatible.
+  rel.appendId(IdKind::Domain, getNumDomainIds());
+
+  mergeLocalIds(rel);
+  append(rel);
+}
+
 void IntegerRelation::inverse() {
   unsigned numRangeIds = getNumIdKind(IdKind::Range);
   convertIdKind(IdKind::Domain, 0, getIdKindEnd(IdKind::Domain), IdKind::Range);

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

@@ -57,3 +57,37 @@ TEST(IntegerRelationTest, inverse) {
 
   EXPECT_TRUE(rel.isEqual(inverseRel));
 }
+
+TEST(IntegerRelationTest, intersectDomainAndRange) {
+  IntegerRelation rel = parseRelationFromSet(
+      "(x, y, z)[N, M]: (y floordiv 2 - N >= 0, z floordiv 5 - M"
+      ">= 0, x + y + z floordiv 7 == 0)",
+      1);
+
+  {
+    IntegerPolyhedron poly = parsePoly("(x)[N, M] : (x >= 0, M - x - 1 >= 0)");
+
+    IntegerRelation expectedRel = parseRelationFromSet(
+        "(x, y, z)[N, M]: (y floordiv 2 - N >= 0, z floordiv 5 - M"
+        ">= 0, x + y + z floordiv 7 == 0, x >= 0, M - x - 1 >= 0)",
+        1);
+
+    IntegerRelation copyRel = rel;
+    copyRel.intersectDomain(poly);
+    EXPECT_TRUE(copyRel.isEqual(expectedRel));
+  }
+
+  {
+    IntegerPolyhedron poly =
+        parsePoly("(y, z)[N, M] : (y >= 0, M - y - 1 >= 0, y + z == 0)");
+
+    IntegerRelation expectedRel = parseRelationFromSet(
+        "(x, y, z)[N, M]: (y floordiv 2 - N >= 0, z floordiv 5 - M"
+        ">= 0, x + y + z floordiv 7 == 0, y >= 0, M - y - 1 >= 0, y + z == 0)",
+        1);
+
+    IntegerRelation copyRel = rel;
+    copyRel.intersectRange(poly);
+    EXPECT_TRUE(copyRel.isEqual(expectedRel));
+  }
+}