[MLIR][Presburger] Implement domain and range restriction for PresburgerRelation

[iambrj]

Summary

This patch implements domain and range restriction for PresburgerRelation

[iambrj]

Fix grammar

[iambrj]

Grammar again

[Groverkss]

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

88–103

These methods should return PresburgerRelation. Methods in PresburgerRelation return a new object unless they have a "inPlace" suffix.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

127–157

I would try to reuse PresburgerRelation::intersect as much as possible. We are implementing simplifications in core methods like intersect and it makes it harder to keep track of these things otherwise.

[iambrj]

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

88–103

Right. What is the rule of thumb for suffixing "inPlace"? For example, inverse, compose, and applyDomain do not have "inPlace" suffixed, even though they mutate the object and do not return a new object.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

127–157

The way I can see to implement intersectRange with intersect is to convert the passed PresburgerSet into a PresburgerRelation, insert domain variables to the converted PresburgerRelation to make it intersectable with this, and then call intersect. Would the above be a good way to do it? Do you have any suggestions on other ways to do it?

[Groverkss]

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

88–103

I think we try to do our best in suffixing things. For older methods, we might have to fix this naming later. But for now, since intersect returns a PresburgerRelation, it's best if intersectDomain/Range also does the same.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

127–157

I think that sounds like a good way to me.

[iambrj]

Address comments

[iambrj]

Remove unused include

[Groverkss]

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

88

91–93

I think using the same set is confusing. I prefer the same documentation as IntegerRelation::intersectRange

/// Formally, let the relation `this` be R: A -> B and poly is C, then this
/// operation modifies R to be A -> (B intersection C).

96

99–101

Similarly for this:

/// Formally, let the relation `this` be R: A -> B and poly is C, then this
/// operation modifies R to be (A intersection C) -> B.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

35–37

Braces can be removed.

[iambrj]

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

91–93

Sure, do you suggest the passed PresburgerSet's name be changed to poly as well?

[Groverkss]

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

91–93

You can just change "poly" to "set" in the documentation. But if you prefer it the other way, you can do that too. As long as the documentation makes sense, it's fine.

[iambrj]

Address comments

[Groverkss]

LGTM

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

public:
   /// Set the space to `oSpace`. `oSpace` should not contain any local ids.
   /// `oSpace` need not have the same number of ids as the current space;
   /// it could have more or less. If it has less, the extra ids become
   /// locals of the disjuncts. It can also have more, in which case the
   /// disjuncts will have fewer locals. If its total number of ids
   /// exceeds that of some disjunct, an assert failure will occur.
   void setSpace(const PresburgerSpace &oSpace);
 
+  void insertVarInPlace(VarKind kind, unsigned pos, unsigned num = 1);
+
   /// Return a reference to the list of disjuncts.
   ArrayRef<IntegerRelation> getAllDisjuncts() const;
 
   /// Return the disjunct at the specified index.
   const IntegerRelation &getDisjunct(unsigned index) const;
 
   /// Mutate this set, turning it into the union of this set and the given
   /// disjunct.
   void unionInPlace(const IntegerRelation &disjunct);
 
   /// Mutate this set, turning it into the union of this set and the given set.
   void unionInPlace(const PresburgerRelation &set);
 
   /// Return the union of this set and the given set.
   PresburgerRelation unionSet(const PresburgerRelation &set) const;
 
   /// Return the intersection of this set and the given set.
   PresburgerRelation intersect(const PresburgerRelation &set) const;
 
+  /// Intersect the given `set` with the range in-place.
+  ///
+  /// Formally, let the relation `this` be R: A -> B and `set` is C, then this
+  /// operation modifies R to be A -> (B intersection C).
+  PresburgerRelation intersectRange(PresburgerSet &set);
+
+  /// Intersect the given `set` with the domain in-place.
+  ///
+  /// Formally, let the relation `this` be R: A -> B and `set` is C, then this
+  /// operation modifies R to be (A intersection C) -> B.
+  PresburgerRelation intersectDomain(const PresburgerSet &set);
+
   /// 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.
   ///
   /// Formally, if `this`: A -> B, and `rel`: B -> C, then this function updates
   /// `this` to `result`: A -> C where a point (a, c) belongs to `result`
   /// iff there exists b such that (a, b) is in `this` and, (b, c) is in rel.
   void compose(const PresburgerRelation &rel);

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 
 void PresburgerRelation::setSpace(const PresburgerSpace &oSpace) {
   assert(space.getNumLocalVars() == 0 && "no locals should be present");
   space = oSpace;
   for (IntegerRelation &disjunct : disjuncts)
     disjunct.setSpaceExceptLocals(space);
 }
 
+void PresburgerRelation::insertVarInPlace(VarKind kind, unsigned pos,
+                                          unsigned num) {
+  for (IntegerRelation &cs : disjuncts)
+    cs.insertVar(kind, pos, num);
+  space.insertVar(kind, pos, num);
+}
+
 unsigned PresburgerRelation::getNumDisjuncts() const {
   return disjuncts.size();
 }
 
 ArrayRef<IntegerRelation> PresburgerRelation::getAllDisjuncts() const {
   return disjuncts;
 }
 
for (const IntegerRelation &csB : set.disjuncts) {
       IntegerRelation intersection = csA.intersect(csB);
       if (!intersection.isEmpty())
         result.unionInPlace(intersection);
     }
   }
   return result;
 }
 
+PresburgerRelation PresburgerRelation::intersectRange(PresburgerSet &set) {
+  assert(space.getRangeSpace().isCompatible(set.getSpace()) &&
+         "Range of `this` must be compatible with range of `set`");
+
+  PresburgerRelation other = set;
+  other.insertVarInPlace(VarKind::Domain, 0, getNumDomainVars());
+  return intersect(other);
+}
+
+PresburgerRelation
+PresburgerRelation::intersectDomain(const PresburgerSet &set) {
+  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);
+}
+
 void PresburgerRelation::inverse() {
   for (IntegerRelation &cs : disjuncts)
     cs.inverse();
 
   if (getNumDisjuncts())
     setSpace(getDisjunct(0).getSpaceWithoutLocals());
 }
 
 void PresburgerRelation::compose(const PresburgerRelation &rel) {
   assert(getSpace().getRangeSpace().isCompatible(
              rel.getSpace().getDomainSpace()) &&
          "Range of `this` should be compatible with domain of `rel`");
 
   PresburgerRelation result =
       PresburgerRelation::getEmpty(PresburgerSpace::getRelationSpace(
           getNumDomainVars(), rel.getNumRangeVars(), getNumSymbolVars()));
   for (const IntegerRelation &csA : disjuncts) {
     for (const IntegerRelation &csB : rel.disjuncts) {
       IntegerRelation composition = csA;

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

parsePresburgerRelationFromPresburgerSet(ArrayRef<StringRef> strs,
   for (unsigned i = 1, e = strs.size(); i < e; ++i) {
     rel = parseIntegerPolyhedron(strs[i]);
     rel.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain);
     result.unionInPlace(rel);
   }
   return result;
 }
 
+TEST(PresburgerRelationTest, intersectDomainAndRange) {
+  PresburgerRelation rel = parsePresburgerRelationFromPresburgerSet(
+      {// (x, y) -> (x + N, y - N)
+       "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0)",
+       // (x, y) -> (x + y, x - y)
+       "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0)",
+       // (x, y) -> (x - y, y - x)}
+       "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0)"},
+      2);
+
+  {
+    PresburgerSet set =
+        parsePresburgerSet({// (2x, x)
+                            "(a, b)[N] : (a - 2 * b == 0)",
+                            // (x, -x)
+                            "(a, b)[N] : (a + b == 0)",
+                            // (N, N)
+                            "(a, b)[N] : (a - N == 0, b - N == 0)"});
+
+    PresburgerRelation expectedRel = parsePresburgerRelationFromPresburgerSet(
+        {"(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, x - 2 * y == 0)",
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, x + y == 0)",
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, x - N == 0, y - N "
+         "== 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, x - 2 * y == 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, x + y == 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, x - N == 0, y - N "
+         "== 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, x - 2 * y == 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, x + y == 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, x - N == 0, y - N "
+         "== 0)"},
+        2);
+
+    PresburgerRelation computedRel = rel.intersectDomain(set);
+    EXPECT_TRUE(computedRel.isEqual(expectedRel));
+  }
+
+  {
+    PresburgerSet set =
+        parsePresburgerSet({// (2x, x)
+                            "(a, b)[N] : (a - 2 * b == 0)",
+                            // (x, -x)
+                            "(a, b)[N] : (a + b == 0)",
+                            // (N, N)
+                            "(a, b)[N] : (a - N == 0, b - N == 0)"});
+
+    PresburgerRelation expectedRel = parsePresburgerRelationFromPresburgerSet(
+        {"(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, a - 2 * b == 0)",
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, a + b == 0)",
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0, a - N == 0, b - N "
+         "== 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, a - 2 * b == 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, a + b == 0)",
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0, a - N == 0, b - N "
+         "== 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, a - 2 * b == 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, a + b == 0)",
+         "(x, y, a, b)[N] : (a - x + y == 0, b - y + x == 0, a - N == 0, b - N "
+         "== 0)"},
+        2);
+
+    PresburgerRelation computedRel = rel.intersectRange(set);
+    EXPECT_TRUE(computedRel.isEqual(expectedRel));
+  }
+}
+
 TEST(PresburgerRelationTest, applyDomainAndRange) {
   {
     PresburgerRelation map1 = parsePresburgerRelationFromPresburgerSet(
         {// (x, y) -> (x + N, y - N)
          "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0)",
          // (x, y) -> (y, x)
          "(x, y, a, b)[N] : (a - y == 0, b - x == 0)",
          // (x, y) -> (x + y, x - y)