1 //===- StorageUniquer.cpp - Common Storage Class Uniquer ------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Support/StorageUniquer.h"
10 
11 #include "mlir/Support/LLVM.h"
12 #include "mlir/Support/TypeID.h"
13 #include "llvm/Support/RWMutex.h"
14 
15 using namespace mlir;
16 using namespace mlir::detail;
17 
18 namespace {
19 /// This class represents a uniquer for storage instances of a specific type
20 /// that has parametric storage. It contains all of the necessary data to unique
21 /// storage instances in a thread safe way. This allows for the main uniquer to
22 /// bucket each of the individual sub-types removing the need to lock the main
23 /// uniquer itself.
24 struct ParametricStorageUniquer {
25   using BaseStorage = StorageUniquer::BaseStorage;
26   using StorageAllocator = StorageUniquer::StorageAllocator;
27 
28   /// A lookup key for derived instances of storage objects.
29   struct LookupKey {
30     /// The known hash value of the key.
31     unsigned hashValue;
32 
33     /// An equality function for comparing with an existing storage instance.
34     function_ref<bool(const BaseStorage *)> isEqual;
35   };
36 
37   /// A utility wrapper object representing a hashed storage object. This class
38   /// contains a storage object and an existing computed hash value.
39   struct HashedStorage {
40     unsigned hashValue;
41     BaseStorage *storage;
42   };
43 
44   /// Storage info for derived TypeStorage objects.
45   struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
46     static HashedStorage getEmptyKey() {
47       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()};
48     }
49     static HashedStorage getTombstoneKey() {
50       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()};
51     }
52 
53     static unsigned getHashValue(const HashedStorage &key) {
54       return key.hashValue;
55     }
56     static unsigned getHashValue(LookupKey key) { return key.hashValue; }
57 
58     static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) {
59       return lhs.storage == rhs.storage;
60     }
61     static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
62       if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
63         return false;
64       // Invoke the equality function on the lookup key.
65       return lhs.isEqual(rhs.storage);
66     }
67   };
68 
69   /// The set containing the allocated storage instances.
70   using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
71   StorageTypeSet instances;
72 
73   /// Allocator to use when constructing derived instances.
74   StorageAllocator allocator;
75 
76   /// A mutex to keep type uniquing thread-safe.
77   llvm::sys::SmartRWMutex<true> mutex;
78 };
79 } // end anonymous namespace
80 
81 namespace mlir {
82 namespace detail {
83 /// This is the implementation of the StorageUniquer class.
84 struct StorageUniquerImpl {
85   using BaseStorage = StorageUniquer::BaseStorage;
86   using StorageAllocator = StorageUniquer::StorageAllocator;
87 
88   //===--------------------------------------------------------------------===//
89   // Parametric Storage
90   //===--------------------------------------------------------------------===//
91 
92   /// Check if an instance of a parametric storage class exists.
93   bool hasParametricStorage(TypeID id) { return parametricUniquers.count(id); }
94 
95   /// Get or create an instance of a parametric type.
96   BaseStorage *
97   getOrCreate(TypeID id, unsigned hashValue,
98               function_ref<bool(const BaseStorage *)> isEqual,
99               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
100     assert(parametricUniquers.count(id) &&
101            "creating unregistered storage instance");
102     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
103     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
104     if (!threadingIsEnabled)
105       return getOrCreateUnsafe(storageUniquer, lookupKey, ctorFn);
106 
107     // Check for an existing instance in read-only mode.
108     {
109       llvm::sys::SmartScopedReader<true> typeLock(storageUniquer.mutex);
110       auto it = storageUniquer.instances.find_as(lookupKey);
111       if (it != storageUniquer.instances.end())
112         return it->storage;
113     }
114 
115     // Acquire a writer-lock so that we can safely create the new type instance.
116     llvm::sys::SmartScopedWriter<true> typeLock(storageUniquer.mutex);
117     return getOrCreateUnsafe(storageUniquer, lookupKey, ctorFn);
118   }
119   /// Get or create an instance of a complex derived type in an thread-unsafe
120   /// fashion.
121   BaseStorage *
122   getOrCreateUnsafe(ParametricStorageUniquer &storageUniquer,
123                     ParametricStorageUniquer::LookupKey &lookupKey,
124                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
125     auto existing = storageUniquer.instances.insert_as({}, lookupKey);
126     if (!existing.second)
127       return existing.first->storage;
128 
129     // Otherwise, construct and initialize the derived storage for this type
130     // instance.
131     BaseStorage *storage = ctorFn(storageUniquer.allocator);
132     *existing.first =
133         ParametricStorageUniquer::HashedStorage{lookupKey.hashValue, storage};
134     return storage;
135   }
136 
137   /// Erase an instance of a parametric derived type.
138   void erase(TypeID id, unsigned hashValue,
139              function_ref<bool(const BaseStorage *)> isEqual,
140              function_ref<void(BaseStorage *)> cleanupFn) {
141     assert(parametricUniquers.count(id) &&
142            "erasing unregistered storage instance");
143     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
144     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
145 
146     // Acquire a writer-lock so that we can safely erase the type instance.
147     llvm::sys::SmartScopedWriter<true> lock(storageUniquer.mutex);
148     auto existing = storageUniquer.instances.find_as(lookupKey);
149     if (existing == storageUniquer.instances.end())
150       return;
151 
152     // Cleanup the storage and remove it from the map.
153     cleanupFn(existing->storage);
154     storageUniquer.instances.erase(existing);
155   }
156 
157   /// Mutates an instance of a derived storage in a thread-safe way.
158   LogicalResult
159   mutate(TypeID id,
160          function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
161     assert(parametricUniquers.count(id) &&
162            "mutating unregistered storage instance");
163     ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
164     if (!threadingIsEnabled)
165       return mutationFn(storageUniquer.allocator);
166 
167     llvm::sys::SmartScopedWriter<true> lock(storageUniquer.mutex);
168     return mutationFn(storageUniquer.allocator);
169   }
170 
171   //===--------------------------------------------------------------------===//
172   // Singleton Storage
173   //===--------------------------------------------------------------------===//
174 
175   /// Get or create an instance of a singleton storage class.
176   BaseStorage *getSingleton(TypeID id) {
177     BaseStorage *singletonInstance = singletonInstances[id];
178     assert(singletonInstance && "expected singleton instance to exist");
179     return singletonInstance;
180   }
181 
182   /// Check if an instance of a singleton storage class exists.
183   bool hasSingleton(TypeID id) { return singletonInstances.count(id); }
184 
185   //===--------------------------------------------------------------------===//
186   // Instance Storage
187   //===--------------------------------------------------------------------===//
188 
189   /// Map of type ids to the storage uniquer to use for registered objects.
190   DenseMap<TypeID, std::unique_ptr<ParametricStorageUniquer>>
191       parametricUniquers;
192 
193   /// Map of type ids to a singleton instance when the storage class is a
194   /// singleton.
195   DenseMap<TypeID, BaseStorage *> singletonInstances;
196 
197   /// Allocator used for uniquing singleton instances.
198   StorageAllocator singletonAllocator;
199 
200   /// Flag specifying if multi-threading is enabled within the uniquer.
201   bool threadingIsEnabled = true;
202 };
203 } // end namespace detail
204 } // namespace mlir
205 
206 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
207 StorageUniquer::~StorageUniquer() {}
208 
209 /// Set the flag specifying if multi-threading is disabled within the uniquer.
210 void StorageUniquer::disableMultithreading(bool disable) {
211   impl->threadingIsEnabled = !disable;
212 }
213 
214 /// Implementation for getting/creating an instance of a derived type with
215 /// parametric storage.
216 auto StorageUniquer::getParametricStorageTypeImpl(
217     TypeID id, unsigned hashValue,
218     function_ref<bool(const BaseStorage *)> isEqual,
219     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
220   return impl->getOrCreate(id, hashValue, isEqual, ctorFn);
221 }
222 
223 /// Implementation for registering an instance of a derived type with
224 /// parametric storage.
225 void StorageUniquer::registerParametricStorageTypeImpl(TypeID id) {
226   impl->parametricUniquers.try_emplace(
227       id, std::make_unique<ParametricStorageUniquer>());
228 }
229 
230 /// Implementation for getting an instance of a derived type with default
231 /// storage.
232 auto StorageUniquer::getSingletonImpl(TypeID id) -> BaseStorage * {
233   return impl->getSingleton(id);
234 }
235 
236 /// Test is the storage singleton is initialized.
237 bool StorageUniquer::isSingletonStorageInitialized(TypeID id) {
238   return impl->hasSingleton(id);
239 }
240 
241 /// Test is the parametric storage is initialized.
242 bool StorageUniquer::isParametricStorageInitialized(TypeID id) {
243   return impl->hasParametricStorage(id);
244 }
245 
246 /// Implementation for registering an instance of a derived type with default
247 /// storage.
248 void StorageUniquer::registerSingletonImpl(
249     TypeID id, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
250   assert(!impl->singletonInstances.count(id) &&
251          "storage class already registered");
252   impl->singletonInstances.try_emplace(id, ctorFn(impl->singletonAllocator));
253 }
254 
255 /// Implementation for erasing an instance of a derived type with parametric
256 /// storage.
257 void StorageUniquer::eraseImpl(TypeID id, unsigned hashValue,
258                                function_ref<bool(const BaseStorage *)> isEqual,
259                                function_ref<void(BaseStorage *)> cleanupFn) {
260   impl->erase(id, hashValue, isEqual, cleanupFn);
261 }
262 
263 /// Implementation for mutating an instance of a derived storage.
264 LogicalResult StorageUniquer::mutateImpl(
265     TypeID id, function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
266   return impl->mutate(id, mutationFn);
267 }
268