[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

156–167

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

192–193

I'd rather have this return some kDimNotFound constant.

194

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).

213

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

236–237

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.

239

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.

245

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

161

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?

192–193

Is something like using kDimNotFound = INT_MAX acceptable?

236–237

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.

239

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.

192–193

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

204

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

236–237

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.

239

Yes, absoutely.

245

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).

236–237

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

236–237

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 values to each variable
+/// in space. `resetValues<ValueType>` enables attaching values to space. All
+/// values must be of same type, which is the type passed as template through
+/// `resetValues<ValueType>`. `ValueType` must be castable to `void*` using
+/// `llvm::PointerLikeTypeTraits`.
 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:
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   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;
 
+  //===--------------------------------------------------------------------===//
+  //     Value Interactions
+  //===--------------------------------------------------------------------===//
+
+  void *&atValue(IdKind kind, unsigned i) {
+    assert(usingValues && "Cannot access values when `usingValues` is false.");
+    assert(kind != IdKind::Local &&
+           "Values cannot be attached to local identifiers.");
+    return values[getIdKindOffset(kind) + i];
+  }
+  void *atValue(IdKind kind, unsigned i) const {
+    assert(usingValues && "Cannot access values when `usingValues` is false.");
+    assert(kind != IdKind::Local &&
+           "Values cannot be attached to local identifiers.");
+    return values[getIdKindOffset(kind) + i];
+  }
+
+  /// Set the value attached to the `i^th` variable to `value`. `T` here should
+  /// match the type used to enable values.
+  template <typename T>
+  void setValue(IdKind kind, unsigned i, T value) {
+    assert(TypeID::get<T>() == valueType && "Type mismatch");
+    atValue(kind, i) = llvm::PointerLikeTypeTraits<T>::getAsVoidPointer(value);
+  }
+
+  /// Get the value attached to the `i^th` variable casted to type `T`. `T` here
+  /// should match the type used to enable values.
+  template <typename T>
+  T getValue(IdKind kind, unsigned i) const {
+    assert(TypeID::get<T>() == valueType && "Type mismatch");
+    return llvm::PointerLikeTypeTraits<T>::getFromVoidPointer(atValue(kind, i));
+  }
+
+  /// Check if the spaces are compatible, as well as have the same values
+  /// attached to each variable.
+  bool isAligned(const PresburgerSpace &other) const;
+  /// Check if the number of variables of the specified kind match, and have
+  /// same values with the other space.
+  bool isAligned(const PresburgerSpace &other, IdKind kind) const;
+
+  /// Find the variable of the specified kind with value `val`. Returns number
+  /// of variables of the specified kind if not found.
+  unsigned findId(IdKind kind, void *val) const;
+
+  /// Returns if values are being used.
+  bool isUsingValues() const { return usingValues; }
+
+  /// Reset the stored values in the space. Enables `usingValues` if it was
+  /// `false` before.
+  template <typename T>
+  void resetValues() {
+    values.clear();
+    values.resize(getNumDimAndSymbolIds());
+#ifndef NDEBUG
+    valueType = TypeID::get<T>();
+#endif
+
+    usingValues = true;
+  }
+
+  /// Disable values being stored in space.
+  void disableValues() { usingValues = false; }
+
 protected:
   PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
                   unsigned numSymbols = 0, unsigned numLocals = 0)
       : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
         numLocals(numLocals) {}
 
 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 values are attached to this space.
+  bool usingValues = false;
+
+#ifndef NDEBUG
+  /// TypeID of the values in space. This should be used in asserts only.
+  TypeID valueType;
+#endif
+
+  /// Stores a value for each non-local identifier as a `void` pointer.
+  llvm::SmallVector<void *> values;
 };
 
 } // 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 values if `usingValues` and variables inserted are not locals.
+  if (usingValues && kind != IdKind::Local)
+    values.insert(values.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 values if `usingValues` and variables removed are not locals.
+  if (usingValues && kind != IdKind::Local)
+    values.erase(values.begin() + getIdKindOffset(kind) + idStart,
+                 values.begin() + getIdKindOffset(kind) + idLimit);
+}
+
+void PresburgerSpace::swapId(IdKind kindA, IdKind kindB, unsigned posA,
+                             unsigned posB) {
+
+  if (!usingValues)
+    return;
+
+  if (kindA == IdKind::Local && kindB == IdKind::Local)
+    return;
+
+  if (kindA == IdKind::Local) {
+    atValue(kindB, posB) = nullptr;
+    return;
+  }
+
+  if (kindB == IdKind::Local) {
+    atValue(kindA, posA) = nullptr;
+    return;
+  }
+
+  std::swap(atValue(kindA, posA), atValue(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(isUsingValues() && other.isUsingValues() &&
+         "Both spaces should be using values to check for "
+         "alignment.");
+  return isCompatible(other) && values == other.values;
+}
+
+bool PresburgerSpace::isAligned(const PresburgerSpace &other,
+                                IdKind kind) const {
+  assert(isUsingValues() && other.isUsingValues() &&
+         "Both spaces should be using values to check for "
+         "alignment.");
+
+  ArrayRef<void *> kindValues = {values.data() + getIdKindOffset(kind),
+                                 values.data() + getIdKindEnd(kind)};
+  ArrayRef<void *> otherKindValues = {
+      other.values.data() + other.getIdKindOffset(kind),
+      other.values.data() + other.getIdKindEnd(kind)};
+
+  return kindValues == otherKindValues;
+}
+
 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
+  // We do not need to change `values` since the ordering of `values` remains
+  // same.
+}
+
+unsigned PresburgerSpace::findId(IdKind kind, void *val) const {
+  assert(usingValues && "Values are not being used in space.");
+  unsigned i = 0;
+  for (unsigned e = getNumIdKind(kind); i < e; ++i)
+    if (atValue(kind, i) != nullptr && atValue(kind, i) == val)
+      break;
+  return i;
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
   os << "Domain: " << getNumDomainIds() << ", "
      << "Range: " << getNumRangeIds() << ", "
      << "Symbols: " << getNumSymbolIds() << ", "
      << "Locals: " << getNumLocalIds() << "\n";
+
+  if (usingValues) {
+#ifndef NDEBUG
+    os << "TypeID of values: " << valueType.getAsOpaquePointer() << "\n";
+#endif
+
+    os << "(";
+    for (void *value : values)
+      os << value << " ";
+    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, insertIdValue) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1, 0);
+  space.resetValues<int *>();
+
+  // Attach value to domain ids.
+  int values[2] = {0, 1};
+  space.setValue<int *>(IdKind::Domain, 0, &values[0]);
+  space.setValue<int *>(IdKind::Domain, 1, &values[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 values for the old ids are still attached properly.
+  EXPECT_EQ(*space.getValue<int *>(IdKind::Domain, 2), values[0]);
+  EXPECT_EQ(*space.getValue<int *>(IdKind::Domain, 3), values[1]);
+}
+
+TEST(PresburgerSpaceTest, removeIdRangeValue) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3, 0);
+  space.resetValues<int *>();
+
+  int values[6] = {0, 1, 2, 3, 4, 5};
+
+  // Attach values to domain identifiers.
+  space.setValue<int *>(IdKind::Domain, 0, &values[0]);
+  space.setValue<int *>(IdKind::Domain, 1, &values[1]);
+
+  // Attach values to range identifiers.
+  space.setValue<int *>(IdKind::Range, 0, &values[2]);
+
+  // Attach values to symbol identifiers.
+  space.setValue<int *>(IdKind::Symbol, 0, &values[3]);
+  space.setValue<int *>(IdKind::Symbol, 1, &values[4]);
+  space.setValue<int *>(IdKind::Symbol, 2, &values[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 values are attached properly.
+  EXPECT_EQ(*space.getValue<int *>(IdKind::Domain, 0), values[1]);
+
+  // Check if symbol values are attached properly.
+  EXPECT_EQ(*space.getValue<int *>(IdKind::Range, 0), values[4]);
+  EXPECT_EQ(*space.getValue<int *>(IdKind::Range, 1), values[5]);
+}