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   //===--------------------------------------------------------------------===//
128   // Instance Storage
129   //===--------------------------------------------------------------------===//
130 
131   /// Utility to create and initialize a storage instance.
132   BaseStorage *
133   initializeStorage(unsigned kind,
134                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
135     BaseStorage *storage = ctorFn(allocator);
136     storage->kind = kind;
137     return storage;
138   }
139 
140   /// Storage info for derived TypeStorage objects.
141   struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
142     static HashedStorage getEmptyKey() {
143       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()};
144     }
145     static HashedStorage getTombstoneKey() {
146       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()};
147     }
148 
149     static unsigned getHashValue(const HashedStorage &key) {
150       return key.hashValue;
151     }
152     static unsigned getHashValue(LookupKey key) { return key.hashValue; }
153 
154     static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) {
155       return lhs.storage == rhs.storage;
156     }
157     static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
158       if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
159         return false;
160       // If the lookup kind matches the kind of the storage, then invoke the
161       // equality function on the lookup key.
162       return lhs.kind == rhs.storage->getKind() && lhs.isEqual(rhs.storage);
163     }
164   };
165 
166   /// Unique types with specific hashing or storage constraints.
167   using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
168   StorageTypeSet storageTypes;
169 
170   /// Unique types with just the kind.
171   DenseMap<unsigned, BaseStorage *> simpleTypes;
172 
173   /// Allocator to use when constructing derived type instances.
174   StorageUniquer::StorageAllocator allocator;
175 
176   /// A mutex to keep type uniquing thread-safe.
177   llvm::sys::SmartRWMutex<true> mutex;
178 
179   /// Flag specifying if multi-threading is enabled within the uniquer.
180   bool threadingIsEnabled = true;
181 };
182 } // end namespace detail
183 } // namespace mlir
184 
185 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
186 StorageUniquer::~StorageUniquer() {}
187 
188 /// Set the flag specifying if multi-threading is disabled within the uniquer.
189 void StorageUniquer::disableMultithreading(bool disable) {
190   impl->threadingIsEnabled = !disable;
191 }
192 
193 /// Implementation for getting/creating an instance of a derived type with
194 /// complex storage.
195 auto StorageUniquer::getImpl(
196     unsigned kind, unsigned hashValue,
197     function_ref<bool(const BaseStorage *)> isEqual,
198     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
199   return impl->getOrCreate(kind, hashValue, isEqual, ctorFn);
200 }
201 
202 /// Implementation for getting/creating an instance of a derived type with
203 /// default storage.
204 auto StorageUniquer::getImpl(
205     unsigned kind, function_ref<BaseStorage *(StorageAllocator &)> ctorFn)
206     -> BaseStorage * {
207   return impl->getOrCreate(kind, ctorFn);
208 }
209 
210 /// Implementation for erasing an instance of a derived type with complex
211 /// storage.
212 void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
213                                function_ref<bool(const BaseStorage *)> isEqual,
214                                function_ref<void(BaseStorage *)> cleanupFn) {
215   impl->erase(kind, hashValue, isEqual, cleanupFn);
216 }
217