diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp --- a/mlir/lib/Support/StorageUniquer.cpp +++ b/mlir/lib/Support/StorageUniquer.cpp @@ -22,7 +22,8 @@ /// 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; @@ -35,6 +36,7 @@ function_ref 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 { @@ -68,20 +70,119 @@ return lhs.isEqual(rhs.storage); } }; - - /// The set containing the allocated storage instances. using StorageTypeSet = DenseSet; - StorageTypeSet instances; + + /// 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; + + /// A mutex to keep uniquing thread-safe. + llvm::sys::SmartRWMutex mutex; + }; + +public: + ParametricStorageUniquer(size_t numShards = 8) + : shards(new std::atomic[numShards] { nullptr }), + numShards(numShards) {} + ~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 isEqual, + function_ref 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 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 type instance. + llvm::sys::SmartScopedWriter typeLock(shard.mutex); + return localInst = getOrCreateUnsafe(shard, lookupKey, ctorFn); + } + /// Mutates an instance of a derived storage in a thread-safe way. + LogicalResult + mutate(bool threadingIsEnabled, + function_ref mutationFn) { + // All we care about for mutation is that we have an allocator and mutex, + // so the shard here doesn't matter. + Shard &shard = getShard(0); + if (!threadingIsEnabled) + return mutationFn(shard.allocator); + + llvm::sys::SmartScopedWriter lock(shard.mutex); + return mutationFn(shard.allocator); + } + +private: + /// Get or create an instance of a param derived type in an thread-unsafe + /// fashion. + BaseStorage * + getOrCreateUnsafe(Shard &shard, LookupKey &key, + function_ref ctorFn) { + auto existing = shard.instances.insert_as({key.hashValue}, key); + BaseStorage *&storage = existing.first->storage; + if (existing.second) + storage = ctorFn(shard.allocator); + return storage; + } + + /// 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; + } /// A thread local cache for storage objects. This helps to reduce the lock /// contention when an object already existing in the cache. ThreadLocalCache 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[]> shards; - /// A mutex to keep type uniquing thread-safe. - llvm::sys::SmartRWMutex mutex; + /// The number of available shards. + size_t numShards; }; } // end anonymous namespace @@ -106,45 +207,9 @@ function_ref 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 typeLock(storageUniquer.mutex); - auto it = storageUniquer.instances.find_as(lookupKey); - if (it != storageUniquer.instances.end()) - return localInst = it->storage; - } - - // Acquire a writer-lock so that we can safely create the new type instance. - llvm::sys::SmartScopedWriter 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 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; + return storageUniquer.getOrCreate(threadingIsEnabled, hashValue, isEqual, + ctorFn); } /// Mutates an instance of a derived storage in a thread-safe way. @@ -154,11 +219,7 @@ assert(parametricUniquers.count(id) && "mutating unregistered storage instance"); ParametricStorageUniquer &storageUniquer = *parametricUniquers[id]; - if (!threadingIsEnabled) - return mutationFn(storageUniquer.allocator); - - llvm::sys::SmartScopedWriter lock(storageUniquer.mutex); - return mutationFn(storageUniquer.allocator); + return storageUniquer.mutate(threadingIsEnabled, mutationFn); } //===--------------------------------------------------------------------===//