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

[Groverkss]

Summary

This patch adds support for applying a relation on domain/range of a relation.

[arjunp]

I think we should add a compose member function, it being a natural opreation on relations. You can move applyRange's documentation to compose and just mention in applyRange that it's the same as compose, but provided for uniformity with applyDomain.

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

534–536

It might be simpler to just say that this returns R2.inverse().compose(R1), and then perhaps provide a concrete example.

541–543

You can just write that it returns R1.compose(R2) and mention that this is provided for uniformity with applyDomain.

[Groverkss]

• Rebase
• Address Arjun's comments

[Groverkss]

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

534–536

I feel the R2.invese().compose() makes it harder to understand than providing a proper mathematical description.

[arjunp]

I think you can remove the backslashes on exists and in, unless doxygen can somehow render them, it seems unnecessary.

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

534–536

R2.inverse().compose() is in fact a proper mathematical description. For me personally, a concrete example was more useful to understand what the point of this is, and then for a formal specification the inverse compose thing will show what the general version of that example is.

[Groverkss]

Address Arjun's comments

[arjunp]

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

533

What do you mean by R1(R2) here? If R3 is A -> C then R3(x) is R2(R1(x)). It seems the standard notation here is R1;R2 (https://en.wikipedia.org/wiki/Composition_of_relations).

535

536

537

542

543

read the original as (i->j), k at first

545

I would prefer this example, it has less "noise" and also shows a little bit happens when the other variables depend on the projected variable

R1 : i -> j : (0 <= i < 2, j = i)
R2: i -> k : (k = i floordiv 2)

R3: k -> j : (0 <= k < 1, 2k <=  j <= 2k + 1)

R1 = {(0, 0), (1, 1)}. R2 maps both 0 and 1 to 0. So R3 = {(0, 0), (0, 1)}

550–557

Not sure if an example is needed for composition but you can do e.g. R2 : j -> k : k = floordiv 2 and R1 from above

552

557

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2148–2156

I think the current implementation should be thought of as converting R1 to A -> (B x C), converting R2 to B x C, and intersecting

[Groverkss]

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

533

I meant it like we represent function composition f(g(x)). But this also works for me.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2148–2156

good catch!

[Groverkss]

Address Arjun's comments

[arjunp]

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

533

Right, I mean that we are doing R2(R1(x)) then, not R1(R2).

I guess you can also add that it R3)(x) = R2(R1(x)), might be more understandable

546

548

[Groverkss]

Address Arjun's comments and rebase

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

public:
   /// to perform much better in the average case. If V is the number of vertices
   /// in the polytope and C is the number of constraints, the algorithm takes
   /// O(VC) time.
   void removeRedundantConstraints();
 
   void removeDuplicateDivs();
 
   /// Converts identifiers of kind srcKind in the range [idStart, idLimit) to
-  /// variables of kind dstKind and placed after all the other variables of kind
-  /// dstKind. The internal ordering among the moved variables is preserved.
+  /// variables of kind dstKind. If `pos` is given, the variables are placed at
+  /// position `pos` of dstKind, otherwise they are placed after all the other
+  /// variables of kind dstKind. The internal ordering among the moved variables
+  /// is preserved.
   void convertIdKind(IdKind srcKind, unsigned idStart, unsigned idLimit,
-                     IdKind dstKind);
+                     IdKind dstKind, unsigned pos);
+  void convertIdKind(IdKind srcKind, unsigned idStart, unsigned idLimit,
+                     IdKind dstKind) {
+    convertIdKind(srcKind, idStart, idLimit, dstKind, getNumIdKind(dstKind));
+  }
   void convertToLocal(IdKind kind, unsigned idStart, unsigned idLimit) {
     convertIdKind(kind, idStart, idLimit, IdKind::Local);
   }
 
   /// Adds additional local ids to the sets such that they both have the union
   /// of the local ids in each set, without changing the set of points that
   /// lie in `this` and `other`.
   ///
public:
   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.
   void inverse();
 
+  /// Let the relation `this` be R1, and the relation `rel` be R2. Modifies R1
+  /// to be the composition of R1 and R2: R1;R2.
+  ///
+  /// Formally, if R1: A -> B, and R2: B -> C, then this function returns a
+  /// relation R3: A -> C such that a point (a, c) belongs to R3 iff there
+  /// exists b such that (a, b) is in R1 and, (b, c) is in R2.
+  void compose(const IntegerRelation &rel);
+
+  /// Given a relation `rel`, apply the relation to the domain of this relation.
+  ///
+  /// R1: i -> j : (0 <= i < 2, j = i)
+  /// R2: i -> k : (k = i floordiv 2)
+  /// R3: k -> j : (0 <= k < 1, 2k <=  j <= 2k + 1)
+  ///
+  /// R1 = {(0, 0), (1, 1)}. R2 maps both 0 and 1 to 0.
+  /// So R3 = {(0, 0), (0, 1)}.
+  ///
+  /// Formally, R1.applyDomain(R2) = R2.inverse().compose(R1).
+  void applyDomain(const IntegerRelation &rel);
+
+  /// Given a relation `rel`, apply the relation to the range of this relation.
+  ///
+  /// Formally, R1.applyRange(R2) is the same as R1.compose(R2) but we provide
+  /// this for uniformity with `applyDomain`.
+  void applyRange(const IntegerRelation &rel);
+
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Checks all rows of equality/inequality constraints for trivial
   /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced
   /// after elimination. Returns true if an invalid constraint is found;
   /// false otherwise.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

while (true) {
 
     // Remove the ith equality and the found local variable.
     removeId(j);
     removeEquality(i);
   }
 }
 
 void IntegerRelation::convertIdKind(IdKind srcKind, unsigned idStart,
-                                    unsigned idLimit, IdKind dstKind) {
+                                    unsigned idLimit, IdKind dstKind,
+                                    unsigned pos) {
   assert(idLimit <= getNumIdKind(srcKind) && "Invalid id range");
 
   if (idStart >= idLimit)
     return;
 
   // Append new local variables corresponding to the dimensions to be converted.
-  unsigned newIdsBegin = getIdKindEnd(dstKind);
   unsigned convertCount = idLimit - idStart;
-  appendId(dstKind, convertCount);
+  unsigned newIdsBegin = insertId(dstKind, pos, convertCount);
 
   // Swap the new local variables with dimensions.
   //
   // Essentially, this moves the information corresponding to the specified ids
   // of kind `srcKind` to the `convertCount` newly created ids of kind
   // `dstKind`. In particular, this moves the columns in the constraint
   // matrices, and zeros out the initially occupied columns (because the newly
   // created ids we're swapping with were zero-initialized).

 }
 
 void IntegerRelation::inverse() {
   unsigned numRangeIds = getNumIdKind(IdKind::Range);
   convertIdKind(IdKind::Domain, 0, getIdKindEnd(IdKind::Domain), IdKind::Range);
   convertIdKind(IdKind::Range, 0, numRangeIds, IdKind::Domain);
 }
 
+void IntegerRelation::compose(const IntegerRelation &rel) {
+  assert(getRangeSet().getSpace().isCompatible(rel.getDomainSet().getSpace()) &&
+         "Range of `this` should be compatible with Domain of `rel`");
+
+  IntegerRelation copyRel = rel;
+
+  // Let relation `this` be R1: A -> B, and `rel` be R2: B -> C.
+  // We convert R1 to A -> (B X C), and R2 to B X C then intersect the range of
+  // R1 with R2. After this, we get R1: A -> C, by projecting out B.
+  // TODO: Using nested spaces here would help, since we could directly
+  // intersect the range with another relation.
+  unsigned numBIds = getNumRangeIds();
+
+  // Convert R1 from A -> B to A -> (B X C).
+  appendId(IdKind::Range, copyRel.getNumRangeIds());
+
+  // Convert R2 to B X C.
+  copyRel.convertIdKind(IdKind::Domain, 0, numBIds, IdKind::Range, 0);
+
+  // Intersect R2 to range of R1.
+  intersectRange(IntegerPolyhedron(copyRel));
+
+  // Project out B in R1.
+  convertIdKind(IdKind::Range, 0, numBIds, IdKind::Local);
+}
+
+void IntegerRelation::applyDomain(const IntegerRelation &rel) {
+  inverse();
+  compose(rel);
+  inverse();
+}
+
+void IntegerRelation::applyRange(const IntegerRelation &rel) { compose(rel); }
+
 void IntegerRelation::printSpace(raw_ostream &os) const {
   space.print(os);
   os << getNumConstraints() << " constraints\n";
 }
 
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

IntegerRelation expectedRel = parseRelationFromSet(
         ">= 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));
   }
 }
+
+TEST(IntegerRelationTest, applyDomainAndRange) {
+
+  {
+    IntegerRelation map1 = parseRelationFromSet(
+        "(x, y, a, b)[N] : (a - x - N == 0, b - y + N == 0)", 2);
+    IntegerRelation map2 =
+        parseRelationFromSet("(x, y, a)[N] : (a - x - y == 0)", 2);
+
+    map1.applyRange(map2);
+
+    IntegerRelation map3 =
+        parseRelationFromSet("(x, y, a)[N] : (a - x - y == 0)", 2);
+
+    EXPECT_TRUE(map1.isEqual(map3));
+  }
+
+  {
+    IntegerRelation map1 = parseRelationFromSet(
+        "(x, y, a, b)[N] : (a - x + N == 0, b - y - N == 0)", 2);
+    IntegerRelation map2 =
+        parseRelationFromSet("(x, y, a, b)[N] : (a - N == 0, b - N == 0)", 2);
+
+    IntegerRelation map3 =
+        parseRelationFromSet("(x, y, a, b)[N] : (x - N == 0, y - N == 0)", 2);
+
+    map1.applyDomain(map2);
+
+    EXPECT_TRUE(map1.isEqual(map3));
+  }
+}