diff --git a/mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h b/mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h --- a/mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h +++ b/mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h @@ -145,6 +145,19 @@ /// 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; diff --git a/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp b/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp --- a/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp +++ b/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp @@ -152,6 +152,31 @@ 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() && diff --git a/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp --- a/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp +++ b/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp @@ -108,3 +108,40 @@ EXPECT_EQ(*space.getId(VarKind::Range, 0), identifiers[4]); EXPECT_EQ(*space.getId(VarKind::Range, 1), identifiers[5]); } + +TEST(PresburgerSpaceTest, convertVarKind) { + PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0); + space.resetIds(); + + int identifiers[6] = {0, 1, 2, 3, 4, 5}; + + // Attach identifiers to domain identifiers. + space.setId(VarKind::Domain, 0, &identifiers[0]); + space.setId(VarKind::Domain, 1, &identifiers[1]); + + // Attach identifiers to range identifiers. + space.setId(VarKind::Range, 0, &identifiers[2]); + + // Attach identifiers to symbol identifiers. + space.setId(VarKind::Symbol, 0, &identifiers[3]); + space.setId(VarKind::Symbol, 1, &identifiers[4]); + space.setId(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(VarKind::Domain, 0), identifiers[1]); + EXPECT_EQ(*space.getId(VarKind::Domain, 1), identifiers[4]); + EXPECT_EQ(*space.getId(VarKind::Domain, 2), identifiers[5]); + EXPECT_EQ(*space.getId(VarKind::Range, 0), identifiers[0]); + EXPECT_EQ(*space.getId(VarKind::Range, 1), identifiers[3]); + EXPECT_EQ(*space.getId(VarKind::Range, 2), identifiers[2]); +}