[mlir][StorageUniquer] Refactor parametric storage to use sharded dense sets

[rriddle]

Summary

This revisions implements sharding in the storage of parametric instances to decrease lock contention by sharding out the allocator/mutex/etc. to use for a specific storage instance based on the hash key. This is a somewhat common approach to reducing lock contention on data structures, and is used by the concurrent hashmaps provided by folly/java/etc. For several compilations tested, this removed all/most lock contention from profiles and reduced compile time by several seconds.

[rriddle]

Right now I've set it so that all types of storage instances use 8 shards, but that isn't that great. Some storage instances have higher concurrency needs, e.g. DictionaryAttr, while others have close to none, e.g. IntegerType. It might be a good idea to expose the level of concurrency expected by specific storage instances, but left that out from the initial patch. Thoughts on this?

[ftynse]

mlir/lib/Support/StorageUniquer.cpp

158

This requires numShards to be a power of 2, worth documenting somewhere + adding an assert. Alternatively, round to the closest power of two before this.

[jpienaar]

mlir/lib/Support/StorageUniquer.cpp

91

Do we need to guard these on when threading is enabled? E.g., have the non threaded case be completely non-sharded and so not even create an array?

132

I don't follow this, why do we not care about which shard for arbitrary mutation?

Also does this need to be public?

[rriddle]

Resolve comments

mlir/lib/Support/StorageUniquer.cpp

91

Updated to wrap all of the complicated bits within an if LLVM_ENABLE_THREADS so that we don't incur any complexity when threading is disabled at build time.

132

I don't follow this, why do we not care about which shard for arbitrary mutation?

Updated to find the shard that allocated the storage instance.

Also does this need to be public?

Yes, but it is fine because everything here is local to the StorageUniquer.cpp file. We could make it private, but only by moving it to be nested within the StorageUniquerImpl class.

158

Added an assert for now given that this is still local to this file. If we expose concurrency to users, I suspect we wouldn't expose exact shard numbers but have something akin to concurrency levels.

[rriddle]

Thanks for the review!

[jpienaar]

Thanks!

mlir/include/mlir/Support/StorageUniquer.h

public:
     /// Allocate an instance of the provided type.
     template <typename T> T *allocate() { return allocator.Allocate<T>(); }
 
     /// Allocate 'size' bytes of 'alignment' aligned memory.
     void *allocate(size_t size, size_t alignment) {
       return allocator.Allocate(size, alignment);
     }
 
+    /// Returns true if this allocator allocated the provided object pointer.
+    bool allocated(const void *ptr) {
+      return allocator.identifyObject(ptr).hasValue();
+    }
+
   private:
     /// The raw allocator for type storage objects.
     llvm::BumpPtrAllocator allocator;
   };
 
   StorageUniquer();
   ~StorageUniquer();
 
public:
   /// Changes the mutable component of 'storage' by forwarding the trailing
   /// arguments to the 'mutate' function of the derived class.
   template <typename Storage, typename... Args>
   LogicalResult mutate(TypeID id, Storage *storage, Args &&...args) {
     auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
       return static_cast<Storage &>(*storage).mutate(
           allocator, std::forward<Args>(args)...);
     };
-    return mutateImpl(id, mutationFn);
+    return mutateImpl(id, storage, mutationFn);
   }
 
 private:
   /// Implementation for getting/creating an instance of a derived type with
   /// parametric storage.
   BaseStorage *getParametricStorageTypeImpl(
       TypeID id, unsigned hashValue,
       function_ref<bool(const BaseStorage *)> isEqual,
private:
   /// Implementation for registering an instance of a derived type with default
   /// storage.
   void
   registerSingletonImpl(TypeID id,
                         function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for mutating an instance of a derived storage.
   LogicalResult
-  mutateImpl(TypeID id,
+  mutateImpl(TypeID id, BaseStorage *storage,
              function_ref<LogicalResult(StorageAllocator &)> mutationFn);
 
   /// The internal implementation class.
   std::unique_ptr<detail::StorageUniquerImpl> impl;
 
   //===--------------------------------------------------------------------===//
   // Key Construction
   //===--------------------------------------------------------------------===//

mlir/lib/Support/StorageUniquer.cpp

 using namespace mlir::detail;
 
 namespace {
 /// This class represents a uniquer for storage instances of a specific type
 /// that has parametric storage. It contains all of the necessary data to unique
 /// storage instances in a thread safe way. This allows for the main uniquer to
 /// bucket each of the individual sub-types removing the need to lock the main
 /// uniquer itself.
-struct ParametricStorageUniquer {
+class ParametricStorageUniquer {
+public:
   using BaseStorage = StorageUniquer::BaseStorage;
   using StorageAllocator = StorageUniquer::StorageAllocator;
 
   /// A lookup key for derived instances of storage objects.
   struct LookupKey {
     /// The known hash value of the key.
     unsigned hashValue;
 
     /// An equality function for comparing with an existing storage instance.
     function_ref<bool(const BaseStorage *)> isEqual;
   };
 
+private:
   /// A utility wrapper object representing a hashed storage object. This class
   /// contains a storage object and an existing computed hash value.
   struct HashedStorage {
     HashedStorage(unsigned hashValue = 0, BaseStorage *storage = nullptr)
         : hashValue(hashValue), storage(storage) {}
     unsigned hashValue;
     BaseStorage *storage;
   };
struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
     }
     static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
       if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
         return false;
       // Invoke the equality function on the lookup key.
       return lhs.isEqual(rhs.storage);
     }
   };
+  using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
 
+  /// This class represents a single shard of the uniquer. The uniquer uses a
+  /// set of shards to allow for multiple threads to create instances with less
+  /// lock contention.
+  struct Shard {
     /// The set containing the allocated storage instances.
-  using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
     StorageTypeSet instances;
 
+    /// Allocator to use when constructing derived instances.
+    StorageAllocator allocator;
+
+#if LLVM_ENABLE_THREADS != 0
+    /// A mutex to keep uniquing thread-safe.
+    llvm::sys::SmartRWMutex<true> mutex;
+#endif
+  };
+
+  /// Get or create an instance of a param derived type in an thread-unsafe
+  /// fashion.
+  BaseStorage *
+  getOrCreateUnsafe(Shard &shard, LookupKey &key,
+                    function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
+    auto existing = shard.instances.insert_as({key.hashValue}, key);
+    BaseStorage *&storage = existing.first->storage;
+    if (existing.second)
+      storage = ctorFn(shard.allocator);
+    return storage;
+  }
+
+public:
+#if LLVM_ENABLE_THREADS != 0
+  /// Initialize the storage uniquer with a given number of storage shards to
+  /// use. The provided shard number is required to be a valid power of 2.
+  ParametricStorageUniquer(size_t numShards = 8)
+      : shards(new std::atomic<Shard *>[numShards]), numShards(numShards) {
+    assert(llvm::isPowerOf2_64(numShards) &&
+           "the number of shards is required to be a power of 2");
+    for (size_t i = 0; i < numShards; i++)
+      shards[i].store(nullptr, std::memory_order_relaxed);
+  }
+  ~ParametricStorageUniquer() {
+    // Free all of the allocated shards.
+    for (size_t i = 0; i != numShards; ++i)
+      if (Shard *shard = shards[i].load())
+        delete shard;
+  }
+  /// Get or create an instance of a parametric type.
+  BaseStorage *
+  getOrCreate(bool threadingIsEnabled, unsigned hashValue,
+              function_ref<bool(const BaseStorage *)> isEqual,
+              function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
+    Shard &shard = getShard(hashValue);
+    ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
+    if (!threadingIsEnabled)
+      return getOrCreateUnsafe(shard, lookupKey, ctorFn);
+
+    // Check for a instance of this object in the local cache.
+    auto localIt = localCache->insert_as({hashValue}, lookupKey);
+    BaseStorage *&localInst = localIt.first->storage;
+    if (localInst)
+      return localInst;
+
+    // Check for an existing instance in read-only mode.
+    {
+      llvm::sys::SmartScopedReader<true> typeLock(shard.mutex);
+      auto it = shard.instances.find_as(lookupKey);
+      if (it != shard.instances.end())
+        return localInst = it->storage;
+    }
+
+    // Acquire a writer-lock so that we can safely create the new storage
+    // instance.
+    llvm::sys::SmartScopedWriter<true> typeLock(shard.mutex);
+    return localInst = getOrCreateUnsafe(shard, lookupKey, ctorFn);
+  }
+  /// Run a mutation function on the provided storage object in a thread-safe
+  /// way.
+  LogicalResult
+  mutate(bool threadingIsEnabled, BaseStorage *storage,
+         function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+    Shard &shard = getShardFor(storage);
+    if (!threadingIsEnabled)
+      return mutationFn(shard.allocator);
+
+    llvm::sys::SmartScopedWriter<true> lock(shard.mutex);
+    return mutationFn(shard.allocator);
+  }
+
+private:
+  /// Return the shard used for the given hash value.
+  Shard &getShard(unsigned hashValue) {
+    // Get a shard number from the provided hashvalue.
+    unsigned shardNum = hashValue & (numShards - 1);
+
+    // Try to acquire an already initialized shard.
+    Shard *shard = shards[shardNum].load(std::memory_order_acquire);
+    if (shard)
+      return *shard;
+
+    // Otherwise, try to allocate a new shard.
+    Shard *newShard = new Shard();
+    if (shards[shardNum].compare_exchange_strong(shard, newShard))
+      return *newShard;
+
+    // If one was allocated before we can initialize ours, delete ours.
+    delete newShard;
+    return *shard;
+  }
+
+  /// Return the shard that allocated the provided storage object.
+  Shard &getShardFor(BaseStorage *storage) {
+    for (size_t i = 0; i != numShards; ++i) {
+      if (Shard *shard = shards[i].load(std::memory_order_acquire)) {
+        llvm::sys::SmartScopedReader<true> lock(shard->mutex);
+        if (shard->allocator.allocated(storage))
+          return *shard;
+      }
+    }
+    llvm_unreachable("expected storage object to have a valid shard");
+  }
+
   /// A thread local cache for storage objects. This helps to reduce the lock
   /// contention when an object already existing in the cache.
   ThreadLocalCache<StorageTypeSet> localCache;
 
-  /// Allocator to use when constructing derived instances.
-  StorageAllocator allocator;
+  /// A set of uniquer shards to allow for further bucketing accesses for
+  /// instances of this storage type. Each shard is lazily initialized to reduce
+  /// the overhead when only a small amount of shards are in use.
+  std::unique_ptr<std::atomic<Shard *>[]> shards;
+
+  /// The number of available shards.
+  size_t numShards;
+
+#else
+  /// If multi-threading is disabled, ignore the shard parameter as we will
+  /// always use one shard.
+  ParametricStorageUniquer(size_t numShards = 0) {}
 
-  /// A mutex to keep type uniquing thread-safe.
-  llvm::sys::SmartRWMutex<true> mutex;
+  /// Get or create an instance of a parametric type.
+  BaseStorage *
+  getOrCreate(bool threadingIsEnabled, unsigned hashValue,
+              function_ref<bool(const BaseStorage *)> isEqual,
+              function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
+    ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
+    return getOrCreateUnsafe(shard, lookupKey, ctorFn);
+  }
+  /// Run a mutation function on the provided storage object in a thread-safe
+  /// way.
+  LogicalResult
+  mutate(bool threadingIsEnabled, BaseStorage *storage,
+         function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+    return mutationFn(shard.allocator);
+  }
+
+private:
+  /// The main uniquer shard that is used for allocating storage instances.
+  Shard shard;
+#endif
 };
 } // end anonymous namespace
 
 namespace mlir {
 namespace detail {
 /// This is the implementation of the StorageUniquer class.
 struct StorageUniquerImpl {
   using BaseStorage = StorageUniquer::BaseStorage;
   using StorageAllocator = StorageUniquer::StorageAllocator;
 
   //===--------------------------------------------------------------------===//
   // Parametric Storage
   //===--------------------------------------------------------------------===//
 
   /// Check if an instance of a parametric storage class exists.
   bool hasParametricStorage(TypeID id) { return parametricUniquers.count(id); }
 
   /// Get or create an instance of a parametric type.
   BaseStorage *
   getOrCreate(TypeID id, unsigned hashValue,
               function_ref<bool(const BaseStorage *)> isEqual,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     assert(parametricUniquers.count(id) &&
            "creating unregistered storage instance");
-    ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
-    if (!threadingIsEnabled)
-      return getOrCreateUnsafe(storageUniquer, lookupKey, ctorFn);
-
-    // Check for a instance of this object in the local cache.
-    auto localIt = storageUniquer.localCache->insert_as({hashValue}, lookupKey);
-    BaseStorage *&localInst = localIt.first->storage;
-    if (localInst)
-      return localInst;
-
-    // Check for an existing instance in read-only mode.
-    {
-      llvm::sys::SmartScopedReader<true> typeLock(storageUniquer.mutex);
-      auto it = storageUniquer.instances.find_as(lookupKey);
-      if (it != storageUniquer.instances.end())
-        return localInst = it->storage;
-    }
+    return storageUniquer.getOrCreate(threadingIsEnabled, hashValue, isEqual,
+                                      ctorFn);
 
-    // Acquire a writer-lock so that we can safely create the new type instance.
-    llvm::sys::SmartScopedWriter<true> typeLock(storageUniquer.mutex);
-    return localInst = getOrCreateUnsafe(storageUniquer, lookupKey, ctorFn);
-  }
-  /// Get or create an instance of a param derived type in an thread-unsafe
-  /// fashion.
-  BaseStorage *
-  getOrCreateUnsafe(ParametricStorageUniquer &storageUniquer,
-                    ParametricStorageUniquer::LookupKey &key,
-                    function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
-    auto existing = storageUniquer.instances.insert_as({key.hashValue}, key);
-    if (!existing.second)
-      return existing.first->storage;
-
-    // Otherwise, construct and initialize the derived storage for this type
-    // instance.
-    BaseStorage *storage = ctorFn(storageUniquer.allocator);
-    *existing.first =
-        ParametricStorageUniquer::HashedStorage{key.hashValue, storage};
-    return storage;
-  }
-
-  /// Mutates an instance of a derived storage in a thread-safe way.
+  /// Run a mutation function on the provided storage object in a thread-safe
+  /// way.
   LogicalResult
-  mutate(TypeID id,
+  mutate(TypeID id, BaseStorage *storage,
          function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
     assert(parametricUniquers.count(id) &&
            "mutating unregistered storage instance");
     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
-    if (!threadingIsEnabled)
-      return mutationFn(storageUniquer.allocator);
-
-    llvm::sys::SmartScopedWriter<true> lock(storageUniquer.mutex);
-    return mutationFn(storageUniquer.allocator);
+    return storageUniquer.mutate(threadingIsEnabled, storage, mutationFn);
   }
 
   //===--------------------------------------------------------------------===//
   // Singleton Storage
   //===--------------------------------------------------------------------===//
 
   /// Get or create an instance of a singleton storage class.
   BaseStorage *getSingleton(TypeID id) {
     BaseStorage *singletonInstance = singletonInstances[id];
     assert(singletonInstance && "expected singleton instance to exist");
     return singletonInstance;
   }
 
   /// Check if an instance of a singleton storage class exists.
-  bool hasSingleton(TypeID id) { return singletonInstances.count(id); }
+  bool hasSingleton(TypeID id) const { return singletonInstances.count(id); }
 
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
   /// Map of type ids to the storage uniquer to use for registered objects.
   DenseMap<TypeID, std::unique_ptr<ParametricStorageUniquer>>
       parametricUniquers;
void StorageUniquer::registerSingletonImpl(
     TypeID id, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
   assert(!impl->singletonInstances.count(id) &&
          "storage class already registered");
   impl->singletonInstances.try_emplace(id, ctorFn(impl->singletonAllocator));
 }
 
 /// Implementation for mutating an instance of a derived storage.
 LogicalResult StorageUniquer::mutateImpl(
-    TypeID id, function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
-  return impl->mutate(id, mutationFn);
+    TypeID id, BaseStorage *storage,
+    function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+  return impl->mutate(id, storage, mutationFn);
 }