[MLIR][Presburger] Add Identifier class to store identifiers and their type

[Groverkss]

Summary

This patch adds a new class Identifier to store identifiers in PresburgerSpace
and their types.

Identifiers were added earlier and were stored as a void pointer, and their type
in the form of mlir::TypeId in PresburgerSpace. To get an identifier, a user of
PresburgerSpace needed to know the type of identifiers. This was a problem for
users of PresburgerSpace like IntegerRelation, which want to work on
identifiers without knowing their type.

The Identifier class allows users like IntegerRelation to work on identifiers
without knowing their type, and also exposes an easier way to work with
Identifiers.

[arjunp]

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

31–35

I think we need much more documentation to explain all this to someone who doesn't already know what it is.

155–159

158

"actually same variable" doesn't seem to make sense within the context of PresburgerSpace.

Can you instead say something like, the user can identify each variable in the space as corresponding to some Identifier, and then explain clearly how and where these identifiers are used, if at all. This should make sense when taken together with the identifier documentation

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

90–93

why was the API changed here?

[bondhugula]

I'd requesting expanding commit summary to include sufficient rationale and summary of the change, and what it enables.

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

31–35

I agree. Examples will also help here.

84–86

What's this?

[Groverkss]

• Address comments

[Groverkss]

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

84–86

oops, removed.

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

90–93

To make it consistent with how we do this through Presburger Library. We don't have setters; instead, we use get as a reference and set it that way.

[arjunp]

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

31

or replace with some other examples you have in mind

33

two variables in what? same set? different sets with same space? different space?

39

this notation is unclear. can you write some kind of example statement S0: A[i][j] = 0 instead

49

The discussion here seems to imply that the names i and j are somehow relevant here. can your clarify this

56

58

60

not clear what "always unique" means. can you just say it's treated as no attachment and then describe the behaviour of such variables

[Groverkss]

• Address more comments
• Address mroe comments
• isEqual type assert

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

33

I don't think it matters? They can be used for two different variables in anything.

[arjunp]

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

41

might be clearer if you use j and k instead of two js

46

I can't understand this syntax. Are they two spaces? Above it says "a PresburgerSpace"

48

56

Documentation is unclear on whether different variables can have different types of objects attached

56–58

"previous example" is unclear whether it is referring to the same code example, or the actual previous example where you attached Values to identifiers. I assume you mean the former, if so can you write it out explicitly.

61

66

Why are the impls in the header?

81

Can you write a longer assert message explaining what's happening

89

Same pointer with different types can only happen due to an internal error in the class, right? So why are we telling the user about that in the API comment?

158

159

286–287

Please clarify what happens with null ids

289–290

303–304

What is this? Please document.

[Groverkss]

• Address more comments

[Groverkss]

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

56

That's a user class detail? I don't see why that should be in the documentation of Identifier.

56–58

I feel that paragraph was unnecessary. I removed that paragraph, as it just caused problems.

66

Templated methods need to be implemented in the header file.

89

I want to allow Identifiers in the same object to be of different types, so I want to keep this in the documentation.

286–287

Currently, null ids are treated as not equal and the check fails. I added a line that specifies what happens when identifiers are not being used. I will specify the null ids behavior more in another patch. But for now, they are treated as not equal.

303–304

Sorry, accidentally committed this.

[arjunp]

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

56

What do you mean user class detail? I just want you to write here that the type of object attached can be different for different variables in the space.

66

I would put the remaining methods implementations in the cpp to reduce compile time, not sure if there are any other considerations in favour of keeping them in the header.

89

My bad, I was assuming the same pointer can't have different types right, and that the assert will only fail if we somehow messed up with storing pointers. But actually we could have released the pointer and been allocated another one with a different type. Just add another sentence explaining why this could happen (only because of realloc of a different type of object to the same pointer). This seems like a weird failure case, I think we should make clear what the behaviour of identifier is with deallocating/reallocating pointers (because it affects things like this).

260–264

kind != VarKind::Local

291

Can you write a sentence explaining why this overload is useful

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

225

226–229

use the getIds function you defined?

234–246

Is the notion of a space "using" or "not using" identifiers explained in any of the class-level docs?

236

personally I prefer !=

238–239

they also have to be compatible right

261

like above, this check needs to be first I believe

280

this can just be a free function

282

i < e

282

Use getIds and range-based for?

284

should we add a contextual (explicit) conversion to bool?

[Groverkss]

• Address comments

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

56

I meant that PresburgerSpace allows for both, but the using class (say IntegerRelation) decides if the objects attached need to be of the same type or not. Anyway, I added the line you wanted.

66

Most of them are one-linears :P I moved isEqual, print and dump in the implementation file.

89

We allow storing anything that has llvm::PointerLikeTypeTraits. I could define PointerLikeTypeTraits for a custom integer and store things there, so there are more cases where this could happen. I made the behavior of equality more clear in the Identifier definition instead.

As long as the equality behavior is clear and we are getting a descriptive assert, I think we don't need to write the realloc sentence.

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

234–246

Yes, last paragraph in PresburgerSpace documentation.

238–239

Identifiers being equal automatically implies compatibility, added a comment about this.

280

Sure, but I don't see why I would need this as a free function instead. I don't think any other function would need printing and it's cleaner to write with a lambda capture.

284

I prefer explicit checking.

[arjunp]

Please add some tests for convertVarKind. Other than that, some last nits below:

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

49

pls change "Are different" to something like "correspond to different SSA values in the program"

56

Yes, so PresburgerSpace should mention that it allows both.

56–57

similarly here -- check they correspond to same variable in the program

58

60–61

62

63

64

70

consider writing this constraint in the class doc instead?

156

286–287

[Groverkss]

• Address comments
• Address comments
• Fix convertVarKind test

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

70

I think of this is a method detail rather than a class detail. The idea of identifier is still the same, this is just how we implement it. I would prefer it to be here.

286–287

"in both spaces have the same identifiers" does not make it clear that they need to be at the same position as well.

[arjunp]

Please add tests for the other two cases (from and to locals). I really find these kinds of things finicky, as we just saw, so I think it's better to test them a bit exhaustively.

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

70

But this is a constraint of how the class can be used, right?

286–287

ok can you write "same identifiers in the same positions" then? I find the current order confusing.

[Groverkss]

• Address comments

[Groverkss]

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

70

Not really? It's just this specific constructor uses it to build an identifier, which is why it's in the constructor documentation. To me, the identifier has no relation with this llvm::PointerLikeTypeTraits. It's just the constructor does. Nothing else in the class cares about it or should ever care about it.

It's not a constraint on the class. You can easily add a constructor that takes a void* pointer and a type and it should be fine as well.

[arjunp]

The class currently uses type traits right? getValue currently uses llvm::PointerLikeTypeTraits<T>::getFromVoidPointer(value);. Please mention the constraint in the class doc before landing.

[Groverkss]

• Documentation

[Groverkss]

• Fix tests
• Fix warning in unittests

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

 namespace mlir {
 namespace presburger {
 
 /// Kind of variable. Implementation wise SetDims are treated as Range
 /// vars, and spaces with no distinction between dimension vars are treated
 /// as relations with zero domain vars.
 enum class VarKind { Symbol, Local, Domain, Range, SetDim = Range };
 
+// An Identifier stores a pointer to an object. Identifiers are meant to be
+// attached to a variable in a PresburgerSpace and can be used to check if
+// two variables correspond to the same object.
+//
+// Take for example the following code:
+//
+// for i = 0 to 100
+//  for j = 0 to 100
+//    S0 --> Statement 0
+//  for j = 0 to 100
+//    S1 --> Statement 1
+//
+// If we represent the space of iteration variables surrounding S0, S1 in a
+// PresburgerSpace, we have:
+//
+// S0 = {i, j}, S1 = {i, j}
+//
+// Since, the variables are in different spaces, without an identifier, there
+// is no way to distinguish if the `j` variables in the two spaces are
+// different. So, we attach an Identifier containing an MLIR Value to them.
+// Now, S0 = {i(Value1), j(Value2)}, S1 = {i(Value1), j(Value3)}.
+//
+// Using the identifier, we can check that the `i` iteration variables in both
+// the spaces are same, while they are different for `j`.
+//
+// Identifiers can also used to map variables to their defining operation. For
+// example, using the Identifier of `i` in the previous example, we can
+// find the for loop corresponding to the iteration varaible.
+//
+// Identifiers treat null pointers as special values which are always unique.
+// This is useful for representing variables which are not associated with
+// any object.
+class Identifier {
+public:
+  Identifier() = default;
+
+  // Create an identifier from a pointer. The type of the pointer must have
+  // a `llvm::PointerLikeTypeTraits` specialization.
+  template <typename T>
+  explicit Identifier(T value)
+      : value(llvm::PointerLikeTypeTraits<T>::getAsVoidPointer(value)) {
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    idType = TypeID::get<T>();
+#endif
+  }
+
+  /// Get the value of the identifier casted to type `T`. `T` here should match
+  /// the type of the identifier used to create it.
+  template <typename T>
+  T getValue() const {
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    assert(TypeID::get<T>() == idType && "Type mismatch");
+#endif
+    return llvm::PointerLikeTypeTraits<T>::getFromVoidPointer(value);
+  }
+
+  bool hasValue() const { return value != nullptr; }
+
+  /// Check if the two identifiers are equal. Null identifiers are considered
+  /// not equal. Does not assert on type mismatch.
+  bool isEqual(const Identifier &other) const {
+    if (value == nullptr || other.value == nullptr)
+      return false;
+    return value == other.value;
+  }
+
+  bool operator==(const Identifier &other) const { return isEqual(other); }
+  bool operator!=(const Identifier &other) const { return !isEqual(other); }
+
+  void print(llvm::raw_ostream &os) const { os << "Id<" << value << ">"; }
+
+  void dump() const {
+    print(llvm::errs());
+    llvm::errs() << "\n";
+  }
+
+private:
+  /// The value of the identifier.
+  void *value = nullptr;
+
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+  /// TypeID of the identifiers in space. This should be used in asserts only.
+  TypeID idType = TypeID::get<void>();
+#endif
+};
+
 /// PresburgerSpace is the space of all possible values of a tuple of integer
 /// valued variables/variables. Each variable has one of the three types:
 ///
 /// Dimension: Ordinary variables over which the space is represented.
 ///
 /// Symbol: Symbol variables correspond to fixed but unknown values.
 /// Mathematically, a space with symbolic variables is like a
 /// family of spaces indexed by the symbolic variables.

 /// A space with no distinction between types of dimension variables can
 /// be implemented as a space with zero domain. VarKind::SetDim should be used
 /// to refer to dimensions in such spaces.
 ///
 /// Compatibility of two spaces implies that number of variables of each kind
 /// other than Locals are equal. Equality of two spaces implies that number of
 /// variables of each kind are equal.
 ///
-/// PresburgerSpace optionally also supports attaching some information to each
-/// variable in space, called "identifier" of that variable. `resetIds<IdType>`
-/// is used to enable/reset these identifiers. All identifiers must be of the
-/// same type, `IdType`. `IdType` must have a `llvm::PointerLikeTypeTraits`
-/// specialization available and should be supported via `mlir::TypeID`.
-///
-/// These identifiers can be used to check if two variables in two different
-/// spaces are actually same variable.
+/// PresburgerSpace optionally also supports attaching an Identifier with each
+/// non-local variable in space. This is disabled by default. `resetIds` is
+/// used to enable/reset these identifiers. The user can identify each variable
+/// in space as corresponding to some Identifier. Some example use cases are
+/// described in the `Identifier` documentation.
 class PresburgerSpace {
 public:
   static PresburgerSpace getRelationSpace(unsigned numDomain = 0,
                                           unsigned numRange = 0,
                                           unsigned numSymbols = 0,
                                           unsigned numLocals = 0) {
     return PresburgerSpace(numDomain, numRange, numSymbols, numLocals);
   }
public:
   /// If identifiers are being used, the newly added variables have no
   /// identifiers.
   unsigned insertVar(VarKind kind, unsigned pos, unsigned num = 1);
 
   /// 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);
 
+  /// 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);
+
+  /// Changes the partition between dimensions and symbols. Depending on the new
+  /// symbol count, either a chunk of dimensional variables immediately before
+  /// the split become symbols, or some of the symbols immediately after the
+  /// split become dimensions.
+  void setVarSymbolSeperation(unsigned newSymbolCount);
+
   /// Swaps the posA^th variable of kindA and posB^th variable of kindB.
   void swapVar(VarKind kindA, VarKind kindB, unsigned posA, unsigned posB);
 
   /// 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;
 
-  /// Changes the partition between dimensions and symbols. Depending on the new
-  /// symbol count, either a chunk of dimensional variables immediately before
-  /// the split become symbols, or some of the symbols immediately after the
-  /// split become dimensions.
-  void setVarSymbolSeperation(unsigned newSymbolCount);
-
-  void print(llvm::raw_ostream &os) const;
-  void dump() const;
-
-  //===--------------------------------------------------------------------===//
-  //     Identifier Interactions
-  //===--------------------------------------------------------------------===//
-
-  /// Set the identifier for `i^th` variable to `id`. `T` here should match the
-  /// type used to enable identifiers.
-  template <typename T>
-  void setId(VarKind kind, unsigned i, T id) {
-#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-    assert(TypeID::get<T>() == idType && "Type mismatch");
-#endif
-    atId(kind, i) = llvm::PointerLikeTypeTraits<T>::getAsVoidPointer(id);
+  /// Get the identifier of the specified variable.
+  Identifier &getId(VarKind kind, unsigned pos) {
+    assert(VarKind::Local != kind && "Local variables have no identifiers");
+    return identifiers[getVarKindOffset(kind) + pos];
+  }
+  Identifier getId(VarKind kind, unsigned pos) const {
+    assert(VarKind::Local != kind && "Local variables have no identifiers");
+    return identifiers[getVarKindOffset(kind) + pos];
   }
 
-  /// Get the identifier for `i^th` variable casted to type `T`. `T` here
-  /// should match the type used to enable identifiers.
-  template <typename T>
-  T getId(VarKind kind, unsigned i) const {
-#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-    assert(TypeID::get<T>() == idType && "Type mismatch");
-#endif
-    return llvm::PointerLikeTypeTraits<T>::getFromVoidPointer(atId(kind, i));
-  }
-
-  /// Check if the i^th variable of the specified kind has a non-null
-  /// identifier.
-  bool hasId(VarKind kind, unsigned i) const {
-    return atId(kind, i) != nullptr;
-  }
-
-  /// Check if the spaces are compatible, as well as have the same identifiers
-  /// for each variable.
-  bool isAligned(const PresburgerSpace &other) const;
-  /// Check if the number of variables of the specified kind match, and have
-  /// same identifiers with the other space.
-  bool isAligned(const PresburgerSpace &other, VarKind kind) const;
-
-  /// Find the variable of the specified kind with identifier `id`.
-  /// Returns PresburgerSpace::kIdNotFound if identifier is not found.
-  template <typename T>
-  unsigned findId(VarKind kind, T id) const {
-    unsigned i = 0;
-    for (unsigned e = getNumVarKind(kind); i < e; ++i)
-      if (hasId(kind, i) && getId<T>(kind, i) == id)
-        return i;
-    return kIdNotFound;
+  ArrayRef<Identifier> getIds(VarKind kind) const {
+    assert(VarKind::Local != kind && "Local variables have no identifiers");
+    return {identifiers.data() + getVarKindOffset(kind), getNumVarKind(kind)};
   }
-  static const unsigned kIdNotFound = UINT_MAX;
 
   /// Returns if identifiers are being used.
   bool isUsingIds() const { return usingIds; }
 
   /// Reset the stored identifiers in the space. Enables `usingIds` if it was
   /// `false` before.
-  template <typename T>
   void resetIds() {
     identifiers.clear();
     identifiers.resize(getNumDimAndSymbolVars());
-#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-    idType = TypeID::get<T>();
-#endif
-
     usingIds = true;
   }
 
   /// Disable identifiers being stored in space.
   void disableIds() {
     identifiers.clear();
     usingIds = false;
   }
 
+  /// Check if the spaces are compatible, as well as have the same identifiers
+  /// for each variable.
+  bool isAligned(const PresburgerSpace &other) const;
+  /// Check if the number of variables of the specified kind match, and have
+  /// same identifiers with the other space.
+  bool isAligned(const PresburgerSpace &other, VarKind kind) const;
+
+  void print(llvm::raw_ostream &os) const;
+  void dump() const;
+
 protected:
-  PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
-                  unsigned numSymbols = 0, unsigned numLocals = 0)
+  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols,
+                  unsigned numLocals)
       : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
         numLocals(numLocals) {}
 
-  void *&atId(VarKind kind, unsigned i) {
-    assert(usingIds && "Cannot access identifiers when `usingIds` is false.");
-    assert(kind != VarKind::Local &&
-           "Local variables cannot have identifiers.");
-    return identifiers[getVarKindOffset(kind) + i];
-  }
-
-  void *atId(VarKind kind, unsigned i) const {
-    assert(usingIds && "Cannot access identifiers when `usingIds` is false.");
-    assert(kind != VarKind::Local &&
-           "Local variables cannot have identifiers.");
-    return identifiers[getVarKindOffset(kind) + i];
-  }
-
 private:
+  unsigned &getVarKind(VarKind kind);
+
   // Number of variables corresponding to domain variables.
   unsigned numDomain;
 
   // Number of variables corresponding to range variables.
   unsigned numRange;
 
   /// Number of variables corresponding to symbols (unknown but constant for
   /// analysis).
   unsigned numSymbols;
 
   /// Number of variables corresponding to locals (variables corresponding
   /// to existentially quantified variables).
   unsigned numLocals;
 
   /// Stores whether or not identifiers are being used in this space.
   bool usingIds = false;
 
-#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-  /// TypeID of the identifiers in space. This should be used in asserts only.
-  TypeID idType;
-#endif
-
   /// Stores an identifier for each non-local variable as a `void` pointer.
-  SmallVector<void *, 0> identifiers;
+  SmallVector<Identifier, 0> identifiers;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

 //===- PresburgerSpace.cpp - MLIR PresburgerSpace Class -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include <algorithm>
 #include <cassert>
 
 using namespace mlir;
 using namespace presburger;
 
 PresburgerSpace PresburgerSpace::getDomainSpace() const {
-  // TODO: Preserve identifiers here.
-  return PresburgerSpace::getSetSpace(numDomain, numSymbols, numLocals);
+  PresburgerSpace newSpace = *this;
+  newSpace.removeVarRange(VarKind::Range, 0, getNumRangeVars());
+  newSpace.convertVarKind(VarKind::Domain, 0, getNumDomainVars(),
+                          VarKind::SetDim, 0);
+  return newSpace;
 }
 
 PresburgerSpace PresburgerSpace::getRangeSpace() const {
-  return PresburgerSpace::getSetSpace(numRange, numSymbols, numLocals);
+  PresburgerSpace newSpace = *this;
+  newSpace.removeVarRange(VarKind::Domain, 0, getNumRangeVars());
+  return newSpace;
 }
 
 PresburgerSpace PresburgerSpace::getSpaceWithoutLocals() const {
   PresburgerSpace space = *this;
-  space.removeVarRange(VarKind::Local, 0, numLocals);
+  space.removeVarRange(VarKind::Local, 0, getNumLocalVars());
   return space;
 }
 
 unsigned PresburgerSpace::getNumVarKind(VarKind kind) const {
   if (kind == VarKind::Domain)
     return getNumDomainVars();
   if (kind == VarKind::Range)
     return getNumRangeVars();
   if (kind == VarKind::Symbol)
     return getNumSymbolVars();
   if (kind == VarKind::Local)
-    return numLocals;
+    return getNumLocalVars();
   llvm_unreachable("VarKind does not exist!");
 }
 
 unsigned PresburgerSpace::getVarKindOffset(VarKind kind) const {
   if (kind == VarKind::Domain)
     return 0;
   if (kind == VarKind::Range)
     return getNumDomainVars();
unsigned PresburgerSpace::insertVar(VarKind kind, unsigned pos, unsigned num) {
   else if (kind == VarKind::Symbol)
     numSymbols += num;
   else
     numLocals += num;
 
   // Insert NULL identifiers if `usingIds` and variables inserted are
   // not locals.
   if (usingIds && kind != VarKind::Local)
-    identifiers.insert(identifiers.begin() + absolutePos, num, nullptr);
+    identifiers.insert(identifiers.begin() + absolutePos, num, Identifier());
 
   return absolutePos;
 }
 
 void PresburgerSpace::removeVarRange(VarKind kind, unsigned varStart,
                                      unsigned varLimit) {
   assert(varLimit <= getNumVarKind(kind) && "invalid var limit");
 
void PresburgerSpace::removeVarRange(VarKind kind, unsigned varStart,
 
   // Remove identifiers if `usingIds` and variables removed are not
   // locals.
   if (usingIds && kind != VarKind::Local)
     identifiers.erase(identifiers.begin() + getVarKindOffset(kind) + varStart,
                       identifiers.begin() + getVarKindOffset(kind) + varLimit);
 }
 
+void PresburgerSpace::convertVarKind(VarKind srcKind, unsigned srcPos,
+                                     unsigned num, VarKind dstKind,
+                                     unsigned dstPos) {
+  assert(srcKind != dstKind && "cannot convert variables to the same kind");
+  assert(srcPos + num <= getNumVarKind(srcKind) &&
+         "invalid range for source variables");
+  assert(dstPos <= getNumVarKind(dstKind) &&
+         "invalid position for destination variables");
+
+  auto addVars = [&](VarKind kind, int num) {
+    switch (kind) {
+    case VarKind::Domain:
+      numDomain += num;
+      break;
+    case VarKind::Range:
+      numRange += num;
+      break;
+    case VarKind::Symbol:
+      numSymbols += num;
+      break;
+    case VarKind::Local:
+      numLocals += num;
+      break;
+    }
+  };
+
+  addVars(srcKind, -(signed)num);
+  addVars(dstKind, num);
+
+  // Move identifiers if `usingIds` and variables moved are not locals.
+  unsigned srcOffset = getVarKindOffset(srcKind) + srcPos;
+  unsigned dstOffset = getVarKindOffset(dstKind) + dstPos;
+  if (isUsingIds() && srcKind != VarKind::Local && dstKind != VarKind::Local) {
+    identifiers.insert(identifiers.begin() + dstOffset,
+                       identifiers.begin() + srcOffset,
+                       identifiers.begin() + srcOffset + num);
+    identifiers.erase(identifiers.begin() + srcOffset,
+                      identifiers.begin() + srcOffset + num);
+  } else if (isUsingIds() && srcKind != VarKind::Local) {
+    identifiers.erase(identifiers.begin() + srcOffset,
+                      identifiers.begin() + srcOffset + num);
+  } else if (isUsingIds() && dstKind != VarKind::Local) {
+    identifiers.insert(identifiers.begin() + dstOffset, num, Identifier());
+  }
+}
+
 void PresburgerSpace::swapVar(VarKind kindA, VarKind kindB, unsigned posA,
                               unsigned posB) {
-
-  if (!usingIds)
+  if (!isUsingIds())
     return;
 
   if (kindA == VarKind::Local && kindB == VarKind::Local)
     return;
 
   if (kindA == VarKind::Local) {
-    atId(kindB, posB) = nullptr;
+    getId(kindB, posB) = Identifier();
     return;
   }
 
   if (kindB == VarKind::Local) {
-    atId(kindA, posA) = nullptr;
+    getId(kindA, posA) = Identifier();
     return;
   }
 
-  std::swap(atId(kindA, posA), atId(kindB, posB));
+  std::swap(getId(kindA, posA), getId(kindB, posB));
 }
 
 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();
 }
 
+/// Checks if the number of ids of the given kind in the two spaces are
+/// equal and if the ids are equal. Assumes that both spaces are using
+/// ids.
+static bool areIdsEqual(const PresburgerSpace &spaceA,
+                        const PresburgerSpace &spaceB, VarKind kind) {
+  assert(spaceA.isUsingIds() && spaceB.isUsingIds() &&
+         "Both spaces should be using ids");
+  if (spaceA.getNumVarKind(kind) != spaceB.getNumVarKind(kind))
+    return false;
+  if (kind == VarKind::Local)
+    return true;
+  unsigned numVars = spaceA.getNumVarKind(kind);
+  for (unsigned i = 0; i < numVars; ++i)
+    if (spaceA.getId(kind, i) != spaceB.getId(kind, i))
+      return false;
+  return true;
+}
+
 bool PresburgerSpace::isAligned(const PresburgerSpace &other) const {
-  assert(isUsingIds() && other.isUsingIds() &&
-         "Both spaces should be using identifiers to check for "
-         "alignment.");
-  return isCompatible(other) && identifiers == other.identifiers;
+  // If only one of the spaces is using identifiers, then they are
+  // not aligned.
+  if (isUsingIds() ^ other.isUsingIds())
+    return false;
+  // If both spaces are using identifiers, then they are aligned if
+  // their identifiers are equal.
+  if (isUsingIds())
+    return areIdsEqual(*this, other, VarKind::Domain) &&
+           areIdsEqual(*this, other, VarKind::Range) &&
+           areIdsEqual(*this, other, VarKind::Symbol);
+  // If neither space is using identifiers, then they are aligned if
+  // they are compatible.
+  return isCompatible(other);
 }
 
 bool PresburgerSpace::isAligned(const PresburgerSpace &other,
                                 VarKind kind) const {
-  assert(isUsingIds() && other.isUsingIds() &&
-         "Both spaces should be using identifiers to check for "
-         "alignment.");
-
-  ArrayRef<void *> kindAttachments =
-      ArrayRef(identifiers).slice(getVarKindOffset(kind), getNumVarKind(kind));
-  ArrayRef<void *> otherKindAttachments =
-      ArrayRef(other.identifiers)
-          .slice(other.getVarKindOffset(kind), other.getNumVarKind(kind));
-  return kindAttachments == otherKindAttachments;
+  // If only one of the spaces is using identifiers, then they are
+  // not aligned.
+  if (isUsingIds() ^ other.isUsingIds())
+    return false;
+  // If both spaces are using identifiers, then they are aligned if
+  // their identifiers are equal.
+  if (isUsingIds())
+    return areIdsEqual(*this, other, kind);
+  // If neither space is using identifiers, then they are aligned if
+  // the number of variable kind is equal.
+  return getNumVarKind(kind) == other.getNumVarKind(kind);
 }
 
 void PresburgerSpace::setVarSymbolSeperation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolVars() &&
          "invalid separation position");
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
   // We do not need to change `identifiers` since the ordering of
   // `identifiers` remains same.
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
   os << "Domain: " << getNumDomainVars() << ", "
      << "Range: " << getNumRangeVars() << ", "
      << "Symbols: " << getNumSymbolVars() << ", "
      << "Locals: " << getNumLocalVars() << "\n";
 
-  if (usingIds) {
-#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-    os << "TypeID of identifiers: " << idType.getAsOpaquePointer() << "\n";
-#endif
+  if (isUsingIds()) {
+    auto printIds = [&](VarKind kind) {
+      os << " ";
+      for (unsigned i = 0; i < getNumVarKind(kind); ++i) {
+        Identifier id = getId(kind, i);
+        if (id.hasValue())
+          id.print(os);
+        else
+          os << "None";
+        os << " ";
+      }
+    };
 
     os << "(";
-    for (void *identifier : identifiers)
-      os << identifier << " ";
-    os << ")\n";
+    printIds(VarKind::Domain);
+    os << ") -> (";
+    printIds(VarKind::Range);
+    os << ") : [";
+    printIds(VarKind::Symbol);
+    os << "]";
   }
 }
 
-void PresburgerSpace::dump() const { print(llvm::errs()); }
+void PresburgerSpace::dump() const {
+  print(llvm::errs());
+  llvm::errs() << "\n";
+}

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

TEST(PresburgerSpaceTest, removeIdRange) {
   space.removeVarRange(VarKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainVars(), 1u);
   EXPECT_EQ(space.getNumRangeVars(), 0u);
   EXPECT_EQ(space.getNumSymbolVars(), 2u);
 }
 
 TEST(PresburgerSpaceTest, insertVarIdentifier) {
   PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1, 0);
-  space.resetIds<int *>();
+  space.resetIds();
 
-  // Attach identifiers to domain ids.
   int identifiers[2] = {0, 1};
-  space.setId<int *>(VarKind::Domain, 0, &identifiers[0]);
-  space.setId<int *>(VarKind::Domain, 1, &identifiers[1]);
+
+  // Attach identifiers to domain ids.
+  space.getId(VarKind::Domain, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Domain, 1) = Identifier(&identifiers[1]);
 
   // Try inserting 2 domain ids.
   space.insertVar(VarKind::Domain, 0, 2);
   EXPECT_EQ(space.getNumDomainVars(), 4u);
 
   // Try inserting 1 range ids.
   space.insertVar(VarKind::Range, 0, 1);
   EXPECT_EQ(space.getNumRangeVars(), 3u);
 
   // Check if the identifiers for the old ids are still attached properly.
-  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 2), identifiers[0]);
-  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 3), identifiers[1]);
+  EXPECT_EQ(space.getId(VarKind::Domain, 2), Identifier(&identifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Domain, 3), Identifier(&identifiers[1]));
 }
 
 TEST(PresburgerSpaceTest, removeVarRangeIdentifier) {
   PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0);
-  space.resetIds<int *>();
+  space.resetIds();
 
   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]);
+  space.getId(VarKind::Domain, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Domain, 1) = Identifier(&identifiers[1]);
 
   // Attach identifiers to range identifiers.
-  space.setId<int *>(VarKind::Range, 0, &identifiers[2]);
+  space.getId(VarKind::Range, 0) = Identifier(&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]);
+  space.getId(VarKind::Symbol, 0) = Identifier(&identifiers[3]);
+  space.getId(VarKind::Symbol, 1) = Identifier(&identifiers[4]);
+  space.getId(VarKind::Symbol, 2) = Identifier(&identifiers[5]);
 
   // Remove 1 domain identifier.
   space.removeVarRange(VarKind::Domain, 0, 1);
   EXPECT_EQ(space.getNumDomainVars(), 1u);
 
   // Remove 1 symbol and 1 range identifier.
   space.removeVarRange(VarKind::Symbol, 0, 1);
   space.removeVarRange(VarKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainVars(), 1u);
   EXPECT_EQ(space.getNumRangeVars(), 0u);
   EXPECT_EQ(space.getNumSymbolVars(), 2u);
 
   // Check if domain identifiers are attached properly.
-  EXPECT_EQ(*space.getId<int *>(VarKind::Domain, 0), identifiers[1]);
+  EXPECT_EQ(space.getId(VarKind::Domain, 0), Identifier(&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]);
+  EXPECT_EQ(space.getId(VarKind::Range, 0), Identifier(&identifiers[4]));
+  EXPECT_EQ(space.getId(VarKind::Range, 1), Identifier(&identifiers[5]));
 }