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 
50     // Check for an existing instance in read-only mode.
51     {
52       llvm::sys::SmartScopedReader<true> typeLock(mutex);
53       auto it = storageTypes.find_as(lookupKey);
54       if (it != storageTypes.end())
55         return it->storage;
56     }
57 
58     // Acquire a writer-lock so that we can safely create the new type instance.
59     llvm::sys::SmartScopedWriter<true> typeLock(mutex);
60 
61     // Check for an existing instance again here, because another writer thread
62     // may have already created one.
63     auto existing = storageTypes.insert_as({}, lookupKey);
64     if (!existing.second)
65       return existing.first->storage;
66 
67     // Otherwise, construct and initialize the derived storage for this type
68     // instance.
69     BaseStorage *storage = initializeStorage(kind, ctorFn);
70     *existing.first = HashedStorage{hashValue, storage};
71     return storage;
72   }
73 
74   /// Get or create an instance of a simple derived type.
75   BaseStorage *
76   getOrCreate(unsigned kind,
77               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
78     // Check for an existing instance in read-only mode.
79     {
80       llvm::sys::SmartScopedReader<true> typeLock(mutex);
81       auto it = simpleTypes.find(kind);
82       if (it != simpleTypes.end())
83         return it->second;
84     }
85 
86     // Acquire a writer-lock so that we can safely create the new type instance.
87     llvm::sys::SmartScopedWriter<true> typeLock(mutex);
88 
89     // Check for an existing instance again here, because another writer thread
90     // may have already created one.
91     auto &result = simpleTypes[kind];
92     if (result)
93       return result;
94 
95     // Otherwise, create and return a new storage instance.
96     return result = initializeStorage(kind, ctorFn);
97   }
98 
99   /// Erase an instance of a complex derived type.
100   void erase(unsigned kind, unsigned hashValue,
101              function_ref<bool(const BaseStorage *)> isEqual,
102              function_ref<void(BaseStorage *)> cleanupFn) {
103     LookupKey lookupKey{kind, hashValue, isEqual};
104 
105     // Acquire a writer-lock so that we can safely erase the type instance.
106     llvm::sys::SmartScopedWriter<true> typeLock(mutex);
107     auto existing = storageTypes.find_as(lookupKey);
108     if (existing == storageTypes.end())
109       return;
110 
111     // Cleanup the storage and remove it from the map.
112     cleanupFn(existing->storage);
113     storageTypes.erase(existing);
114   }
115 
116   //===--------------------------------------------------------------------===//
117   // Instance Storage
118   //===--------------------------------------------------------------------===//
119 
120   /// Utility to create and initialize a storage instance.
121   BaseStorage *
122   initializeStorage(unsigned kind,
123                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
124     BaseStorage *storage = ctorFn(allocator);
125     storage->kind = kind;
126     return storage;
127   }
128 
129   /// Storage info for derived TypeStorage objects.
130   struct StorageKeyInfo : DenseMapInfo<HashedStorage> {
131     static HashedStorage getEmptyKey() {
132       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getEmptyKey()};
133     }
134     static HashedStorage getTombstoneKey() {
135       return HashedStorage{0, DenseMapInfo<BaseStorage *>::getTombstoneKey()};
136     }
137 
138     static unsigned getHashValue(const HashedStorage &key) {
139       return key.hashValue;
140     }
141     static unsigned getHashValue(LookupKey key) { return key.hashValue; }
142 
143     static bool isEqual(const HashedStorage &lhs, const HashedStorage &rhs) {
144       return lhs.storage == rhs.storage;
145     }
146     static bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
147       if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
148         return false;
149       // If the lookup kind matches the kind of the storage, then invoke the
150       // equality function on the lookup key.
151       return lhs.kind == rhs.storage->getKind() && lhs.isEqual(rhs.storage);
152     }
153   };
154 
155   // Unique types with specific hashing or storage constraints.
156   using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
157   StorageTypeSet storageTypes;
158 
159   // Unique types with just the kind.
160   DenseMap<unsigned, BaseStorage *> simpleTypes;
161 
162   // Allocator to use when constructing derived type instances.
163   StorageUniquer::StorageAllocator allocator;
164 
165   // A mutex to keep type uniquing thread-safe.
166   llvm::sys::SmartRWMutex<true> mutex;
167 };
168 } // end namespace detail
169 } // namespace mlir
170 
171 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
172 StorageUniquer::~StorageUniquer() {}
173 
174 /// Implementation for getting/creating an instance of a derived type with
175 /// complex storage.
176 auto StorageUniquer::getImpl(
177     unsigned kind, unsigned hashValue,
178     function_ref<bool(const BaseStorage *)> isEqual,
179     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
180   return impl->getOrCreate(kind, hashValue, isEqual, ctorFn);
181 }
182 
183 /// Implementation for getting/creating an instance of a derived type with
184 /// default storage.
185 auto StorageUniquer::getImpl(
186     unsigned kind, function_ref<BaseStorage *(StorageAllocator &)> ctorFn)
187     -> BaseStorage * {
188   return impl->getOrCreate(kind, ctorFn);
189 }
190 
191 /// Implementation for erasing an instance of a derived type with complex
192 /// storage.
193 void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
194                                function_ref<bool(const BaseStorage *)> isEqual,
195                                function_ref<void(BaseStorage *)> cleanupFn) {
196   impl->erase(kind, hashValue, isEqual, cleanupFn);
197 }
198