[MLIR][Presburger] Implement convertVarKind for PresburgerSpace

[iambrj]

Summary

This patch implements convertVarKind for PresburgerSpace

[Groverkss]

This was already added as part of https://reviews.llvm.org/rGd222b69093d69e034b7f01d9ceedce28d48fbb0c.

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

public:
 
   /// Removes variables of the specified kind in the column range [varStart,
   /// varLimit). The range is relative to the kind of variable.
   void removeVarRange(VarKind kind, unsigned varStart, unsigned varLimit);
 
   /// Swaps the posA^th variable of kindA and posB^th variable of kindB.
   void swapVar(VarKind kindA, VarKind kindB, unsigned posA, unsigned posB);
 
+  /// Converts variables of kind srcKind in the range [varStart, varLimit) to
+  /// variables of kind dstKind. If `pos` is given, the variables are placed at
+  /// position `pos` of dstKind, otherwise they are placed after all the other
+  /// variables of kind dstKind. The internal ordering among the moved variables
+  /// is preserved.
+  void convertVarKind(VarKind srcKind, unsigned varStart, unsigned varLimit,
+                      VarKind dstKind, unsigned pos);
+  void convertVarKind(VarKind srcKind, unsigned varStart, unsigned varLimit,
+                      VarKind dstKind) {
+    convertVarKind(srcKind, varStart, varLimit, dstKind,
+                   getNumVarKind(dstKind));
+  }
+
   /// Returns true if both the spaces are compatible i.e. if both spaces have
   /// the same number of variables of each kind (excluding locals).
   bool isCompatible(const PresburgerSpace &other) const;
 
   /// Returns true if both the spaces are equal including local variables i.e.
   /// if both spaces have the same number of variables of each kind (including
   /// locals).
   bool isEqual(const PresburgerSpace &other) const;

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

void PresburgerSpace::swapVar(VarKind kindA, VarKind kindB, unsigned posA,
   if (kindB == VarKind::Local) {
     atId(kindA, posA) = nullptr;
     return;
   }
 
   std::swap(atId(kindA, posA), atId(kindB, posB));
 }
 
+void PresburgerSpace::convertVarKind(VarKind srcKind, unsigned varStart,
+                                     unsigned varLimit, VarKind dstKind,
+                                     unsigned pos) {
+  assert(varLimit <= getNumVarKind(srcKind) && "Invalid id range");
+
+  if (varStart >= varLimit)
+    return;
+
+  // Append new variables corresponding to the dimensions to be converted.
+  unsigned convertCount = varLimit - varStart;
+  unsigned newVarsBegin =
+      insertVar(dstKind, pos, convertCount) - getVarKindOffset(dstKind);
+
+  // Swap the new variables with old variables.
+  //
+  // Essentially, this moves the identifier information corresponding to the
+  // specified ids of kind `srcKind` to the `convertCount` newly created ids of
+  // kind `dstKind`.
+  for (unsigned i = 0; i < convertCount; ++i)
+    swapVar(srcKind, dstKind, varStart + i, newVarsBegin + i);
+
+  // Complete the move by deleting the now redundant variables.
+  removeVarRange(srcKind, varStart, varLimit);
+}
+
 bool PresburgerSpace::isCompatible(const PresburgerSpace &other) const {
   return getNumDomainVars() == other.getNumDomainVars() &&
          getNumRangeVars() == other.getNumRangeVars() &&
          getNumSymbolVars() == other.getNumSymbolVars();
 }
 
 bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
   return isCompatible(other) && getNumLocalVars() == other.getNumLocalVars();

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

TEST(PresburgerSpaceTest, removeVarRangeIdentifier) {
 
   // Check if domain identifiers are attached properly.
   EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 0), identifiers[1]);
 
   // Check if symbol identifiers are attached properly.
   EXPECT_EQ(*space.getId<int *>(VarKind::Range, 0), identifiers[4]);
   EXPECT_EQ(*space.getId<int *>(VarKind::Range, 1), identifiers[5]);
 }
+
+TEST(PresburgerSpaceTest, convertVarKind) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0);
+  space.resetIds<int *>();
+
+  int identifiers[6] = {0, 1, 2, 3, 4, 5};
+
+  // Attach identifiers to domain identifiers.
+  space.setId<int *>(VarKind::Domain, 0, &identifiers[0]);
+  space.setId<int *>(VarKind::Domain, 1, &identifiers[1]);
+
+  // Attach identifiers to range identifiers.
+  space.setId<int *>(VarKind::Range, 0, &identifiers[2]);
+
+  // Attach identifiers to symbol identifiers.
+  space.setId<int *>(VarKind::Symbol, 0, &identifiers[3]);
+  space.setId<int *>(VarKind::Symbol, 1, &identifiers[4]);
+  space.setId<int *>(VarKind::Symbol, 2, &identifiers[5]);
+
+  // Make a few kind conversions.
+  space.convertVarKind(VarKind::Domain, 0, 1, VarKind::Range, 0);
+  space.convertVarKind(VarKind::Symbol, 1, 3, VarKind::Domain, 1);
+  space.convertVarKind(VarKind::Symbol, 0, 1, VarKind::Range, 1);
+
+  // Check if var counts are correct.
+  EXPECT_EQ(space.getNumDomainVars(), 3u);
+  EXPECT_EQ(space.getNumRangeVars(), 3u);
+  EXPECT_EQ(space.getNumSymbolVars(), 0u);
+
+  // Check if identifiers are transferred correctly.
+  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 0), identifiers[1]);
+  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 1), identifiers[4]);
+  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 2), identifiers[5]);
+  EXPECT_EQ(*space.getId<int *>(VarKind::Range, 0), identifiers[0]);
+  EXPECT_EQ(*space.getId<int *>(VarKind::Range, 1), identifiers[3]);
+  EXPECT_EQ(*space.getId<int *>(VarKind::Range, 2), identifiers[2]);
+}