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