[MLIR][Presburger] Implement composition for PresburgerRelation

[iambrj]

Summary

This patch implements range and domain composition for PresburgerRelations

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

113–119

Why do we create a new relation here and not just do it in place?

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

39

nit: Prefer using -> instead of <->. <-> makes it look like the map is surjective, which may not be true.

[Groverkss]

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

99–101

Sentences in comments should use full stops at end.

[iambrj]

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

113–119

I considered that but missed that I can change the space using setSpace, so resorted to creating a new relation instead. Will rewrite.

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

39

I considered ->, but felt -> would be too close to function notation. I wanted to emphasize that these are relations, but I agree <-> makes it look like the map is surjective. Any suggestions for a more relational notation than -> ?

[iambrj]

Address comments

[Groverkss]

LGTM. Just some style comments.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

113–115

For a single statement, you can remove braces.

113–119

Ah. Ideally, we should have an inverse in PresburgerSpace as well, which would allow this. I can do that later though.

116–118

Same here, you can remove the braces since its a single statement.

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

39

Right. MLIR and Presburger already use -> syntax so for consistency sake, it's better to use -> maybe for now.

[iambrj]

Fix style mistakes

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

public:
   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;
 
+  /// 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);
+
+  /// Apply the domain of given relation `rel` to `this` relation.
+  ///
+  /// Formally, R1.applyDomain(R2) = R2.inverse().compose(R1).
+  void applyDomain(const PresburgerRelation &rel);
+
+  /// Same as compose, provided for uniformity with applyDomain.
+  void applyRange(const PresburgerRelation &rel);
+
   /// Return true if the set contains the given point, and false otherwise.
   bool containsPoint(ArrayRef<MPInt> point) const;
   bool containsPoint(ArrayRef<int64_t> point) const {
     return containsPoint(getMPIntVec(point));
   }
 
   /// Return the complement of this set. All local variables in the set must
   /// correspond to floor divisions.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 //===- PresburgerRelation.cpp - MLIR PresburgerRelation Class -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include <optional>
 
 using namespace mlir;
for (const IntegerRelation &csB : set.disjuncts) {
       IntegerRelation intersection = csA.intersect(csB);
       if (!intersection.isEmpty())
         result.unionInPlace(intersection);
     }
   }
   return result;
 }
 
+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;
+      composition.compose(csB);
+      if (!composition.isEmpty())
+        result.unionInPlace(composition);
+    }
+  }
+  *this = result;
+}
+
+void PresburgerRelation::applyDomain(const PresburgerRelation &rel) {
+  assert(getSpace().getDomainSpace().isCompatible(
+             rel.getSpace().getDomainSpace()) &&
+         "Domain of `this` should be compatible with domain of `rel`");
+
+  inverse();
+  compose(rel);
+  inverse();
+}
+
+void PresburgerRelation::applyRange(const PresburgerRelation &rel) {
+  compose(rel);
+}
+
 /// Return the coefficients of the ineq in `rel` specified by  `idx`.
 /// `idx` can refer not only to an actual inequality of `rel`, but also
 /// to either of the inequalities that make up an equality in `rel`.
 ///
 /// When 0 <= idx < rel.getNumInequalities(), this returns the coeffs of the
 /// idx-th inequality of `rel`.
 ///
 /// Otherwise, it is then considered to index into the ineqs corresponding to

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
   IntegerRelationTest.cpp
   LinearTransformTest.cpp
   MatrixTest.cpp
   MPIntTest.cpp
   Parser.h
   ParserTest.cpp
   PresburgerSetTest.cpp
+  PresburgerRelationTest.cpp
   PresburgerSpaceTest.cpp
   PWMAFunctionTest.cpp
   SimplexTest.cpp
   UtilsTest.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
   PRIVATE MLIRPresburger
   MLIRAffineAnalysis
   MLIRParser
   )

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

+//===- PresburgerRelationTest.cpp - Tests for PresburgerRelation class ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/Presburger/PresburgerRelation.h"
+#include "Parser.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <iostream>
+
+using namespace mlir;
+using namespace presburger;
+
+static PresburgerRelation
+parsePresburgerRelationFromPresburgerSet(ArrayRef<StringRef> strs,
+                                         unsigned numDomain) {
+  assert(!strs.empty() && "strs should not be empty");
+
+  IntegerRelation rel = parseIntegerPolyhedron(strs[0]);
+  rel.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain);
+  PresburgerRelation result(rel);
+  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, 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)
+         "(x, y, a, b)[N] : (a - x - y == 0, b - x + y == 0)"},
+        2);
+    PresburgerRelation map2 = parsePresburgerRelationFromPresburgerSet(
+        {// (x, y) -> (x + y)
+         "(x, y, r)[N] : (r - x - y == 0)",
+         // (x, y) -> (N)
+         "(x, y, r)[N] : (r - N == 0)",
+         // (x, y) -> (y - x)
+         "(x, y, r)[N] : (r + x - y == 0)"},
+        2);
+
+    map1.applyRange(map2);
+
+    PresburgerRelation map3 = parsePresburgerRelationFromPresburgerSet(
+        {
+            // (x, y) -> (x + y)
+            "(x, y, r)[N] : (r - x - y == 0)",
+            // (x, y) -> (N)
+            "(x, y, r)[N] : (r - N == 0)",
+            // (x, y) -> (y - x - 2N)
+            "(x, y, r)[N] : (r - y + x + 2 * N == 0)",
+            // (x, y) -> (x - y)
+            "(x, y, r)[N] : (r - x + y == 0)",
+            // (x, y) -> (2x)
+            "(x, y, r)[N] : (r - 2 * x == 0)",
+            // (x, y) -> (-2y)
+            "(x, y, r)[N] : (r + 2 * y == 0)",
+        },
+        2);
+
+    EXPECT_TRUE(map1.isEqual(map3));
+  }
+
+  {
+    PresburgerRelation map1 = parsePresburgerRelationFromPresburgerSet(
+        {// (x, y) -> (y, x)
+         "(x, y, a, b)[N] : (y - a == 0, x - b == 0)",
+         // (x, y) -> (x + N, y - N)
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0)"},
+        2);
+    PresburgerRelation map2 = parsePresburgerRelationFromPresburgerSet(
+        {// (x, y) -> (x - y)
+         "(x, y, r)[N] : (x - y - r == 0)",
+         // (x, y) -> N
+         "(x, y, r)[N] : (N - r == 0)"},
+        2);
+
+    map1.applyDomain(map2);
+
+    PresburgerRelation map3 = parsePresburgerRelationFromPresburgerSet(
+        {// (y - x) -> (x, y)
+         "(r, x, y)[N] : (y - x - r == 0)",
+         // (x - y - 2N) -> (x, y)
+         "(r, x, y)[N] : (x - y - 2 * N - r == 0)",
+         // (x, y) -> N
+         "(r, x, y)[N] : (N - r == 0)"},
+        1);
+
+    EXPECT_TRUE(map1.isEqual(map3));
+  }
+}
+
+TEST(PresburgerRelationTest, inverse) {
+  {
+    PresburgerRelation rel = parsePresburgerRelationFromPresburgerSet(
+        {// (x, y) -> (-y, -x)
+         "(x, y, a, b)[N] : (y + a == 0, x + b == 0)",
+         // (x, y) -> (x + N, y - N)
+         "(x, y, a, b)[N] : (x - a + N == 0, y - b - N == 0)"},
+        2);
+
+    rel.inverse();
+
+    PresburgerRelation inverseRel = parsePresburgerRelationFromPresburgerSet(
+        {// (x, y) -> (-y, -x)
+         "(x, y, a, b)[N] : (y + a == 0, x + b == 0)",
+         // (x, y) -> (x - N, y + N)
+         "(x, y, a, b)[N] : (x - N - a == 0, y + N - b == 0)"},
+        2);
+
+    EXPECT_TRUE(rel.isEqual(inverseRel));
+  }
+}