[MLIR][Presburger] Implement convertVarKind for PresburgerRelation

Author: iambrj

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

@@ -66,6 +66,14 @@ class PresburgerRelation {
 
   void insertVarInPlace(VarKind kind, unsigned pos, unsigned num = 1);
 
+  /// Converts variables of the specified kind in the column range [srcPos,
+  /// srcPos + num) to variables of the specified kind at position dstPos. The
+  /// ranges are relative to the kind of variable.
+  ///
+  /// srcKind and dstKind must be different.
+  void convertVarKind(VarKind srcKind, unsigned srcPos, unsigned num,
+                      VarKind dstKind, unsigned dstPos);
+
   /// Return a reference to the list of disjuncts.
   ArrayRef<IntegerRelation> getAllDisjuncts() const;
 

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

@@ -9,6 +9,7 @@
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
@@ -38,6 +39,23 @@ void PresburgerRelation::insertVarInPlace(VarKind kind, unsigned pos,
   space.insertVar(kind, pos, num);
 }
 
+void PresburgerRelation::convertVarKind(VarKind srcKind, unsigned srcPos,
+                                        unsigned num, VarKind dstKind,
+                                        unsigned dstPos) {
+  assert(srcKind != VarKind::Local && dstKind != VarKind::Local &&
+      "srcKind/dstKind cannot be local");
+  assert(srcKind != dstKind && "cannot convert variables to the same kind");
+  assert(srcPos + num <= space.getNumVarKind(srcKind) &&
+         "invalid range for source variables");
+  assert(dstPos <= space.getNumVarKind(dstKind) &&
+         "invalid position for destination variables");
+
+  space.convertVarKind(srcKind, srcPos, num, dstKind, dstPos);
+
+  for (IntegerRelation &disjunct : disjuncts)
+    disjunct.convertVarKind(srcKind, srcPos, srcPos + num, dstKind, dstPos);
+}
+
 unsigned PresburgerRelation::getNumDisjuncts() const {
   return disjuncts.size();
 }

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

@@ -16,14 +16,6 @@
 using namespace mlir;
 using namespace presburger;
 
-static IntegerRelation parseRelationFromSet(StringRef set, unsigned numDomain) {
-  IntegerRelation rel = parseIntegerPolyhedron(set);
-
-  rel.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain);
-
-  return rel;
-}
-
 TEST(IntegerRelationTest, getDomainAndRangeSet) {
   IntegerRelation rel = parseRelationFromSet(
       "(x, xr)[N] : (xr - x - 10 == 0, xr >= 0, N - xr >= 0)", 1);

mlir/unittests/Analysis/Presburger/Parser.h

@@ -80,6 +80,27 @@ parsePWMAF(ArrayRef<std::pair<StringRef, StringRef>> pieces) {
   return func;
 }
 
+inline IntegerRelation parseRelationFromSet(StringRef set, unsigned numDomain) {
+  IntegerRelation rel = parseIntegerPolyhedron(set);
+
+  rel.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain);
+
+  return rel;
+}
+
+inline PresburgerRelation
+parsePresburgerRelationFromPresburgerSet(ArrayRef<StringRef> strs,
+                                         unsigned numDomain) {
+  assert(!strs.empty() && "strs should not be empty");
+
+  IntegerRelation rel = parseIntegerPolyhedron(strs[0]);
+  PresburgerRelation result(rel);
+  for (unsigned i = 1, e = strs.size(); i < e; ++i)
+    result.unionInPlace(parseIntegerPolyhedron(strs[i]));
+  result.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain, 0);
+  return result;
+}
+
 } // namespace presburger
 } // namespace mlir
 

mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "Parser.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 
 #include <gmock/gmock.h>
@@ -16,22 +17,6 @@
 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(
@@ -291,3 +276,43 @@ TEST(PresburgerRelationTest, getDomainAndRangeSet) {
 
   EXPECT_TRUE(rangeSet.isEqual(expectedRangeSet));
 }
+
+TEST(PresburgerRelationTest, convertVarKind) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0);
+
+  IntegerRelation disj1 = parseRelationFromSet(
+                      "(x, y, a)[U, V, W] : (x - U == 0, y + a - W == 0,"
+                      "U - V >= 0, y - a >= 0)",
+                      2),
+                  disj2 = parseRelationFromSet(
+                      "(x, y, a)[U, V, W] : (x + y - U == 0, x - a + V == 0,"
+                      "V - U >= 0, y + a >= 0)",
+                      2);
+
+  PresburgerRelation rel(disj1);
+  rel.unionInPlace(disj2);
+
+  // Make a few kind conversions.
+  rel.convertVarKind(VarKind::Domain, 0, 1, VarKind::Range, 0);
+  rel.convertVarKind(VarKind::Symbol, 1, 2, VarKind::Domain, 1);
+  rel.convertVarKind(VarKind::Symbol, 0, 1, VarKind::Range, 1);
+
+  // Expected rel.
+  disj1.convertVarKind(VarKind::Domain, 0, 1, VarKind::Range, 0);
+  disj1.convertVarKind(VarKind::Symbol, 1, 3, VarKind::Domain, 1);
+  disj1.convertVarKind(VarKind::Symbol, 0, 1, VarKind::Range, 1);
+  disj2.convertVarKind(VarKind::Domain, 0, 1, VarKind::Range, 0);
+  disj2.convertVarKind(VarKind::Symbol, 1, 3, VarKind::Domain, 1);
+  disj2.convertVarKind(VarKind::Symbol, 0, 1, VarKind::Range, 1);
+
+  PresburgerRelation expectedRel(disj1);
+  expectedRel.unionInPlace(disj2);
+
+  // Check if var counts are correct.
+  EXPECT_EQ(rel.getNumDomainVars(), 3u);
+  EXPECT_EQ(rel.getNumRangeVars(), 3u);
+  EXPECT_EQ(rel.getNumSymbolVars(), 0u);
+
+  // Check if identifiers are transferred correctly.
+  EXPECT_TRUE(expectedRel.isEqual(rel));
+}