[MLIR][Presburger] Add values to PresburgerSpace

[Groverkss]

Summary

This patch allows attaching user information, called "values" to each
identifier. The values are used to carry information along with variables and
are also used to determine if two variables are identical.

This patch is part of a series of patches to allow attaching user information
with variables in Presburger library.

[ftynse]

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

69

Nit: I'm trying to think of a better name than "values" to avoid potential confusion with mlir::Value. Something generic like info or data.

70

Btw, is there any expectations that values are unique for all dimensions?

71

It is worth mentioning that ValueType should be connected to MLIR's TypeID mechanism.

72–73

163–174

These look like they should be private, we don't want the give the caller access without type safety.

199–200

I'd rather have this return some kDimNotFound constant.

201

Let's template this method too, no need to show users void * (and many value types that wrap pointers aren't implicitly convertible to pointers).

220

Should this also clear the list of values, just in case?

269–270

Have you considered alternative storage schemes that don't require using 4 bytes (because of alignment) on each object to store a flag? One could steal a bit from one of the ints above and provide accessor functions with bit manipulation, for example.

272

This should be #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS (https://llvm.org/docs/ProgrammersManual.html#abi-breaking-checks) because it changes the size of the object that may be passed around in function arguments.

278

We likely want SmallVector<void *, 0> so we don't stack-allocate values for 6 pointers when they are not used.

Also nit: don't prefix SmallVector with llvm::.

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

163

makeArrayRef(values).slice(getIdKindOffset(kind), getNumIdKind(kind)) looks more readable

[Groverkss]

Address Alex's comments

[Groverkss]

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

69

The idea I had was to call all dimensions/symbols/locals as "variables" and the values as the "identifiers" since they "identify" variables. But I will be fine with anything else also like info or data.

70

There are no such requirements in the space. But practically, they should be unique so that two variables can be distinguished in a single space.

71

AFAIU mlir::TypeID has a fallback in case the ValueType is not connected to mlir::TypeID and should work for all types. So it is not required to be connected to mlir::TypeID right?

199–200

Is something like using kDimNotFound = INT_MAX acceptable?

269–270

We could use a flags variable to store these flags, with each bit corresponding to a different flag. I'm not sure if the stealing bit approach is as scalable in case we have to add more flags.

272

The reason I see this flag is used is so that the user can have an ABI compatible build between with asserts and without asserts. If a user disables LLVM_ENABLE_ABI_BREAKING_CHECKS but keeps LLVM_ENABLE_ASSERTIONS enabled, it will fail to compile, since we use valueType in asserts. Should we also guard the asserts with this macro?

[ftynse]

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

69

I am fine with any suggestion, avoiding potential confusion is the most important part.

70

Aren't dimensions distinguished by IdKind + position? This makes it sounds like values are required to distinguish dimensions, but it is not really the case, is it?

71

The fallback is based on names and doesn't guarantee actual uniqueness. It is only relevant in case of shared libraries so may be a bit of a stretch, but it's better to at least tell the user that MLIR's TypeID mechanism is being used so they know how to debug.

199–200

UINT_MAX probably, or static_cast<unsigned>(-1) sounds fine.

211

This now also needs LLVM_ENABLE_ABI_BREAKING_CHECKS since the variable definition is guarded by this flag...

269–270

Are you planning to add more flags? If so, how many?

In general, I would say we expect numDomain + numRange + numSymbols + numLocals <= UNIT_MAX, otherwise we would not be able to reason about the total dimensionality of the space. Under this assumption, we can safely use 2 bits (take one bit from two of the values above) and be sure we are not losing any representational capacity.

272

Yes, absoutely.

278

llvm:: is unnecessary

[Groverkss]

• Address most comments (2 remaining)

[Groverkss]

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

70

Right, I think I should have been more clear. What I mean was, let's say we are comparing two spaces and we are looking to align an identifier at the i^th position of some kind in space1 to an identifier in space2 of the same kind. We need the identifiers in space2 to be unique, so we have only one matching identifier in space2. This is why, practically, you would want the identifiers to be unique.

[ftynse]

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

70

Makes sense. I suggest documenting this in the class documentation and in the relevant functions + adding the assertion in the function that values are indeed unique. May also be worth adding a documentation hint to the functions that add dimensions (and set values to null) so the user doesn't forget to add the values later on (a stricter alternative is to take the value or a list thereof as an optional extra parameter).

[Groverkss]

• Address remaining comments

[Groverkss]

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

69

Changed it to` attachments` for now. I'll send a bigger patch after this for naming since it would affect the whole Presburger library (variables/identifiers).

269–270

There is no current plan to add any flags, however, if we ever decide to move to nested spaces, named spaces, etc. we could need more flags.

For now, I would like to avoid stealing a bit from these ints since I'm not sure of the impact of having to extract the sub-ranges of the ints every time we want the variable. I can put this as a TODO, and we can decide this after benchmarking with and without this change. Would that be acceptable?

[ftynse]

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

269–270

There is no current plan to add any flags, however, if we ever decide to move to nested spaces, named spaces, etc. we could need more flags.

When we do, we can reconsider. This sounds rather distant in the future so sounds like premature "optimization".

For now, I would like to avoid stealing a bit from these ints since I'm not sure of the impact of having to extract the sub-ranges of the ints every time we want the variable. I can put this as a TODO, and we can decide this after benchmarking with and without this change. Would that be acceptable?

This will cost one trivial integer arithmetic operation (right shift or bitwise "and" assuming you are stealing the least significant bit), pretty much the cheapest operation imaginable. Orders of magnitude less than a memory access. We use this trick (via PointerIntPair), e.g., to represent the small result number in mlir::OpResult (most common mlir::Value) without any visible performance problems, among many other places in LLVM. But sure, if you want to benchmark, I am happy to see the numbers.

[Groverkss]

Fix tests

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

 //===- PresburgerSpace.h - MLIR PresburgerSpace Class -----------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Classes representing space information like number of identifiers and kind of
 // identifiers.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 #define MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 
+#include "mlir/Support/TypeID.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 /// Kind of identifier. Implementation wise SetDims are treated as Range
 /// ids, and spaces with no distinction between dimension ids are treated
 /// as relations with zero domain ids.

 ///
 /// A space with no distinction between types of dimension identifiers can
 /// be implemented as a space with zero domain. IdKind::SetDim should be used
 /// to refer to dimensions in such spaces.
 ///
 /// Compatibility of two spaces implies that number of identifiers of each kind
 /// other than Locals are equal. Equality of two spaces implies that number of
 /// identifiers of each kind are equal.
+///
+/// PresburgerSpace optionally also supports attaching attachements to each
+/// variable in space. `resetAttachements<AttachementType>` enables attaching
+/// attachements to space. All attachements must be of the same type,
+/// `AttachementType`. `AttachementType` must have a
+/// `llvm::PointerLikeTypeTraits` specialization available and should be
+/// supported via mlir::TypeID.
+///
+/// These attachements can be used to check if two variables in two different
+/// spaces correspond to the same variable.
 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:
 
   /// Return the IdKind of the id at the specified position.
   IdKind getIdKindAt(unsigned pos) const;
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
+  ///
+  /// If attachements are being used, the newly added variables have no
+  /// attachements.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
 
   /// Removes identifiers of the specified kind in the column range [idStart,
   /// idLimit). The range is relative to the kind of identifier.
   void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
+  /// Swaps the posA^th variable of kindA and posB^th variable of kindB.
+  void swapId(IdKind kindA, IdKind kindB, unsigned posA, unsigned posB);
+
   /// Returns true if both the spaces are compatible i.e. if both spaces have
   /// the same number of identifiers of each kind (excluding locals).
   bool isCompatible(const PresburgerSpace &other) const;
 
   /// Returns true if both the spaces are equal including local identifiers i.e.
   /// if both spaces have the same number of identifiers 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 identifiers immediately before
   /// the split become symbols, or some of the symbols immediately after the
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 
+  //===--------------------------------------------------------------------===//
+  //     Attachement Interactions
+  //===--------------------------------------------------------------------===//
+
+  /// Set the attachement for `i^th` variable to `attachement`. `T` here should
+  /// match the type used to enable attachements.
+  template <typename T>
+  void setAttachement(IdKind kind, unsigned i, T attachement) {
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    assert(TypeID::get<T>() == attachementType && "Type mismatch");
+#endif
+    atAttachement(kind, i) =
+        llvm::PointerLikeTypeTraits<T>::getAsVoidPointer(attachement);
+  }
+
+  /// Get the attachement for `i^th` variable casted to type `T`. `T` here
+  /// should match the type used to enable attachements.
+  template <typename T>
+  T getAttachement(IdKind kind, unsigned i) const {
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    assert(TypeID::get<T>() == attachementType && "Type mismatch");
+#endif
+    return llvm::PointerLikeTypeTraits<T>::getFromVoidPointer(
+        atAttachement(kind, i));
+  }
+
+  /// Check if the i^th variable of the specified kind has a non-null
+  /// attachement.
+  bool hasAttachement(IdKind kind, unsigned i) const {
+    return atAttachement(kind, i) != nullptr;
+  }
+
+  /// Check if the spaces are compatible, as well as have the same attachements
+  /// for each variable.
+  bool isAligned(const PresburgerSpace &other) const;
+  /// Check if the number of variables of the specified kind match, and have
+  /// same attachements with the other space.
+  bool isAligned(const PresburgerSpace &other, IdKind kind) const;
+
+  /// Find the variable of the specified kind with attachement `val`.
+  /// PresburgerSpace::kIdNotFound if attachement is not found.
+  template <typename T>
+  unsigned findId(IdKind kind, T val) const {
+    unsigned i = 0;
+    for (unsigned e = getNumIdKind(kind); i < e; ++i)
+      if (hasAttachement(kind, i) && getAttachement<T>(kind, i) == val)
+        return i;
+    return kIdNotFound;
+  }
+  static const unsigned kIdNotFound = UINT_MAX;
+
+  /// Returns if attachements are being used.
+  bool isUsingAttachements() const { return usingAttachements; }
+
+  /// Reset the stored attachements in the space. Enables `usingAttachements` if
+  /// it was `false` before.
+  template <typename T>
+  void resetAttachements() {
+    attachements.clear();
+    attachements.resize(getNumDimAndSymbolIds());
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    attachementType = TypeID::get<T>();
+#endif
+
+    usingAttachements = true;
+  }
+
+  /// Disable attachements being stored in space.
+  void disableAttachements() {
+    attachements.clear();
+    usingAttachements = false;
+  }
+
 protected:
   PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
                   unsigned numSymbols = 0, unsigned numLocals = 0)
       : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
         numLocals(numLocals) {}
 
+  void *&atAttachement(IdKind kind, unsigned i) {
+    assert(usingAttachements &&
+           "Cannot access attachements when `usingAttachements` is false.");
+    assert(kind != IdKind::Local &&
+           "Local variables cannot have attachements.");
+    return attachements[getIdKindOffset(kind) + i];
+  }
+
+  void *atAttachement(IdKind kind, unsigned i) const {
+    assert(usingAttachements &&
+           "Cannot access attachements when `usingAttachements` is false.");
+    assert(kind != IdKind::Local &&
+           "Local variables cannot have attachements.");
+    return attachements[getIdKindOffset(kind) + i];
+  }
+
 private:
   // Number of identifiers corresponding to domain identifiers.
   unsigned numDomain;
 
   // Number of identifiers corresponding to range identifiers.
   unsigned numRange;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
   /// analysis).
   unsigned numSymbols;
 
   /// Number of identifers corresponding to locals (identifiers corresponding
   /// to existentially quantified variables).
   unsigned numLocals;
+
+  /// Stores whether or not attachements are being used in this space.
+  bool usingAttachements = false;
+
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+  /// TypeID of the attachements in space. This should be used in asserts only.
+  TypeID attachementType;
+#endif
+
+  /// Stores a attachement for each non-local identifier as a `void` pointer.
+  SmallVector<void *, 0> attachements;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

if (kind == IdKind::Domain)
     numDomain += num;
   else if (kind == IdKind::Range)
     numRange += num;
   else if (kind == IdKind::Symbol)
     numSymbols += num;
   else
     numLocals += num;
 
+  // Insert NULL attachements if `usingAttachements` and variables inserted are
+  // not locals.
+  if (usingAttachements && kind != IdKind::Local)
+    attachements.insert(attachements.begin() + absolutePos, num, nullptr);
+
   return absolutePos;
 }
 
 void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
                                     unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   unsigned numIdsEliminated = idLimit - idStart;
   if (kind == IdKind::Domain)
     numDomain -= numIdsEliminated;
   else if (kind == IdKind::Range)
     numRange -= numIdsEliminated;
   else if (kind == IdKind::Symbol)
     numSymbols -= numIdsEliminated;
   else
     numLocals -= numIdsEliminated;
+
+  // Remove attachements if `usingAttachements` and variables removed are not
+  // locals.
+  if (usingAttachements && kind != IdKind::Local)
+    attachements.erase(attachements.begin() + getIdKindOffset(kind) + idStart,
+                       attachements.begin() + getIdKindOffset(kind) + idLimit);
+}
+
+void PresburgerSpace::swapId(IdKind kindA, IdKind kindB, unsigned posA,
+                             unsigned posB) {
+
+  if (!usingAttachements)
+    return;
+
+  if (kindA == IdKind::Local && kindB == IdKind::Local)
+    return;
+
+  if (kindA == IdKind::Local) {
+    atAttachement(kindB, posB) = nullptr;
+    return;
+  }
+
+  if (kindB == IdKind::Local) {
+    atAttachement(kindA, posA) = nullptr;
+    return;
+  }
+
+  std::swap(atAttachement(kindA, posA), atAttachement(kindB, posB));
 }
 
 bool PresburgerSpace::isCompatible(const PresburgerSpace &other) const {
   return getNumDomainIds() == other.getNumDomainIds() &&
          getNumRangeIds() == other.getNumRangeIds() &&
          getNumSymbolIds() == other.getNumSymbolIds();
 }
 
 bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
   return isCompatible(other) && getNumLocalIds() == other.getNumLocalIds();
 }
 
+bool PresburgerSpace::isAligned(const PresburgerSpace &other) const {
+  assert(isUsingAttachements() && other.isUsingAttachements() &&
+         "Both spaces should be using attachements to check for "
+         "alignment.");
+  return isCompatible(other) && attachements == other.attachements;
+}
+
+bool PresburgerSpace::isAligned(const PresburgerSpace &other,
+                                IdKind kind) const {
+  assert(isUsingAttachements() && other.isUsingAttachements() &&
+         "Both spaces should be using attachements to check for "
+         "alignment.");
+
+  ArrayRef<void *> kindAttachements =
+      makeArrayRef(attachements)
+          .slice(getIdKindOffset(kind), getNumIdKind(kind));
+  ArrayRef<void *> otherKindAttachements =
+      makeArrayRef(other.attachements)
+          .slice(other.getIdKindOffset(kind), other.getNumIdKind(kind));
+  return kindAttachements == otherKindAttachements;
+}
+
 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
+  // We do not need to change `attachements` since the ordering of
+  // `attachements` remains same.
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
   os << "Domain: " << getNumDomainIds() << ", "
      << "Range: " << getNumRangeIds() << ", "
      << "Symbols: " << getNumSymbolIds() << ", "
      << "Locals: " << getNumLocalIds() << "\n";
+
+  if (usingAttachements) {
+#ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
+    os << "TypeID of attachements: " << attachementType.getAsOpaquePointer()
+       << "\n";
+#endif
+
+    os << "(";
+    for (void *attachement : attachements)
+      os << attachement << " ";
+    os << ")\n";
+  }
 }
 
 void PresburgerSpace::dump() const { print(llvm::errs()); }

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

TEST(PresburgerSpaceTest, removeIdRange) {
 
   // Remove 1 symbol and 1 range identifier.
   space.removeIdRange(IdKind::Symbol, 0, 1);
   space.removeIdRange(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
   EXPECT_EQ(space.getNumRangeIds(), 0u);
   EXPECT_EQ(space.getNumSymbolIds(), 2u);
 }
+
+TEST(PresburgerSpaceTest, insertIdAttachement) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1, 0);
+  space.resetAttachements<int *>();
+
+  // Attach attachement to domain ids.
+  int attachements[2] = {0, 1};
+  space.setAttachement<int *>(IdKind::Domain, 0, &attachements[0]);
+  space.setAttachement<int *>(IdKind::Domain, 1, &attachements[1]);
+
+  // Try inserting 2 domain ids.
+  space.insertId(IdKind::Domain, 0, 2);
+  EXPECT_EQ(space.getNumDomainIds(), 4u);
+
+  // Try inserting 1 range ids.
+  space.insertId(IdKind::Range, 0, 1);
+  EXPECT_EQ(space.getNumRangeIds(), 3u);
+
+  // Check if the attachements for the old ids are still attached properly.
+  EXPECT_EQ(*space.getAttachement<int *>(IdKind::Domain, 2), attachements[0]);
+  EXPECT_EQ(*space.getAttachement<int *>(IdKind::Domain, 3), attachements[1]);
+}
+
+TEST(PresburgerSpaceTest, removeIdRangeAttachement) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0);
+  space.resetAttachements<int *>();
+
+  int attachements[6] = {0, 1, 2, 3, 4, 5};
+
+  // Attach attachements to domain identifiers.
+  space.setAttachement<int *>(IdKind::Domain, 0, &attachements[0]);
+  space.setAttachement<int *>(IdKind::Domain, 1, &attachements[1]);
+
+  // Attach attachements to range identifiers.
+  space.setAttachement<int *>(IdKind::Range, 0, &attachements[2]);
+
+  // Attach attachements to symbol identifiers.
+  space.setAttachement<int *>(IdKind::Symbol, 0, &attachements[3]);
+  space.setAttachement<int *>(IdKind::Symbol, 1, &attachements[4]);
+  space.setAttachement<int *>(IdKind::Symbol, 2, &attachements[5]);
+
+  // Remove 1 domain identifier.
+  space.removeIdRange(IdKind::Domain, 0, 1);
+  EXPECT_EQ(space.getNumDomainIds(), 1u);
+
+  // Remove 1 symbol and 1 range identifier.
+  space.removeIdRange(IdKind::Symbol, 0, 1);
+  space.removeIdRange(IdKind::Range, 0, 1);
+  EXPECT_EQ(space.getNumDomainIds(), 1u);
+  EXPECT_EQ(space.getNumRangeIds(), 0u);
+  EXPECT_EQ(space.getNumSymbolIds(), 2u);
+
+  // Check if domain attachements are attached properly.
+  EXPECT_EQ(*space.getAttachement<int *>(IdKind::Domain, 0), attachements[1]);
+
+  // Check if symbol attachements are attached properly.
+  EXPECT_EQ(*space.getAttachement<int *>(IdKind::Range, 0), attachements[4]);
+  EXPECT_EQ(*space.getAttachement<int *>(IdKind::Range, 1), attachements[5]);
+}