diff --git a/mlir/include/mlir/Support/StorageUniquer.h b/mlir/include/mlir/Support/StorageUniquer.h --- a/mlir/include/mlir/Support/StorageUniquer.h +++ b/mlir/include/mlir/Support/StorageUniquer.h @@ -115,6 +115,11 @@ 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; @@ -227,7 +232,7 @@ return static_cast(*storage).mutate( allocator, std::forward(args)...); }; - return mutateImpl(id, mutationFn); + return mutateImpl(id, storage, mutationFn); } private: @@ -254,7 +259,7 @@ /// Implementation for mutating an instance of a derived storage. LogicalResult - mutateImpl(TypeID id, + mutateImpl(TypeID id, BaseStorage *storage, function_ref mutationFn); /// The internal implementation class. 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,165 @@ 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; + +#if LLVM_ENABLE_THREADS != 0 + /// A mutex to keep uniquing thread-safe. + llvm::sys::SmartRWMutex 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 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[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 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 storage + // instance. + llvm::sys::SmartScopedWriter 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 mutationFn) { + Shard &shard = getShardFor(storage); + if (!threadingIsEnabled) + return mutationFn(shard.allocator); + + llvm::sys::SmartScopedWriter 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 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 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; + + /// 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) {} + + /// Get or create an instance of a parametric type. + BaseStorage * + getOrCreate(bool threadingIsEnabled, unsigned hashValue, + function_ref isEqual, + function_ref 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 mutationFn) { + return mutationFn(shard.allocator); + } - /// A mutex to keep type uniquing thread-safe. - llvm::sys::SmartRWMutex mutex; +private: + /// The main uniquer shard that is used for allocating storage instances. + Shard shard; +#endif }; } // end anonymous namespace @@ -106,59 +253,20 @@ 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); + return storageUniquer.getOrCreate(threadingIsEnabled, hashValue, isEqual, + 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; - } - - /// 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 mutationFn) { 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, storage, mutationFn); } //===--------------------------------------------------------------------===// @@ -173,7 +281,7 @@ } /// 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 @@ -247,6 +355,7 @@ /// Implementation for mutating an instance of a derived storage. LogicalResult StorageUniquer::mutateImpl( - TypeID id, function_ref mutationFn) { - return impl->mutate(id, mutationFn); + TypeID id, BaseStorage *storage, + function_ref mutationFn) { + return impl->mutate(id, storage, mutationFn); }