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