| //===- StorageUniquer.cpp - Common Storage Class Uniquer ------------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Support/StorageUniquer.h" |
| |
| #include "mlir/Support/LLVM.h" |
| #include "llvm/Support/RWMutex.h" |
| |
| using namespace mlir; |
| using namespace mlir::detail; |
| |
| namespace mlir { |
| namespace detail { |
| /// This is the implementation of the StorageUniquer class. |
| struct StorageUniquerImpl { |
| using BaseStorage = StorageUniquer::BaseStorage; |
| using StorageAllocator = StorageUniquer::StorageAllocator; |
| |
| /// A lookup key for derived instances of storage objects. |
| struct LookupKey { |
| /// The known derived kind for the storage. |
| unsigned kind; |
| |
| /// 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; |
| }; |
| |
| /// A utility wrapper object representing a hashed storage object. This class |
| /// contains a storage object and an existing computed hash value. |
| struct HashedStorage { |
| unsigned hashValue; |
| BaseStorage *storage; |
| }; |
| |
| /// Get or create an instance of a complex derived type. |
| BaseStorage * |
| getOrCreate(unsigned kind, unsigned hashValue, |
| function_ref<bool(const BaseStorage *)> isEqual, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) { |
| LookupKey lookupKey{kind, hashValue, isEqual}; |
| if (!threadingIsEnabled) |
| return getOrCreateUnsafe(kind, hashValue, lookupKey, ctorFn); |
| |
| // Check for an existing instance in read-only mode. |
| { |
| llvm::sys::SmartScopedReader<true> typeLock(mutex); |
| auto it = storageTypes.find_as(lookupKey); |
| if (it != storageTypes.end()) |
| return it->storage; |
| } |
| |
| // Acquire a writer-lock so that we can safely create the new type instance. |
| llvm::sys::SmartScopedWriter<true> typeLock(mutex); |
| return getOrCreateUnsafe(kind, hashValue, lookupKey, ctorFn); |
| } |
| /// Get or create an instance of a complex derived type in an unsafe fashion. |
| BaseStorage * |
| getOrCreateUnsafe(unsigned kind, unsigned hashValue, LookupKey &lookupKey, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) { |
| auto existing = storageTypes.insert_as({}, lookupKey); |
| if (!existing.second) |
| return existing.first->storage; |
| |
| // Otherwise, construct and initialize the derived storage for this type |
| // instance. |
| BaseStorage *storage = initializeStorage(kind, ctorFn); |
| *existing.first = HashedStorage{hashValue, storage}; |
| return storage; |
| } |
| |
| /// Get or create an instance of a simple derived type. |
| BaseStorage * |
| getOrCreate(unsigned kind, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) { |
| if (!threadingIsEnabled) |
| return getOrCreateUnsafe(kind, ctorFn); |
| |
| // Check for an existing instance in read-only mode. |
| { |
| llvm::sys::SmartScopedReader<true> typeLock(mutex); |
| auto it = simpleTypes.find(kind); |
| if (it != simpleTypes.end()) |
| return it->second; |
| } |
| |
| // Acquire a writer-lock so that we can safely create the new type instance. |
| llvm::sys::SmartScopedWriter<true> typeLock(mutex); |
| return getOrCreateUnsafe(kind, ctorFn); |
| } |
| /// Get or create an instance of a simple derived type in an unsafe fashion. |
| BaseStorage * |
| getOrCreateUnsafe(unsigned kind, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) { |
| auto &result = simpleTypes[kind]; |
| if (result) |
| return result; |
| |
| // Otherwise, create and return a new storage instance. |
| return result = initializeStorage(kind, ctorFn); |
| } |
| |
| /// Erase an instance of a complex derived type. |
| void erase(unsigned kind, unsigned hashValue, |
| function_ref<bool(const BaseStorage *)> isEqual, |
| function_ref<void(BaseStorage *)> cleanupFn) { |
| LookupKey lookupKey{kind, hashValue, isEqual}; |
| |
| // Acquire a writer-lock so that we can safely erase the type instance. |
| llvm::sys::SmartScopedWriter<true> typeLock(mutex); |
| auto existing = storageTypes.find_as(lookupKey); |
| if (existing == storageTypes.end()) |
| return; |
| |
| // Cleanup the storage and remove it from the map. |
| cleanupFn(existing->storage); |
| storageTypes.erase(existing); |
| } |
| |
| //===--------------------------------------------------------------------===// |
| // Instance Storage |
| //===--------------------------------------------------------------------===// |
| |
| /// Utility to create and initialize a storage instance. |
| BaseStorage * |
| initializeStorage(unsigned kind, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) { |
| BaseStorage *storage = ctorFn(allocator); |
| storage->kind = kind; |
| return storage; |
| } |
| |
| /// Storage info for derived TypeStorage objects. |
| struct StorageKeyInfo : DenseMapInfo<HashedStorage> { |
| static HashedStorage getEmptyKey() { |
| return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()}; |
| } |
| static HashedStorage getTombstoneKey() { |
| return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()}; |
| } |
| |
| static unsigned getHashValue(const HashedStorage &key) { |
| return key.hashValue; |
| } |
| static unsigned getHashValue(LookupKey key) { return key.hashValue; } |
| |
| static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) { |
| return lhs.storage == rhs.storage; |
| } |
| static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) { |
| if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey())) |
| return false; |
| // If the lookup kind matches the kind of the storage, then invoke the |
| // equality function on the lookup key. |
| return lhs.kind == rhs.storage->getKind() && lhs.isEqual(rhs.storage); |
| } |
| }; |
| |
| /// Unique types with specific hashing or storage constraints. |
| using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>; |
| StorageTypeSet storageTypes; |
| |
| /// Unique types with just the kind. |
| DenseMap<unsigned, BaseStorage *> simpleTypes; |
| |
| /// Allocator to use when constructing derived type instances. |
| StorageUniquer::StorageAllocator allocator; |
| |
| /// A mutex to keep type uniquing thread-safe. |
| llvm::sys::SmartRWMutex<true> mutex; |
| |
| /// Flag specifying if multi-threading is enabled within the uniquer. |
| bool threadingIsEnabled = true; |
| }; |
| } // end namespace detail |
| } // namespace mlir |
| |
| StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {} |
| StorageUniquer::~StorageUniquer() {} |
| |
| /// Set the flag specifying if multi-threading is disabled within the uniquer. |
| void StorageUniquer::disableMultithreading(bool disable) { |
| impl->threadingIsEnabled = !disable; |
| } |
| |
| /// Implementation for getting/creating an instance of a derived type with |
| /// complex storage. |
| auto StorageUniquer::getImpl( |
| unsigned kind, unsigned hashValue, |
| function_ref<bool(const BaseStorage *)> isEqual, |
| function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * { |
| return impl->getOrCreate(kind, hashValue, isEqual, ctorFn); |
| } |
| |
| /// Implementation for getting/creating an instance of a derived type with |
| /// default storage. |
| auto StorageUniquer::getImpl( |
| unsigned kind, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) |
| -> BaseStorage * { |
| return impl->getOrCreate(kind, ctorFn); |
| } |
| |
| /// Implementation for erasing an instance of a derived type with complex |
| /// storage. |
| void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue, |
| function_ref<bool(const BaseStorage *)> isEqual, |
| function_ref<void(BaseStorage *)> cleanupFn) { |
| impl->erase(kind, hashValue, isEqual, cleanupFn); |
| } |