Actual source code: object_pool.hpp

  1: #pragma once

  3: #include <petsc/private/petscimpl.h>
  4: #include <petsc/private/mempoison.h>

  6: #include <petsc/private/cpp/register_finalize.hpp>
  7: #include <petsc/private/cpp/utility.hpp>
  8: #include <petsc/private/cpp/unordered_map.hpp>
  9: #include <petsc/private/cpp/memory.hpp>

 11: #include <cstddef> // std::size_t
 12: #include <vector>  // std::take_a_wild_guess

 14: namespace Petsc
 15: {

 17: namespace memory
 18: {

 20: enum class align_val_t : std::size_t {
 21: };

 23: } // namespace memory

 25: } // namespace Petsc

 27: namespace std
 28: {

 30: template <>
 31: struct hash<::Petsc::memory::align_val_t> {
 32: #if PETSC_CPP_VERSION < 17
 33:   using argument_type = ::Petsc::memory::align_val_t;
 34:   using result_type   = size_t;
 35: #endif

 37:   constexpr size_t operator()(const ::Petsc::memory::align_val_t &x) const noexcept { return static_cast<size_t>(x); }
 38: };

 40: } // namespace std

 42: namespace Petsc
 43: {

 45: namespace memory
 46: {

 48: // ==========================================================================================
 49: // PoolAllocator
 50: //
 51: // A general purpose memory pool. It internally maintains an array of allocated memory regions
 52: // and their sizes. Currently does not prune the allocated memory in any way.
 53: // ==========================================================================================

 55: class PoolAllocator : public RegisterFinalizeable<PoolAllocator> {
 56:   using base_type = RegisterFinalizeable<PoolAllocator>;
 57:   friend base_type;

 59: public:
 60:   // define the size and alignment as separate types, this helps to disambiguate them at the
 61:   // callsite!
 62:   using size_type  = std::size_t;
 63:   using align_type = align_val_t;
 64:   using pool_type  = std::vector<std::pair<align_type, UnorderedMap<size_type, std::vector<void *>>>>;

 66:   PoolAllocator() noexcept                            = default;
 67:   PoolAllocator(PoolAllocator &&) noexcept            = default;
 68:   PoolAllocator &operator=(PoolAllocator &&) noexcept = default;

 70:   // the pool carries raw memory and is not copyable
 71:   PoolAllocator(const PoolAllocator &)            = delete;
 72:   PoolAllocator &operator=(const PoolAllocator &) = delete;

 74:   ~PoolAllocator() noexcept;

 76:   PetscErrorCode try_allocate(void **, size_type, align_type, bool *) noexcept;
 77:   PetscErrorCode allocate(void **, size_type, align_type, bool * = nullptr) noexcept;
 78:   PetscErrorCode deallocate(void **, size_type, align_type) noexcept;

 80:   static PetscErrorCode get_attributes(const void *, size_type *, align_type *) noexcept;
 81:   static PetscErrorCode unpoison(const void *, size_type *) noexcept;
 82:   static PetscErrorCode repoison(const void *, size_type) noexcept;

 84:   template <typename T>
 85:   PetscErrorCode for_each(T &&) noexcept;

 87:   class LockGuard {
 88:   public:
 89:     LockGuard() = delete;

 91:   private:
 92:     friend class PoolAllocator;

 94:     explicit LockGuard(PoolAllocator *pool) noexcept : pool_{pool} { ++pool_->locked_; }

 96:     struct PoolUnlocker {
 97:       void operator()(PoolAllocator *pool) const noexcept { --pool->locked_; }
 98:     };

100:     std::unique_ptr<PoolAllocator, PoolUnlocker> pool_{};
101:   };

103:   LockGuard lock_guard() noexcept;

105: private:
106:   pool_type pool_;
107:   int       locked_ = 0;

109:   struct AllocationHeader;

111:   PETSC_NODISCARD static constexpr size_type total_size_(size_type, align_type) noexcept;

113:   static PetscErrorCode valid_alignment_(align_type) noexcept;
114:   static PetscErrorCode extract_header_(void *, AllocationHeader **, bool = true) noexcept;
115:   static PetscErrorCode allocate_ptr_(size_type, align_type, void **) noexcept;

117:   static PetscErrorCode delete_ptr_(void **) noexcept;

119:   PETSC_NODISCARD pool_type       &pool() noexcept { return pool_; }
120:   PETSC_NODISCARD const pool_type &pool() const noexcept { return pool_; }

122:   PETSC_NODISCARD typename pool_type::iterator       find_align_(align_type) noexcept;
123:   PETSC_NODISCARD typename pool_type::const_iterator find_align_(align_type) const noexcept;

125: public:
126:   PetscErrorCode clear_(size_type * = nullptr) noexcept;
127:   PetscErrorCode finalize_() noexcept;
128: };

130: /*
131:   PoolAllocator::for_each - Perform an action on every allocation in the currently in the pool

133:   Input Parameter:
134: . callable - The callable used to perform the action, should accept a void *&.

136:   Notes:
137:   The callable may delete the pointer, but MUST set the pointer to nullptr in this case. If the
138:   pointer is deleted, it is removed from the pool.

140:   The pointers are walked in LIFO order from most recent first to least recent deallocation
141:   last.
142: */
143: template <typename T>
144: inline PetscErrorCode PoolAllocator::for_each(T &&callable) noexcept
145: {
146:   PetscFunctionBegin;
147:   for (auto &&align_it : pool()) {
148:     for (auto &&size_it : align_it.second) {
149:       auto &&ptr_stack = size_it.second;

151:       for (auto it = ptr_stack.rbegin(); it != ptr_stack.rend();) {
152:         size_type size;
153:         auto    &&ptr = *it;

155:         PetscCall(unpoison(ptr, &size));
156:         PetscCall(callable(ptr));
157:         if (ptr) {
158:           PetscCall(repoison(ptr, size));
159:           ++it;
160:         } else {
161:           // the callable has deleted the pointer, so we should remove it. it is a reverse
162:           // iterator though so we:
163:           //
164:           // 1. std::next(it).base() -> convert to forward iterator
165:           // 2. it = decltype(it){...} -> convert returned iterator back to reverse iterator
166:           PetscCallCXX(it = decltype(it){ptr_stack.erase(std::next(it).base())});
167:         }
168:       }
169:     }
170:   }
171:   PetscFunctionReturn(PETSC_SUCCESS);
172: }

174: // ==========================================================================================
175: // PoolAllocated
176: //
177: // A simple mixin to enable allocating a class from a Pool. That is, it provides a static
178: // operator new() and operator delete() member functions such that
179: //
180: // auto foo = new ClassDerivedFromPoolAllocated(args...);
181: //
182: // Allocates the new object from a pool and
183: //
184: // delete foo;
185: //
186: // Returns the memory to the pool.
187: // ==========================================================================================

189: class PoolAllocated {
190: public:
191:   using allocator_type = PoolAllocator;
192:   using size_type      = typename allocator_type::size_type;
193:   using align_type     = typename allocator_type::align_type;

195:   PETSC_NODISCARD static void *operator new(size_type) noexcept;
196:   static void operator delete(void *) noexcept;

198: #if PETSC_CPP_VERSION >= 17
199:   PETSC_NODISCARD static void *operator new(size_type, std::align_val_t) noexcept;
200:   static void operator delete(void *, std::align_val_t) noexcept;
201: #endif

203: protected:
204:   PETSC_NODISCARD static allocator_type &pool() noexcept;

206: private:
207:   static allocator_type pool_;
208: };

210: } // namespace memory

212: // ==========================================================================================
213: // ConstructorInterface
214: //
215: // Provides a common interface for constructors and destructors for use with the object
216: // pools. Specifically, each interface may provide the following functions:
217: //
218: // construct_(T *):
219: // Given a pointer to an allocated object, construct the object in that memory. This may
220: // allocate memory *within* the object, but must not reallocate the pointer itself. Defaults to
221: // placement new.
222: //
223: // destroy_(T *):
224: // Given a pointer to an object, destroy the object completely. This should clean up the
225: // pointed-to object but not deallocate the pointer itself. Any resources not cleaned up by
226: // this function will be leaked. Defaults to calling the objects destructor.
227: //
228: // invalidate_(T *):
229: // Similar to destroy_(), but you are allowed to leave resources behind in the
230: // object. Essentially puts the object into a "zombie" state to be reused later. Defaults to
231: // calling destroy_().
232: //
233: // reset_():
234: // Revives a previously invalidated object. Should restore the object to a "factory new" state
235: // as-if it had just been newly allocated, but may take advantage of the fact that some
236: // resources need not be re-aquired from scratch. Defaults to calling construct_().
237: // ==========================================================================================

239: template <typename T, typename Derived>
240: class ConstructorInterface {
241: public:
242:   using value_type = T;

244:   template <typename... Args>
245:   PetscErrorCode construct(value_type *ptr, Args &&...args) const noexcept
246:   {
247:     PetscFunctionBegin;
248:     PetscCall(this->underlying().construct_(ptr, std::forward<Args>(args)...));
249:     PetscFunctionReturn(PETSC_SUCCESS);
250:   }

252:   PetscErrorCode destroy(value_type *ptr) const noexcept
253:   {
254:     const Derived &underlying = this->underlying();

256:     PetscFunctionBegin;
257:     PetscCall(underlying.destroy_(ptr));
258:     PetscFunctionReturn(PETSC_SUCCESS);
259:   }

261:   template <typename... Args>
262:   PetscErrorCode reset(value_type *val, Args &&...args) const noexcept
263:   {
264:     const Derived &underlying = this->underlying();

266:     PetscFunctionBegin;
267:     PetscCall(underlying.reset_(val, std::forward<Args>(args)...));
268:     PetscFunctionReturn(PETSC_SUCCESS);
269:   }

271:   PetscErrorCode invalidate(value_type *ptr) const noexcept
272:   {
273:     const Derived &underlying = this->underlying();

275:     PetscFunctionBegin;
276:     PetscCall(underlying.invalidate_(ptr));
277:     PetscFunctionReturn(PETSC_SUCCESS);
278:   }

280: protected:
281:   template <typename... Args>
282:   static PetscErrorCode construct_(value_type *ptr, Args &&...args) noexcept
283:   {
284:     PetscFunctionBegin;
285:     PetscAssertPointer(ptr, 1);
286:     PetscCallCXX(util::construct_at(ptr, std::forward<Args>(args)...));
287:     PetscFunctionReturn(PETSC_SUCCESS);
288:   }

290:   static PetscErrorCode destroy_(value_type *ptr) noexcept
291:   {
292:     PetscFunctionBegin;
293:     if (ptr) {
294:       PetscAssertPointer(ptr, 1);
295:       PetscCallCXX(util::destroy_at(ptr));
296:     }
297:     PetscFunctionReturn(PETSC_SUCCESS);
298:   }

300:   template <typename... Args>
301:   PetscErrorCode reset_(value_type *val, Args &&...args) const noexcept
302:   {
303:     PetscFunctionBegin;
304:     PetscCall(this->underlying().construct(val, std::forward<Args>(args)...));
305:     PetscFunctionReturn(PETSC_SUCCESS);
306:   }

308:   PetscErrorCode invalidate_(value_type *ptr) const noexcept
309:   {
310:     PetscFunctionBegin;
311:     PetscCall(this->underlying().destroy(ptr));
312:     PetscFunctionReturn(PETSC_SUCCESS);
313:   }

315:   PETSC_NODISCARD Derived       &underlying() noexcept { return static_cast<Derived &>(*this); }
316:   PETSC_NODISCARD const Derived &underlying() const noexcept { return static_cast<const Derived &>(*this); }
317: };

319: template <typename T>
320: struct DefaultConstructor : ConstructorInterface<T, DefaultConstructor<T>> { };

322: // ==========================================================================================
323: // ObjectPool
324: //
325: // multi-purpose basic object-pool, useful for recirculating old "destroyed" objects. Registers
326: // all objects to be cleaned up on PetscFinalize()
327: // ==========================================================================================

329: template <typename T, typename Constructor = DefaultConstructor<T>>
330: class ObjectPool;

332: template <typename T, typename Constructor>
333: class ObjectPool : public RegisterFinalizeable<ObjectPool<T, Constructor>> {
334:   using base_type = RegisterFinalizeable<ObjectPool<T, Constructor>>;

336: public:
337:   using value_type       = T;
338:   using constructor_type = Constructor;
339:   using allocator_type   = memory::PoolAllocator;

341:   ObjectPool()                                  = default;
342:   ObjectPool(ObjectPool &&) noexcept            = default;
343:   ObjectPool &operator=(ObjectPool &&) noexcept = default;

345:   ~ObjectPool() noexcept;

347:   template <typename... Args>
348:   PetscErrorCode allocate(value_type **, Args &&...) noexcept;
349:   PetscErrorCode deallocate(value_type **) noexcept;

351:   PETSC_NODISCARD constructor_type       &constructor() noexcept { return pair_.first(); }
352:   PETSC_NODISCARD const constructor_type &constructor() const noexcept { return pair_.first(); }
353:   PETSC_NODISCARD allocator_type         &allocator() noexcept { return pair_.second(); }
354:   PETSC_NODISCARD const allocator_type   &allocator() const noexcept { return pair_.second(); }

356: private:
357:   util::compressed_pair<constructor_type, allocator_type> pair_{};

359:   using align_type = typename allocator_type::align_type;
360:   using size_type  = typename allocator_type::size_type;

362:   friend base_type;
363:   PetscErrorCode finalize_() noexcept;
364: };

366: // ==========================================================================================
367: // ObjectPool -- Private API
368: // ==========================================================================================

370: template <typename T, typename Constructor>
371: inline PetscErrorCode ObjectPool<T, Constructor>::finalize_() noexcept
372: {
373:   PetscFunctionBegin;
374:   {
375:     auto _ = this->allocator().lock_guard();

377:     // clang-format off
378:     PetscCall(
379:       this->allocator().for_each([&, this](void *ptr)
380:       {
381:         PetscFunctionBegin;
382:         PetscCall(this->constructor().destroy(static_cast<value_type *>(ptr)));
383:         PetscFunctionReturn(PETSC_SUCCESS);
384:       })
385:     );
386:     // clang-format on
387:   }
388:   PetscCall(this->allocator().clear_());
389:   PetscFunctionReturn(PETSC_SUCCESS);
390: }

392: // ==========================================================================================
393: // ObjectPool -- Public API
394: // ==========================================================================================

396: template <typename T, typename Constructor>
397: inline ObjectPool<T, Constructor>::~ObjectPool() noexcept
398: {
399:   PetscFunctionBegin;
400:   PetscCallAbort(PETSC_COMM_SELF, this->finalize());
401:   PetscFunctionReturnVoid();
402: }

404: /*
405:   ObjectPool::allocate - Allocate an object from the pool

407:   Input Parameters:
408: . args... - The arguments to be passed to the constructor for the object

410:   Output Parameter:
411: . obj - The pointer to the constructed object

413:   Notes:
414:   The user must deallocate the object using deallocate() and cannot free it themselves.
415: */
416: template <typename T, typename Constructor>
417: template <typename... Args>
418: inline PetscErrorCode ObjectPool<T, Constructor>::allocate(value_type **obj, Args &&...args) noexcept
419: {
420:   auto  allocated_from_pool = true;
421:   void *mem                 = nullptr;

423:   PetscFunctionBegin;
424:   PetscAssertPointer(obj, 1);
425:   // order is deliberate! We register our finalizer before the pool does so since we need to
426:   // destroy the objects within it before it deletes their memory.
427:   PetscCall(this->allocator().allocate(&mem, sizeof(value_type), static_cast<align_type>(alignof(value_type)), &allocated_from_pool));
428:   PetscCall(this->register_finalize());
429:   *obj = static_cast<value_type *>(mem);
430:   if (allocated_from_pool) {
431:     // if the allocation reused memory from the pool then this indicates the object is resettable.
432:     PetscCall(this->constructor().reset(*obj, std::forward<Args>(args)...));
433:   } else {
434:     PetscCall(this->constructor().construct(*obj, std::forward<Args>(args)...));
435:   }
436:   PetscFunctionReturn(PETSC_SUCCESS);
437: }

439: /*
440:   ObjectPool::deallocate - Return an object to the pool

442:   Input Parameter:
443: . obj - The pointer to the object to return

445:   Notes:
446:   Sets obj to nullptr on return. obj must have been allocated by the pool in order to be
447:   deallocated this way.
448: */
449: template <typename T, typename Constructor>
450: inline PetscErrorCode ObjectPool<T, Constructor>::deallocate(value_type **obj) noexcept
451: {
452:   PetscFunctionBegin;
453:   PetscAssertPointer(obj, 1);
454:   PetscCall(this->register_finalize());
455:   PetscCall(this->constructor().invalidate(*obj));
456:   PetscCall(this->allocator().deallocate(reinterpret_cast<void **>(obj), sizeof(value_type), static_cast<align_type>(alignof(value_type))));
457:   PetscFunctionReturn(PETSC_SUCCESS);
458: }

460: } // namespace Petsc