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