[Mlir-commits] [mlir] 67f52f3 - [mlir][StorageUniquer] Refactor parametric storage to use sharded dense sets
River Riddle
llvmlistbot at llvm.org
Mon Oct 26 19:40:41 PDT 2020
Author: River Riddle
Date: 2020-10-26T19:40:19-07:00
New Revision: 67f52f35d62b25f929646e972287c7b5397a044e
URL: https://github.com/llvm/llvm-project/commit/67f52f35d62b25f929646e972287c7b5397a044e
DIFF: https://github.com/llvm/llvm-project/commit/67f52f35d62b25f929646e972287c7b5397a044e.diff
LOG: [mlir][StorageUniquer] Refactor parametric storage to use sharded dense sets
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.
Differential Revision: https://reviews.llvm.org/D89659
Added:
Modified:
mlir/include/mlir/Support/StorageUniquer.h
mlir/lib/Support/StorageUniquer.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Support/StorageUniquer.h b/mlir/include/mlir/Support/StorageUniquer.h
index 36cfaf6359b3d..fc7ffa74f3b5e 100644
--- a/mlir/include/mlir/Support/StorageUniquer.h
+++ b/mlir/include/mlir/Support/StorageUniquer.h
@@ -115,6 +115,11 @@ class StorageUniquer {
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 @@ class StorageUniquer {
return static_cast<Storage &>(*storage).mutate(
allocator, std::forward<Args>(args)...);
};
- return mutateImpl(id, mutationFn);
+ return mutateImpl(id, storage, mutationFn);
}
private:
@@ -254,7 +259,7 @@ class StorageUniquer {
/// 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.
diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp
index 8e0ef6b8f2765..7a802430f010c 100644
--- a/mlir/lib/Support/StorageUniquer.cpp
+++ b/mlir/lib/Support/StorageUniquer.cpp
@@ -22,7 +22,8 @@ namespace {
/// 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 @@ struct ParametricStorageUniquer {
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 {
@@ -68,20 +70,165 @@ struct ParametricStorageUniquer {
return lhs.isEqual(rhs.storage);
}
};
-
- /// The set containing the allocated storage instances.
using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
- 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<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) {}
+
+ /// 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);
+ }
- /// A mutex to keep type uniquing thread-safe.
- llvm::sys::SmartRWMutex<true> mutex;
+private:
+ /// The main uniquer shard that is used for allocating storage instances.
+ Shard shard;
+#endif
};
} // end anonymous namespace
@@ -106,59 +253,20 @@ struct StorageUniquerImpl {
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;
- }
-
- // 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);
+ 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<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);
}
//===--------------------------------------------------------------------===//
@@ -173,7 +281,7 @@ struct StorageUniquerImpl {
}
/// 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 @@ void StorageUniquer::registerSingletonImpl(
/// 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);
}
More information about the Mlir-commits
mailing list