[mlir][StorageUniquer] Use allocators per thread instead of per shard

[rriddle]

Summary

This greatly reduces the number of allocators we create, while still
retaining thread safety. Reducing the number of allocators is much
better for locality and memory usage; this revision drops memory
usage for some MLIR heavy workloads (with lots of attributes/types)
by >=5%. This is due to the observation that the number of threads
is effectively always smaller than the number of parametric attributes/types.

[mehdi_amini]

This assumes less threads than shards overall right?

[rriddle]

In D145991#4191347, @mehdi_amini wrote:

This assumes less threads than shards overall right?

Yeah, which is effectively always a guarantee (even on my beefy desktop) given the number of parametric builtin attributes/types we have.

[mehdi_amini]

mlir/lib/Support/StorageUniquer.cpp

176

Can you get the shard after this if?

[mehdi_amini]

In D145991#4191352, @rriddle wrote:

In D145991#4191347, @mehdi_amini wrote:

This assumes less threads than shards overall right?

Yeah, which is effectively always a guarantee (even on my beefy desktop) given the number of parametric builtin attributes/types we have.

Makes sense, maybe worth a note in the commit message to call it out explicitly.

[rriddle]

mlir/lib/Support/StorageUniquer.cpp

176

Nice catch, thanks!

[Mogball]

amazing

mlir/lib/Support/StorageUniquer.cpp

private:
 
   /// 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.
     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) {
+                    function_ref<BaseStorage *(StorageAllocator &)> ctorFn,
+                    function_ref<StorageAllocator &()> getAllocFn) {
     auto existing = shard.instances.insert_as({key.hashValue}, key);
     BaseStorage *&storage = existing.first->storage;
     if (existing.second)
-      storage = ctorFn(shard.allocator);
+      storage = ctorFn(getAllocFn());
     return storage;
   }
 
   /// Destroy all of the storage instances within the given shard.
   void destroyShardInstances(Shard &shard) {
     if (!destructorFn)
       return;
     for (HashedStorage &instance : shard.instances)
for (size_t i = 0; i != numShards; ++i) {
         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) {
+              function_ref<BaseStorage *(StorageAllocator &)> ctorFn,
+              function_ref<StorageAllocator &()> getAllocFn) {
     Shard &shard = getShard(hashValue);
     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
     if (!threadingIsEnabled)
-      return getOrCreateUnsafe(shard, lookupKey, ctorFn);
+      return getOrCreateUnsafe(shard, lookupKey, ctorFn, getAllocFn);
 
     // 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);
+    return localInst = getOrCreateUnsafe(shard, lookupKey, ctorFn, getAllocFn);
   }
+
   /// Run a mutation function on the provided storage object in a thread-safe
   /// way.
   LogicalResult
-  mutate(bool threadingIsEnabled, BaseStorage *storage,
+  mutate(bool threadingIsEnabled, StorageAllocator &allocator,
+         BaseStorage *storage,
          function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
-    Shard &shard = getShardFor(storage);
+    // Get a shard to use for mutating this storage instance. It doesn't need to
+    // be the same shard as the original allocation, but does need to be
+    // deterministic.
+    Shard &shard = getShard(llvm::hash_value(storage));
     if (!threadingIsEnabled)
-      return mutationFn(shard.allocator);
+      return mutationFn(allocator);
 
     llvm::sys::SmartScopedWriter<true> lock(shard.mutex);
-    return mutationFn(shard.allocator);
+    return mutationFn(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;
 
   /// 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;
struct StorageUniquerImpl {
   /// 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 &storageUniquer = *parametricUniquers[id];
-    return storageUniquer.getOrCreate(threadingIsEnabled, hashValue, isEqual,
-                                      ctorFn);
+    return storageUniquer.getOrCreate(
+        threadingIsEnabled, hashValue, isEqual, ctorFn,
+        [&]() -> StorageAllocator & { return getThreadSafeAllocator(); });
   }
 
   /// Run a mutation function on the provided storage object in a thread-safe
   /// way.
   LogicalResult
   mutate(TypeID id, BaseStorage *storage,
          function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
     assert(parametricUniquers.count(id) &&
            "mutating unregistered storage instance");
     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
-    return storageUniquer.mutate(threadingIsEnabled, storage, mutationFn);
+    return storageUniquer.mutate(threadingIsEnabled, getThreadSafeAllocator(),
+                                 storage, mutationFn);
+  }
+
+  /// Return an allocator that can be used to safely allocate instances on the
+  /// current thread.
+  StorageAllocator &getThreadSafeAllocator() {
+#if LLVM_ENABLE_THREADS != 0
+    StorageAllocator *&allocator = threadSafeAllocator.get();
+
+    // If the allocator has not been initialized, create a new one.
+    if (!allocator) {
+      allocator = new StorageAllocator();
+
+      // Record this allocator, given that we don't want it to be destroyed when
+      // the thread dies.
+      llvm::sys::SmartScopedLock<true> lock(threadAllocatorMutex);
+      threadAllocators.push_back(std::unique_ptr<StorageAllocator>(allocator));
+    }
+
+    return *allocator;
+#else
+    return allocator;
+#endif
   }
 
   //===--------------------------------------------------------------------===//
   // 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) const { return singletonInstances.count(id); }
 
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
+#if LLVM_ENABLE_THREADS != 0
+  /// A thread local set of allocators used for uniquing parametric instances,
+  /// or other data allocated in thread volatile situations.
+  ThreadLocalCache<StorageAllocator *> threadSafeAllocator;
+
+  /// All of the allocators that have been created for thread based allocation.
+  std::vector<std::unique_ptr<StorageAllocator>> threadAllocators;
+
+  /// A mutex used for safely adding a new thread allocator.
+  llvm::sys::SmartMutex<true> threadAllocatorMutex;
+#endif
+
   /// Map of type ids to the storage uniquer to use for registered objects.
   DenseMap<TypeID, std::unique_ptr<ParametricStorageUniquer>>
       parametricUniquers;
 
   /// Map of type ids to a singleton instance when the storage class is a
   /// singleton.
   DenseMap<TypeID, BaseStorage *> singletonInstances;
 
-  /// Allocator used for uniquing singleton instances.
-  StorageAllocator singletonAllocator;
+  /// Main allocator used for uniquing singleton instances, and other state when
+  /// thread safety is guaranteed.
+  StorageAllocator allocator;
 
   /// Flag specifying if multi-threading is enabled within the uniquer.
   bool threadingIsEnabled = true;
 };
 } // namespace detail
 } // namespace mlir
 
 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}

 }
 
 /// Implementation for registering an instance of a derived type with default
 /// storage.
 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));
+  impl->singletonInstances.try_emplace(id, ctorFn(impl->allocator));
 }
 
 /// Implementation for mutating an instance of a derived storage.
 LogicalResult StorageUniquer::mutateImpl(
     TypeID id, BaseStorage *storage,
     function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
   return impl->mutate(id, storage, mutationFn);
 }