[MLIR][Presburger] Implement getAsRelation for PWMAFunction

[iambrj]

Summary

This patch implements getAsRelation for PWMAFunction

[iambrj]

Fix code style formatting

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

162

Could you send this bug fix as a separate patch?

mlir/unittests/Analysis/Presburger/Parser.h

50–64

in this implementation, I would just call parsePresburgerSet, and the call convertVarKind. (If PresburgerRelation does not have convertVarKind, this would be the ideal time to implement it :) )

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

public:
   /// space.
   PresburgerSpace getDomainSpace() const { return space.getDomainSpace(); }
   PresburgerSpace getOutputSpace() const { return space.getDomainSpace(); }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
   /// union of the domains of all the pieces.
   PresburgerSet getDomain() const;
 
+  /// Return the piece-wise MultiAffineFunction as a PresburgerRelation.
+  PresburgerRelation getAsRelation() const;
+
   /// Return the output of the function at the given point.
   std::optional<SmallVector<MPInt, 8>> valueAt(ArrayRef<MPInt> point) const;
   std::optional<SmallVector<MPInt, 8>> valueAt(ArrayRef<int64_t> point) const {
     return valueAt(getMPIntVec(point));
   }
 
   /// Return all the pieces of this piece-wise function.
   ArrayRef<Piece> getAllPieces() const { return pieces; }

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

 
 PresburgerSet PWMAFunction::getDomain() const {
   PresburgerSet domain = PresburgerSet::getEmpty(getDomainSpace());
   for (const Piece &piece : pieces)
     domain.unionInPlace(piece.domain);
   return domain;
 }
 
+PresburgerRelation PWMAFunction::getAsRelation() const {
+  PresburgerRelation result = PresburgerRelation::getEmpty(space);
+  for (const Piece &piece : pieces) {
+    PresburgerRelation rel(piece.output.getAsRelation());
+    rel = rel.intersectDomain(piece.domain);
+    result.unionInPlace(rel);
+  }
+  return result;
+}
+
 void MultiAffineFunction::print(raw_ostream &os) const {
   space.print(os);
   os << "Division Representation:\n";
   divs.print(os);
   os << "Output:\n";
   output.print(os);
 }
 

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 }
 
 PresburgerRelation
 PresburgerRelation::intersectDomain(const PresburgerSet &set) const {
   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.insertVarInPlace(VarKind::Domain, 0, getNumRangeVars());
   other.inverse();
   return intersect(other);
 }
 
 PresburgerSet PresburgerRelation::getDomainSet() const {
   PresburgerSet result = PresburgerSet::getEmpty(space.getDomainSpace());
   for (const IntegerRelation &cs : disjuncts)
     result.unionInPlace(cs.getDomainSet());

mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp

PWMAFunction result = parsePWMAF({
          "(x) -> (x floordiv 16)"},
         {"(x) : (x >= 0, x <= 1000, x floordiv 16 >= (x + 10) floordiv 17 + 1)",
          "(x) -> ((x + 10) floordiv 17)"},
     });
 
     EXPECT_TRUE(func1.unionLexMin(func2).isEqual(result));
   }
 }
+
+TEST(PWAFunctionTest, getAsRelation) {
+  PWMAFunction func = parsePWMAF({
+      {"(x, y)[N] : (N >= 0, x - N - 1 >= 0, y - N - 1 >= 0)",
+       "(x, y)[N] -> (x + y)"},
+      {"(x, y)[N] : (N >= 0, x - N >= 0, -y - N >= 0)", "(x, y)[N] -> (x - y)"},
+      {"(x, y)[N] : (N >= 0, -x - N - 1 >= 0, y - N >= 0)",
+       "(x, y)[N] -> (y - x)"},
+      {"(x, y)[N] : (N >= 0, -x - N - 1 >= 0, -y - N - 1 >= 0)",
+       "(x, y)[N] -> (-x - y)"},
+  });
+  PresburgerRelation rel = parsePresburgerRelationFromPresburgerSet(
+      {"(x, y, a)[N] : (N >= 0, x - N - 1 >= 0, y - N - 1 >= 0, a - x - y == "
+       "0)",
+       "(x, y, a)[N] : (N >= 0, x - N >= 0, -y - N >= 0, a - x + y == 0)",
+       "(x, y, a)[N] : (N >= 0, -x - N - 1 >= 0, y - N >= 0, a + x - y == 0)",
+       "(x, y, a)[N] : (N >= 0, -x - N - 1 >= 0, -y - N - 1 >= 0, a + x + y == "
+       "0)"},
+      2);
+
+  EXPECT_TRUE(rel.isEqual(func.getAsRelation()));
+}

mlir/unittests/Analysis/Presburger/Parser.h

inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
 
   IntegerPolyhedron initPoly = parseIntegerPolyhedron(strs[0]);
   PresburgerSet result(initPoly);
   for (unsigned i = 1, e = strs.size(); i < e; ++i)
     result.unionInPlace(parseIntegerPolyhedron(strs[i]));
   return result;
 }
 
+inline 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;
+}
+
 inline MultiAffineFunction parseMultiAffineFunction(StringRef str) {
   MLIRContext context(MLIRContext::Threading::DISABLED);
 
   // TODO: Add default constructor for MultiAffineFunction.
   MultiAffineFunction multiAff(PresburgerSpace::getRelationSpace(),
                                Matrix(0, 1));
   if (getMultiAffineFunctionFromMap(parseAffineMap(str, &context), multiAff)
           .failed())

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

 
 #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, 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)"},