MLIR  16.0.0git
StorageUniquer.cpp
Go to the documentation of this file.
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 
10 
11 #include "mlir/Support/LLVM.h"
13 #include "mlir/Support/TypeID.h"
14 #include "llvm/Support/RWMutex.h"
15 
16 using namespace mlir;
17 using namespace mlir::detail;
18 
19 namespace {
20 /// This class represents a uniquer for storage instances of a specific type
21 /// that has parametric storage. It contains all of the necessary data to unique
22 /// storage instances in a thread safe way. This allows for the main uniquer to
23 /// bucket each of the individual sub-types removing the need to lock the main
24 /// uniquer itself.
25 class ParametricStorageUniquer {
26 public:
27  using BaseStorage = StorageUniquer::BaseStorage;
28  using StorageAllocator = StorageUniquer::StorageAllocator;
29 
30  /// A lookup key for derived instances of storage objects.
31  struct LookupKey {
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 private:
40  /// A utility wrapper object representing a hashed storage object. This class
41  /// contains a storage object and an existing computed hash value.
42  struct HashedStorage {
43  HashedStorage(unsigned hashValue = 0, BaseStorage *storage = nullptr)
44  : hashValue(hashValue), storage(storage) {}
45  unsigned hashValue;
46  BaseStorage *storage;
47  };
48 
49  /// Storage info for derived TypeStorage objects.
50  struct StorageKeyInfo {
51  static inline HashedStorage getEmptyKey() {
52  return HashedStorage(0, DenseMapInfo<BaseStorage *>::getEmptyKey());
53  }
54  static inline HashedStorage getTombstoneKey() {
55  return HashedStorage(0, DenseMapInfo<BaseStorage *>::getTombstoneKey());
56  }
57 
58  static inline unsigned getHashValue(const HashedStorage &key) {
59  return key.hashValue;
60  }
61  static inline unsigned getHashValue(const LookupKey &key) {
62  return key.hashValue;
63  }
64 
65  static inline bool isEqual(const HashedStorage &lhs,
66  const HashedStorage &rhs) {
67  return lhs.storage == rhs.storage;
68  }
69  static inline bool isEqual(const LookupKey &lhs, const HashedStorage &rhs) {
70  if (isEqual(rhs, getEmptyKey()) || isEqual(rhs, getTombstoneKey()))
71  return false;
72  // Invoke the equality function on the lookup key.
73  return lhs.isEqual(rhs.storage);
74  }
75  };
76  using StorageTypeSet = DenseSet<HashedStorage, StorageKeyInfo>;
77 
78  /// This class represents a single shard of the uniquer. The uniquer uses a
79  /// set of shards to allow for multiple threads to create instances with less
80  /// lock contention.
81  struct Shard {
82  /// The set containing the allocated storage instances.
83  StorageTypeSet instances;
84 
85  /// Allocator to use when constructing derived instances.
86  StorageAllocator allocator;
87 
88 #if LLVM_ENABLE_THREADS != 0
89  /// A mutex to keep uniquing thread-safe.
90  llvm::sys::SmartRWMutex<true> mutex;
91 #endif
92  };
93 
94  /// Get or create an instance of a param derived type in an thread-unsafe
95  /// fashion.
96  BaseStorage *
97  getOrCreateUnsafe(Shard &shard, LookupKey &key,
98  function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
99  auto existing = shard.instances.insert_as({key.hashValue}, key);
100  BaseStorage *&storage = existing.first->storage;
101  if (existing.second)
102  storage = ctorFn(shard.allocator);
103  return storage;
104  }
105 
106  /// Destroy all of the storage instances within the given shard.
107  void destroyShardInstances(Shard &shard) {
108  if (!destructorFn)
109  return;
110  for (HashedStorage &instance : shard.instances)
111  destructorFn(instance.storage);
112  }
113 
114 public:
115 #if LLVM_ENABLE_THREADS != 0
116  /// Initialize the storage uniquer with a given number of storage shards to
117  /// use. The provided shard number is required to be a valid power of 2. The
118  /// destructor function is used to destroy any allocated storage instances.
119  ParametricStorageUniquer(function_ref<void(BaseStorage *)> destructorFn,
120  size_t numShards = 8)
121  : shards(new std::atomic<Shard *>[numShards]), numShards(numShards),
122  destructorFn(destructorFn) {
123  assert(llvm::isPowerOf2_64(numShards) &&
124  "the number of shards is required to be a power of 2");
125  for (size_t i = 0; i < numShards; i++)
126  shards[i].store(nullptr, std::memory_order_relaxed);
127  }
128  ~ParametricStorageUniquer() {
129  // Free all of the allocated shards.
130  for (size_t i = 0; i != numShards; ++i) {
131  if (Shard *shard = shards[i].load()) {
132  destroyShardInstances(*shard);
133  delete shard;
134  }
135  }
136  }
137  /// Get or create an instance of a parametric type.
138  BaseStorage *
139  getOrCreate(bool threadingIsEnabled, unsigned hashValue,
140  function_ref<bool(const BaseStorage *)> isEqual,
141  function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
142  Shard &shard = getShard(hashValue);
143  ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
144  if (!threadingIsEnabled)
145  return getOrCreateUnsafe(shard, lookupKey, ctorFn);
146 
147  // Check for a instance of this object in the local cache.
148  auto localIt = localCache->insert_as({hashValue}, lookupKey);
149  BaseStorage *&localInst = localIt.first->storage;
150  if (localInst)
151  return localInst;
152 
153  // Check for an existing instance in read-only mode.
154  {
155  llvm::sys::SmartScopedReader<true> typeLock(shard.mutex);
156  auto it = shard.instances.find_as(lookupKey);
157  if (it != shard.instances.end())
158  return localInst = it->storage;
159  }
160 
161  // Acquire a writer-lock so that we can safely create the new storage
162  // instance.
163  llvm::sys::SmartScopedWriter<true> typeLock(shard.mutex);
164  return localInst = getOrCreateUnsafe(shard, lookupKey, ctorFn);
165  }
166  /// Run a mutation function on the provided storage object in a thread-safe
167  /// way.
169  mutate(bool threadingIsEnabled, BaseStorage *storage,
170  function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
171  Shard &shard = getShardFor(storage);
172  if (!threadingIsEnabled)
173  return mutationFn(shard.allocator);
174 
175  llvm::sys::SmartScopedWriter<true> lock(shard.mutex);
176  return mutationFn(shard.allocator);
177  }
178 
179 private:
180  /// Return the shard used for the given hash value.
181  Shard &getShard(unsigned hashValue) {
182  // Get a shard number from the provided hashvalue.
183  unsigned shardNum = hashValue & (numShards - 1);
184 
185  // Try to acquire an already initialized shard.
186  Shard *shard = shards[shardNum].load(std::memory_order_acquire);
187  if (shard)
188  return *shard;
189 
190  // Otherwise, try to allocate a new shard.
191  Shard *newShard = new Shard();
192  if (shards[shardNum].compare_exchange_strong(shard, newShard))
193  return *newShard;
194 
195  // If one was allocated before we can initialize ours, delete ours.
196  delete newShard;
197  return *shard;
198  }
199 
200  /// Return the shard that allocated the provided storage object.
201  Shard &getShardFor(BaseStorage *storage) {
202  for (size_t i = 0; i != numShards; ++i) {
203  if (Shard *shard = shards[i].load(std::memory_order_acquire)) {
204  llvm::sys::SmartScopedReader<true> lock(shard->mutex);
205  if (shard->allocator.allocated(storage))
206  return *shard;
207  }
208  }
209  llvm_unreachable("expected storage object to have a valid shard");
210  }
211 
212  /// A thread local cache for storage objects. This helps to reduce the lock
213  /// contention when an object already existing in the cache.
215 
216  /// A set of uniquer shards to allow for further bucketing accesses for
217  /// instances of this storage type. Each shard is lazily initialized to reduce
218  /// the overhead when only a small amount of shards are in use.
219  std::unique_ptr<std::atomic<Shard *>[]> shards;
220 
221  /// The number of available shards.
222  size_t numShards;
223 
224  /// Function to used to destruct any allocated storage instances.
225  function_ref<void(BaseStorage *)> destructorFn;
226 
227 #else
228  /// If multi-threading is disabled, ignore the shard parameter as we will
229  /// always use one shard. The destructor function is used to destroy any
230  /// allocated storage instances.
231  ParametricStorageUniquer(function_ref<void(BaseStorage *)> destructorFn,
232  size_t numShards = 0)
233  : destructorFn(destructorFn) {}
234  ~ParametricStorageUniquer() { destroyShardInstances(shard); }
235 
236  /// Get or create an instance of a parametric type.
237  BaseStorage *
238  getOrCreate(bool threadingIsEnabled, unsigned hashValue,
239  function_ref<bool(const BaseStorage *)> isEqual,
240  function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
241  ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
242  return getOrCreateUnsafe(shard, lookupKey, ctorFn);
243  }
244  /// Run a mutation function on the provided storage object in a thread-safe
245  /// way.
247  mutate(bool threadingIsEnabled, BaseStorage *storage,
248  function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
249  return mutationFn(shard.allocator);
250  }
251 
252 private:
253  /// The main uniquer shard that is used for allocating storage instances.
254  Shard shard;
255 
256  /// Function to used to destruct any allocated storage instances.
257  function_ref<void(BaseStorage *)> destructorFn;
258 #endif
259 };
260 } // namespace
261 
262 namespace mlir {
263 namespace detail {
264 /// This is the implementation of the StorageUniquer class.
268 
269  //===--------------------------------------------------------------------===//
270  // Parametric Storage
271  //===--------------------------------------------------------------------===//
272 
273  /// Check if an instance of a parametric storage class exists.
274  bool hasParametricStorage(TypeID id) { return parametricUniquers.count(id); }
275 
276  /// Get or create an instance of a parametric type.
277  BaseStorage *
278  getOrCreate(TypeID id, unsigned hashValue,
279  function_ref<bool(const BaseStorage *)> isEqual,
281  assert(parametricUniquers.count(id) &&
282  "creating unregistered storage instance");
283  ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
284  return storageUniquer.getOrCreate(threadingIsEnabled, hashValue, isEqual,
285  ctorFn);
286  }
287 
288  /// Run a mutation function on the provided storage object in a thread-safe
289  /// way.
291  mutate(TypeID id, BaseStorage *storage,
293  assert(parametricUniquers.count(id) &&
294  "mutating unregistered storage instance");
295  ParametricStorageUniquer &storageUniquer = *parametricUniquers[id];
296  return storageUniquer.mutate(threadingIsEnabled, storage, mutationFn);
297  }
298 
299  //===--------------------------------------------------------------------===//
300  // Singleton Storage
301  //===--------------------------------------------------------------------===//
302 
303  /// Get or create an instance of a singleton storage class.
305  BaseStorage *singletonInstance = singletonInstances[id];
306  assert(singletonInstance && "expected singleton instance to exist");
307  return singletonInstance;
308  }
309 
310  /// Check if an instance of a singleton storage class exists.
311  bool hasSingleton(TypeID id) const { return singletonInstances.count(id); }
312 
313  //===--------------------------------------------------------------------===//
314  // Instance Storage
315  //===--------------------------------------------------------------------===//
316 
317  /// Map of type ids to the storage uniquer to use for registered objects.
320 
321  /// Map of type ids to a singleton instance when the storage class is a
322  /// singleton.
324 
325  /// Allocator used for uniquing singleton instances.
327 
328  /// Flag specifying if multi-threading is enabled within the uniquer.
329  bool threadingIsEnabled = true;
330 };
331 } // namespace detail
332 } // namespace mlir
333 
336 
337 /// Set the flag specifying if multi-threading is disabled within the uniquer.
339  impl->threadingIsEnabled = !disable;
340 }
341 
342 /// Implementation for getting/creating an instance of a derived type with
343 /// parametric storage.
344 auto StorageUniquer::getParametricStorageTypeImpl(
345  TypeID id, unsigned hashValue,
346  function_ref<bool(const BaseStorage *)> isEqual,
347  function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
348  return impl->getOrCreate(id, hashValue, isEqual, ctorFn);
349 }
350 
351 /// Implementation for registering an instance of a derived type with
352 /// parametric storage.
353 void StorageUniquer::registerParametricStorageTypeImpl(
354  TypeID id, function_ref<void(BaseStorage *)> destructorFn) {
355  impl->parametricUniquers.try_emplace(
356  id, std::make_unique<ParametricStorageUniquer>(destructorFn));
357 }
358 
359 /// Implementation for getting an instance of a derived type with default
360 /// storage.
361 auto StorageUniquer::getSingletonImpl(TypeID id) -> BaseStorage * {
362  return impl->getSingleton(id);
363 }
364 
365 /// Test is the storage singleton is initialized.
367  return impl->hasSingleton(id);
368 }
369 
370 /// Test is the parametric storage is initialized.
372  return impl->hasParametricStorage(id);
373 }
374 
375 /// Implementation for registering an instance of a derived type with default
376 /// storage.
377 void StorageUniquer::registerSingletonImpl(
378  TypeID id, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
379  assert(!impl->singletonInstances.count(id) &&
380  "storage class already registered");
381  impl->singletonInstances.try_emplace(id, ctorFn(impl->singletonAllocator));
382 }
383 
384 /// Implementation for mutating an instance of a derived storage.
385 LogicalResult StorageUniquer::mutateImpl(
386  TypeID id, BaseStorage *storage,
387  function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
388  return impl->mutate(id, storage, mutationFn);
389 }
This class acts as the base storage that all storage classes must derived from.
This is a utility allocator used to allocate memory for instances of derived types.
void disableMultithreading(bool disable=true)
Set the flag specifying if multi-threading is disabled within the uniquer.
bool isSingletonStorageInitialized(TypeID id)
Test if there is a singleton storage uniquer initialized for the provided TypeID.
bool isParametricStorageInitialized(TypeID id)
Test if there is a parametric storage uniquer initialized for the provided TypeID.
This class provides support for defining a thread local object with non static storage duration.
This class provides an efficient unique identifier for a specific C++ type.
Definition: TypeID.h:104
Detect if any of the given parameter types has a sub-element handler.
Include the generated interface declarations.
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
This is the implementation of the StorageUniquer class.
bool hasSingleton(TypeID id) const
Check if an instance of a singleton storage class exists.
DenseMap< TypeID, std::unique_ptr< ParametricStorageUniquer > > parametricUniquers
Map of type ids to the storage uniquer to use for registered objects.
LogicalResult mutate(TypeID id, BaseStorage *storage, function_ref< LogicalResult(StorageAllocator &)> mutationFn)
Run a mutation function on the provided storage object in a thread-safe way.
BaseStorage * getOrCreate(TypeID id, unsigned hashValue, function_ref< bool(const BaseStorage *)> isEqual, function_ref< BaseStorage *(StorageAllocator &)> ctorFn)
Get or create an instance of a parametric type.
BaseStorage * getSingleton(TypeID id)
Get or create an instance of a singleton storage class.
StorageAllocator singletonAllocator
Allocator used for uniquing singleton instances.
DenseMap< TypeID, BaseStorage * > singletonInstances
Map of type ids to a singleton instance when the storage class is a singleton.
bool hasParametricStorage(TypeID id)
Check if an instance of a parametric storage class exists.