[MLIR] Use a shared uniquer for affine maps and integer sets.

[ulysseB]

Summary

Affine maps and integer sets previously relied on a single lock for creating unique instances. In a multi-threaded setting, this lock becomes a contention point. This commit updates AffineMap and IntegerSet to use StorageUniquer instead. StorageUniquer internally uses sharded locks and thread-local caches to reduce contention. It is already used for affine expressions, types and attributes. On my local machine, this gives me a 5X speedup for an application that manipulates a lot of affine maps and integer sets.

This commit also removes the integer set uniquer threshold. The threshold was used to avoid adding integer sets with a lot of constraints to the hash_map containing unique instances, but the constraints and the integer set were still allocated in the same allocator and never freed, thus not saving any space expect for the hash-map entry.

[rriddle]

Nice cleanup, thanks!

mlir/include/mlir/IR/IntegerSet.h

public:
   }
   static IntegerSet getFromOpaquePointer(const void *pointer) {
     return IntegerSet(
         reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
   }
 
 private:
   ImplType *set;
-  /// Sets with constraints fewer than kUniquingThreshold are uniqued.
-  constexpr static unsigned kUniquingThreshold = 4;
 };
 
 // Make AffineExpr hashable.
 inline ::llvm::hash_code hash_value(IntegerSet arg) {
   return ::llvm::hash_value(arg.set);
 }
 
 } // end namespace mlir

mlir/lib/IR/AffineMapDetail.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef AFFINEMAPDETAIL_H_
 #define AFFINEMAPDETAIL_H_
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/Support/StorageUniquer.h"
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 namespace detail {
 
-struct AffineMapStorage {
+struct AffineMapStorage : public StorageUniquer::BaseStorage {
+  /// The hash key used for uniquing.
+  using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>;
+
   unsigned numDims;
   unsigned numSymbols;
 
   /// The affine expressions for this (multi-dimensional) map.
   /// TODO: use trailing objects for this.
   ArrayRef<AffineExpr> results;
 
   MLIRContext *context;
+
+  bool operator==(const KeyTy &key) const {
+    return std::get<0>(key) == numDims && std::get<1>(key) == numSymbols &&
+           std::get<2>(key) == results;
+  }
+
+  // Constructs an AffineMapStorage from a key. The context must be set by the
+  // caller.
+  static AffineMapStorage *
+  construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) {
+    auto *res = new (allocator.allocate<AffineMapStorage>()) AffineMapStorage();
+    res->numDims = std::get<0>(key);
+    res->numSymbols = std::get<1>(key);
+    res->results = allocator.copyInto(std::get<2>(key));
+    return res;
+  }
 };
 
 } // end namespace detail
 } // end namespace mlir
 
 #endif // AFFINEMAPDETAIL_H_

mlir/lib/IR/IntegerSet.cpp

unsigned IntegerSet::getNumEqualities() const {
   return numEqualities;
 }
 
 unsigned IntegerSet::getNumInequalities() const {
   return getNumConstraints() - getNumEqualities();
 }
 
 bool IntegerSet::isEmptyIntegerSet() const {
-  // This will only work if uniquing is on.
-  static_assert(kUniquingThreshold >= 1,
-                "uniquing threshold should be at least one");
   return *this == getEmptySet(set->dimCount, set->symbolCount, getContext());
 }
 
 ArrayRef<AffineExpr> IntegerSet::getConstraints() const {
   return set->constraints;
 }
 
 AffineExpr IntegerSet::getConstraint(unsigned idx) const {

mlir/lib/IR/IntegerSetDetail.h

 //===- IntegerSetDetail.h - MLIR IntegerSet storage details -----*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This holds implementation details of IntegerSet.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef INTEGERSETDETAIL_H_
 #define INTEGERSETDETAIL_H_
 
 #include "mlir/IR/AffineExpr.h"
+#include "mlir/Support/StorageUniquer.h"
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 namespace detail {
 
-struct IntegerSetStorage {
+struct IntegerSetStorage : public StorageUniquer::BaseStorage {
+  /// The hash key used for uniquing.
+  using KeyTy =
+      std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>, ArrayRef<bool>>;
+
   unsigned dimCount;
   unsigned symbolCount;
 
   /// Array of affine constraints: a constraint is either an equality
   /// (affine_expr == 0) or an inequality (affine_expr >= 0).
   ArrayRef<AffineExpr> constraints;
 
   // Bits to check whether a constraint is an equality or an inequality.
   ArrayRef<bool> eqFlags;
+
+  bool operator==(const KeyTy &key) const {
+    return std::get<0>(key) == dimCount && std::get<1>(key) == symbolCount &&
+           std::get<2>(key) == constraints && std::get<3>(key) == eqFlags;
+  }
+
+  static IntegerSetStorage *
+  construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) {
+    auto *res =
+        new (allocator.allocate<IntegerSetStorage>()) IntegerSetStorage();
+    res->dimCount = std::get<0>(key);
+    res->symbolCount = std::get<1>(key);
+    res->constraints = allocator.copyInto(std::get<2>(key));
+    res->eqFlags = allocator.copyInto(std::get<3>(key));
+    return res;
+  }
 };
 
 } // end namespace detail
 } // end namespace mlir
 #endif // INTEGERSETDETAIL_H_

mlir/lib/IR/MLIRContext.cpp

~ScopedWriterLock() {
     if (mutex)
       mutex->unlock();
   }
   llvm::sys::SmartRWMutex<true> *mutex;
 };
 } // end anonymous namespace.
 
 //===----------------------------------------------------------------------===//
-// AffineMap and IntegerSet hashing
-//===----------------------------------------------------------------------===//
-
-/// A utility function to safely get or create a uniqued instance within the
-/// given set container.
-template <typename ValueT, typename DenseInfoT, typename KeyT,
-          typename ConstructorFn>
-static ValueT safeGetOrCreate(DenseSet<ValueT, DenseInfoT> &container,
-                              KeyT &&key, llvm::sys::SmartRWMutex<true> &mutex,
-                              bool threadingIsEnabled,
-                              ConstructorFn &&constructorFn) {
-  // Check for an existing instance in read-only mode.
-  if (threadingIsEnabled) {
-    llvm::sys::SmartScopedReader<true> instanceLock(mutex);
-    auto it = container.find_as(key);
-    if (it != container.end())
-      return *it;
-  }
-
-  // Acquire a writer-lock so that we can safely create the new instance.
-  ScopedWriterLock instanceLock(mutex, threadingIsEnabled);
-
-  // Check for an existing instance again here, because another writer thread
-  // may have already created one. Otherwise, construct a new instance.
-  auto existing = container.insert_as(ValueT(), key);
-  if (existing.second)
-    return *existing.first = constructorFn();
-  return *existing.first;
-}
-
-namespace {
-struct AffineMapKeyInfo : DenseMapInfo<AffineMap> {
-  // Affine maps are uniqued based on their dim/symbol counts and affine
-  // expressions.
-  using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>;
-  using DenseMapInfo<AffineMap>::isEqual;
-
-  static unsigned getHashValue(const AffineMap &key) {
-    return getHashValue(
-        KeyTy(key.getNumDims(), key.getNumSymbols(), key.getResults()));
-  }
-
-  static unsigned getHashValue(KeyTy key) {
-    return hash_combine(
-        std::get<0>(key), std::get<1>(key),
-        hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()));
-  }
-
-  static bool isEqual(const KeyTy &lhs, AffineMap rhs) {
-    if (rhs == getEmptyKey() || rhs == getTombstoneKey())
-      return false;
-    return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(),
-                                  rhs.getResults());
-  }
-};
-
-struct IntegerSetKeyInfo : DenseMapInfo<IntegerSet> {
-  // Integer sets are uniqued based on their dim/symbol counts, affine
-  // expressions appearing in the LHS of constraints, and eqFlags.
-  using KeyTy =
-      std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>, ArrayRef<bool>>;
-  using DenseMapInfo<IntegerSet>::isEqual;
-
-  static unsigned getHashValue(const IntegerSet &key) {
-    return getHashValue(KeyTy(key.getNumDims(), key.getNumSymbols(),
-                              key.getConstraints(), key.getEqFlags()));
-  }
-
-  static unsigned getHashValue(KeyTy key) {
-    return hash_combine(
-        std::get<0>(key), std::get<1>(key),
-        hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()),
-        hash_combine_range(std::get<3>(key).begin(), std::get<3>(key).end()));
-  }
-
-  static bool isEqual(const KeyTy &lhs, IntegerSet rhs) {
-    if (rhs == getEmptyKey() || rhs == getTombstoneKey())
-      return false;
-    return lhs == std::make_tuple(rhs.getNumDims(), rhs.getNumSymbols(),
-                                  rhs.getConstraints(), rhs.getEqFlags());
-  }
-};
-} // end anonymous namespace.
-
-//===----------------------------------------------------------------------===//
 // MLIRContextImpl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 /// This is the implementation of the MLIRContext class, using the pImpl idiom.
 /// This class is completely private to this file, so everything is public.
 class MLIRContextImpl {
 public:
#endif
 
   /// A mutex used when accessing operation information.
   llvm::sys::SmartRWMutex<true> operationInfoMutex;
 
   //===--------------------------------------------------------------------===//
   // Affine uniquing
   //===--------------------------------------------------------------------===//
 
-  // Affine allocator and mutex for thread safety.
-  llvm::BumpPtrAllocator affineAllocator;
-  llvm::sys::SmartRWMutex<true> affineMutex;
-
-  // Affine map uniquing.
-  using AffineMapSet = DenseSet<AffineMap, AffineMapKeyInfo>;
-  AffineMapSet affineMaps;
-
-  // Integer set uniquing.
-  using IntegerSets = DenseSet<IntegerSet, IntegerSetKeyInfo>;
-  IntegerSets integerSets;
-
-  // Affine expression uniquing.
+  // Affine expression, map and integer set uniquing.
   StorageUniquer affineUniquer;
 
   //===--------------------------------------------------------------------===//
   // Type uniquing
   //===--------------------------------------------------------------------===//
 
   DenseMap<TypeID, AbstractType *> registeredTypes;
   StorageUniquer typeUniquer;
MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
   impl->emptyStringAttr = StringAttr::getEmptyStringAttrUnchecked(this);
 
   // Register the affine storage objects with the uniquer.
   impl->affineUniquer
       .registerParametricStorageType<AffineBinaryOpExprStorage>();
   impl->affineUniquer
       .registerParametricStorageType<AffineConstantExprStorage>();
   impl->affineUniquer.registerParametricStorageType<AffineDimExprStorage>();
+  impl->affineUniquer.registerParametricStorageType<AffineMapStorage>();
+  impl->affineUniquer.registerParametricStorageType<IntegerSetStorage>();
 }
 
 MLIRContext::~MLIRContext() {}
 
 /// Copy the specified array of elements into memory managed by the provided
 /// bump pointer allocator.  This assumes the elements are all PODs.
 template <typename T>
 static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,

 StorageUniquer &MLIRContext::getAffineUniquer() {
   return getImpl().affineUniquer;
 }
 
 AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
                              ArrayRef<AffineExpr> results,
                              MLIRContext *context) {
   auto &impl = context->getImpl();
-  auto key = std::make_tuple(dimCount, symbolCount, results);
-
-  // Safely get or create an AffineMap instance.
-  return safeGetOrCreate(
-      impl.affineMaps, key, impl.affineMutex, impl.threadingIsEnabled, [&] {
-        auto *res = impl.affineAllocator.Allocate<detail::AffineMapStorage>();
-
-        // Copy the results into the bump pointer.
-        results = copyArrayRefInto(impl.affineAllocator, results);
-
-        // Initialize the memory using placement new.
-        new (res)
-            detail::AffineMapStorage{dimCount, symbolCount, results, context};
-        return AffineMap(res);
-      });
+  auto *storage = impl.affineUniquer.get<AffineMapStorage>(
+      [&](AffineMapStorage *storage) { storage->context = context; }, dimCount,
+      symbolCount, results);
+  return AffineMap(storage);
 }
 
 /// Check whether the arguments passed to the AffineMap::get() are consistent.
 /// This method checks whether the highest index of dimensional identifier
 /// present in result expressions is less than `dimCount` and the highest index
 /// of symbolic identifier present in result expressions is less than
 /// `symbolCount`.
 LLVM_NODISCARD static bool willBeValidAffineMap(unsigned dimCount,

 IntegerSet IntegerSet::get(unsigned dimCount, unsigned symbolCount,
                            ArrayRef<AffineExpr> constraints,
                            ArrayRef<bool> eqFlags) {
   // The number of constraints can't be zero.
   assert(!constraints.empty());
   assert(constraints.size() == eqFlags.size());
 
   auto &impl = constraints[0].getContext()->getImpl();
-
-  // A utility function to construct a new IntegerSetStorage instance.
-  auto constructorFn = [&] {
-    auto *res = impl.affineAllocator.Allocate<detail::IntegerSetStorage>();
-
-    // Copy the results and equality flags into the bump pointer.
-    constraints = copyArrayRefInto(impl.affineAllocator, constraints);
-    eqFlags = copyArrayRefInto(impl.affineAllocator, eqFlags);
-
-    // Initialize the memory using placement new.
-    new (res)
-        detail::IntegerSetStorage{dimCount, symbolCount, constraints, eqFlags};
-    return IntegerSet(res);
-  };
-
-  // If this instance is uniqued, then we handle it separately so that multiple
-  // threads may simultaneously access existing instances.
-  if (constraints.size() < IntegerSet::kUniquingThreshold) {
-    auto key = std::make_tuple(dimCount, symbolCount, constraints, eqFlags);
-    return safeGetOrCreate(impl.integerSets, key, impl.affineMutex,
-                           impl.threadingIsEnabled, constructorFn);
-  }
-
-  // Otherwise, acquire a writer-lock so that we can safely create the new
-  // instance.
-  ScopedWriterLock affineLock(impl.affineMutex, impl.threadingIsEnabled);
-  return constructorFn();
+  auto *storage = impl.affineUniquer.get<IntegerSetStorage>(
+      [](IntegerSetStorage *) {}, dimCount, symbolCount, constraints, eqFlags);
+  return IntegerSet(storage);
 }
 
 //===----------------------------------------------------------------------===//
 // StorageUniquerSupport
 //===----------------------------------------------------------------------===//
 
 /// Utility method to generate a callback that can be used to generate a
 /// diagnostic when checking the construction invariants of a storage object.