[MLIR][Presburger] support IntegerRelation::convertIdKind

[Groverkss]

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

397–399

Can we just have convertIdKind? I don't see the reason to have convertToLocal if we have a more general function here.

[arjunp]

clang-format

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

397–399

Seems like a common use case, and making a default parameter for the other one doesn't make sense to me.

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1133–1135

Can you put a comment here on why we have to do this swapping? The reasoning IIRC is that there can be some information attached to identifiers which we need to carry over.

[arjunp]

• add comments
• clang-format

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

public:
 
   /// Removes redundant constraints using Simplex. Although the algorithm can
   /// theoretically take exponential time in the worst case (rare), it is known
   /// 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();
 
-  /// Converts identifiers in the column range [idStart, idLimit) to local
-  /// variables.
-  void convertDimToLocal(unsigned dimStart, unsigned dimLimit);
+  /// 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.
+  void convertIdKind(IdKind srcKind, unsigned idStart, unsigned idLimit,
+                     IdKind 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`. The ordering of the local ids in the
   /// sets may also be changed. After merging, if the `i^th` local variable in
   /// one set has a known division representation, then the `i^th` local
   /// variable in the other set either has the same division representation or
   /// no known division representation.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

for (unsigned k = 0, t = getNumInequalities(); k < t; ++k)
       eliminateFromConstraint(this, k, i, j, j, /*isEq=*/false);
 
     // Remove the ith equality and the found local variable.
     removeId(j);
     removeEquality(i);
   }
 }
 
-void IntegerRelation::convertDimToLocal(unsigned dimStart, unsigned dimLimit) {
-  assert(dimLimit <= getNumDimIds() && "Invalid dim pos range");
+void IntegerRelation::convertIdKind(IdKind srcKind, unsigned idStart,
+                                    unsigned idLimit, IdKind dstKind) {
+  assert(idLimit <= getNumIdKind(srcKind) && "Invalid id range");
 
-  if (dimStart >= dimLimit)
+  if (idStart >= idLimit)
     return;
 
   // Append new local variables corresponding to the dimensions to be converted.
-  unsigned convertCount = dimLimit - dimStart;
-  unsigned newLocalIdStart = getNumIds();
-  appendId(IdKind::Local, convertCount);
+  unsigned newIdsBegin = getIdKindEnd(dstKind);
+  unsigned convertCount = idLimit - idStart;
+  appendId(dstKind, 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).
+  unsigned offset = getIdKindOffset(srcKind);
   for (unsigned i = 0; i < convertCount; ++i)
-    swapId(i + dimStart, i + newLocalIdStart);
+    swapId(offset + idStart + i, newIdsBegin + i);
 
-  // Remove dimensions converted to local variables.
-  removeIdRange(IdKind::SetDim, dimStart, dimLimit);
+  // Complete the move by deleting the initially occupied columns.
+  removeIdRange(srcKind, idStart, idLimit);
 }
 
 void IntegerRelation::addBound(BoundType type, unsigned pos, int64_t value) {
   assert(pos < getNumCols());
   if (type == BoundType::EQ) {
     unsigned row = equalities.appendExtraRow();
     equalities(row, pos) = 1;
     equalities(row, getNumCols() - 1) = -value;

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

AffineMap mlir::alignAffineMapWithValues(AffineMap map, ValueRange operands,
 
   return map.replaceDimsAndSymbols(dimReplacements, symReplacements,
                                    dims.size(), numSymbols);
 }
 
 FlatAffineValueConstraints FlatAffineRelation::getDomainSet() const {
   FlatAffineValueConstraints domain = *this;
   // Convert all range variables to local variables.
-  domain.convertDimToLocal(getNumDomainDims(),
-                           getNumDomainDims() + getNumRangeDims());
+  domain.convertToLocal(IdKind::SetDim, getNumDomainDims(),
   return domain;
 }
 
 FlatAffineValueConstraints FlatAffineRelation::getRangeSet() const {
   FlatAffineValueConstraints range = *this;
   // Convert all domain variables to local variables.
-  range.convertDimToLocal(0, getNumDomainDims());
+  range.convertToLocal(IdKind::SetDim, 0, getNumDomainDims());
   return range;
 }
 
 void FlatAffineRelation::compose(const FlatAffineRelation &other) {
   assert(getNumDomainDims() == other.getNumRangeDims() &&
          "Domain of this and range of other do not match");
   assert(std::equal(values.begin(), values.begin() + getNumDomainDims(),
                     other.values.begin() + other.getNumDomainDims()) &&
void FlatAffineRelation::compose(const FlatAffineRelation &other) {
 
   // Merge symbol and local identifiers.
   mergeSymbolIds(rel);
   mergeLocalIds(rel);
 
   // Convert `rel` from [otherDomain] -> [otherRange thisRange] to
   // [otherDomain] -> [thisRange] by converting first otherRange range ids
   // to local ids.
-  rel.convertDimToLocal(rel.getNumDomainDims(),
-                        rel.getNumDomainDims() + removeDims);
+  rel.convertToLocal(IdKind::SetDim, rel.getNumDomainDims(),
   // Convert `this` from [otherDomain thisDomain] -> [thisRange] to
   // [otherDomain] -> [thisRange] by converting last thisDomain domain ids
   // to local ids.
-  convertDimToLocal(getNumDomainDims() - removeDims, getNumDomainDims());
+  convertToLocal(IdKind::SetDim, getNumDomainDims() - removeDims,
+                 getNumDomainDims());
 
   auto thisMaybeValues = getMaybeDimValues();
   auto relMaybeValues = rel.getMaybeDimValues();
 
   // Add and match domain of `rel` to domain of `this`.
   for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
     if (relMaybeValues[i].hasValue())
       setValue(i, relMaybeValues[i].getValue());

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

TEST(IntegerPolyhedronTest, computeLocalReprTightUpperBound) {
     // Check if the divisions can be computed even with a tighter upper bound.
     checkDivisionRepresentation(poly, divisions, denoms);
   }
 
   {
     IntegerPolyhedron poly =
         parsePoly("(i, j, q) : (4*q - i - j + 2 >= 0, -4*q + i + j >= 0)");
     // Convert `q` to a local variable.
-    poly.convertDimToLocal(2, 3);
+    poly.convertToLocal(IdKind::SetDim, 2, 3);
 
     std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, 1}};
     SmallVector<unsigned, 8> denoms = {4};
 
     // Check if the divisions can be computed even with a tighter upper bound.
     checkDivisionRepresentation(poly, divisions, denoms);
   }
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprFromEquality) {
   {
     IntegerPolyhedron poly = parsePoly("(i, j, q) : (-4*q + i + j == 0)");
     // Convert `q` to a local variable.
-    poly.convertDimToLocal(2, 3);
+    poly.convertToLocal(IdKind::SetDim, 2, 3);
 
     std::vector<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
     SmallVector<unsigned, 8> denoms = {4};
 
     checkDivisionRepresentation(poly, divisions, denoms);
   }
   {
     IntegerPolyhedron poly = parsePoly("(i, j, q) : (4*q - i - j == 0)");
     // Convert `q` to a local variable.
-    poly.convertDimToLocal(2, 3);
+    poly.convertToLocal(IdKind::SetDim, 2, 3);
 
     std::vector<SmallVector<int64_t, 8>> divisions = {{-1, -1, 0, 0}};
     SmallVector<unsigned, 8> denoms = {4};
 
     checkDivisionRepresentation(poly, divisions, denoms);
   }
   {
     IntegerPolyhedron poly = parsePoly("(i, j, q) : (3*q + i + j - 2 == 0)");
     // Convert `q` to a local variable.
-    poly.convertDimToLocal(2, 3);
+    poly.convertToLocal(IdKind::SetDim, 2, 3);
 
     std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, -2}};
     SmallVector<unsigned, 8> denoms = {3};
 
     checkDivisionRepresentation(poly, divisions, denoms);
   }
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprFromEqualityAndInequality) {
   {
     IntegerPolyhedron poly =
         parsePoly("(i, j, q, k) : (-3*k + i + j == 0, 4*q - "
                   "i - j + 2 >= 0, -4*q + i + j >= 0)");
     // Convert `q` and `k` to local variables.
-    poly.convertDimToLocal(2, 4);
+    poly.convertToLocal(IdKind::SetDim, 2, 4);
 
     std::vector<SmallVector<int64_t, 8>> divisions = {{1, 1, 0, 0, 1},
                                                       {-1, -1, 0, 0, 0}};
     SmallVector<unsigned, 8> denoms = {4, 3};
 
     checkDivisionRepresentation(poly, divisions, denoms);
   }
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprNoRepr) {
   IntegerPolyhedron poly =
       parsePoly("(x, q) : (x - 3 * q >= 0, -x + 3 * q + 3 >= 0)");
   // Convert q to a local variable.
-  poly.convertDimToLocal(1, 2);
+  poly.convertToLocal(IdKind::SetDim, 1, 2);
 
   std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0}};
   SmallVector<unsigned, 8> denoms = {0};
 
   // Check that no division is computed.
   checkDivisionRepresentation(poly, divisions, denoms);
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprNegConstNormalize) {
   IntegerPolyhedron poly =
       parsePoly("(x, q) : (-1 - 3*x - 6 * q >= 0, 6 + 3*x + 6*q >= 0)");
   // Convert q to a local variable.
-  poly.convertDimToLocal(1, 2);
+  poly.convertToLocal(IdKind::SetDim, 1, 2);
 
   // q = floor((-1/3 - x)/2)
   //   = floor((1/3) + (-1 - x)/2)
   //   = floor((-1 - x)/2).
   std::vector<SmallVector<int64_t, 8>> divisions = {{-1, 0, -1}};
   SmallVector<unsigned, 8> denoms = {2};
   checkDivisionRepresentation(poly, divisions, denoms);
 }