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

[Groverkss]

Summary

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

[ftynse]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2119

Not blocking, but I would be interested to see a more efficient implementation. Currently, it will first copy the constraints from poly to rel, then resize rel to accommodate range variables that likely leads to copies, then resize it *again* to accommodate local ids, and then copy the constraints from rel to this. Sounds like a whole lot of copying for a relatively simple operation. Perhaps we could get this done with only one resize (of the constraints of this relation, which increases both the number of rows and the number of ranges) and a direct copy of coefficients from poly to this. It may be less concise and use non-trivial indexing, but likely more efficient.

[Groverkss]

Rebase

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2119

The reason we need to do the space alignment is for the call to mergeLocalIds. Without merging the locals, we cannot directly append constraints, so I don't think this will work.

[ftynse]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2119

That's why I mentioned the increase of both the rows and columns in one sweep. I had drilled down a bit in the code and had seen that mereLocalIds does additional column manipulations itself (with even more copies!), which led me to conclude that decreasing the number of copies required non-trivial work and it was not reasonable to block this patch until that work is done. I am quite positive that what I have in mind will ultimately work, it's just harder to implement; if one transitively inlined mergeLocalIds into this call, it would become more obvious.

On a general design point, I think more care should be taken, eventually, about the memory effects of copies hidden under the layers of API. Copies are not free (although small copies may be cheap) and may have larger effects on memory behavior of the objects copied take, e.g., a little bit more than half of L1. It is not clear to me how much this affects performance of some compilation flow, largely because I don't have access to an end-to-end flow that extensively uses affine, especially compared to algorithmically-expensive things like the simplex, but I take small gains whenever possible. (One could imagine some sort of transactional protocol for updating the contraint matrices that has an indirection table for column accesses and doesn't actually remove/reshuffle the columns until the transaction is committed. Depending on the amount of data copied vs. the cost of extra indirection in address computation, this may end up beneficial.)

There is another small trade-off here: increase the API surface by having overloads that take rvalue-references but let the user spare the copy if they know they don't need the object anymore after the operation.

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2119

Right, I see what you mean.

We have some more work in the pipeline to improve the interface and de-virtualize the hierarchy. After that, we can benchmark, try out the transaction idea, try to remove copies as much as possible and see how it impacts performance.

The rvalue-reference idea also seems interesting and I think we can try to prototype that.

[arjunp]

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

508–511

Can you explicitly write that these mutates this, and also write a formal description? i.e., if this is R1: A -> B and poly is C, return R2: (A intersection C) -> B such that ...

508–520

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

2111–2112

use inverse here?

2114

2128

2131–2132

Not for this patch, but I would like an intersect function to mergeLocalIds and append.

[Groverkss]

Address Arjun's comments

[arjunp]

In the operations that were added previously, we've been following a system where most operations return their results, with operations that mutate in-place having a suffix InPlace, like unionInPlace. Not for this patch, but we should have uniformity in naming here.

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

public:
   }
 
   /// Return a set corresponding to all points in the domain of the relation.
   IntegerPolyhedron getDomainSet() const;
 
   /// 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.
   void inverse();
 
   void print(raw_ostream &os) const;
   void dump() const;
 

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

copyRel.convertIdKind(IdKind::Domain, 0, getNumIdKind(IdKind::Domain),
                         IdKind::Local);
 
   // We do not need to do anything to Range variables since they are already in
   // SetDim position.
 
   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);
   convertIdKind(IdKind::Range, 0, numRangeIds, IdKind::Domain);
 }
 
 void IntegerRelation::printSpace(raw_ostream &os) const {
   space.print(os);

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

IntegerRelation inverseRel =
       parseRelationFromSet("(z, x, y)[N, M]  : (x >= 0, N - x >= 0, y >= 0, M "
                            "- y >= 0, x + y - z == 0)",
                            1);
 
   rel.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));
+  }
+}