Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include "BitUtil.h"
       4             : #include "MemoryResource.h"
       5             : #include <unordered_map>
       6             : #include <vector>
       7             : #include <memory>
       8             : 
       9             : namespace elsa::mr
      10             : {
      11             :     namespace pool_resource
      12             :     {
      13             :         template <typename T>
      14             :         class PoolResource;
      15             : 
      16             :         // must be set to a power of 2, with sufficient unused bits for the bitfield
      17             :         // => granularity must be at least 8!
      18             :         // 256 is chosen to make sure types for cuda kernels are sufficiently aligned.
      19             :         static constexpr size_t BLOCK_GRANULARITY = 256;
      20             : 
      21             :         static constexpr size_t MIN_BLOCK_SIZE = BLOCK_GRANULARITY;
      22             : 
      23             :         static constexpr size_t MIN_BLOCK_SIZE_LOG = 8;
      24             : 
      25             :         static constexpr size_t BITFIELD_MASK = BLOCK_GRANULARITY - 1;
      26             : 
      27             :         static constexpr size_t SIZE_MASK = ~BITFIELD_MASK;
      28             : 
      29             :         static constexpr size_t FREE_BIT = 1 << 0;
      30             : 
      31             :         static constexpr size_t PREV_FREE_BIT = 1 << 1;
      32             : 
      33             :         static constexpr size_t CHUNK_START_BIT = 1 << 2;
      34             :     } // namespace pool_resource
      35             : 
      36             :     class PoolResourceConfig
      37             :     {
      38             :     private:
      39             :         template <typename T>
      40             :         friend class pool_resource::PoolResource;
      41             : 
      42             :         size_t maxChunkSize;
      43             : 
      44             :         // size of the large allocation chunks that are requested from the underlying allocator
      45             :         size_t chunkSize;
      46             : 
      47             :         size_t maxCachedChunks;
      48             : 
      49             :         constexpr PoolResourceConfig(size_t maxChunkSize, size_t chunkSize, size_t maxCachedChunks)
      50             :             : maxChunkSize{maxChunkSize}, chunkSize{chunkSize}, maxCachedChunks{maxCachedChunks}
      51          60 :         {
      52          60 :         }
      53             : 
      54             :     public:
      55             :         /// @brief Default configuration for a pool resource with (hopefully) sensible defaults.
      56             :         /// @return Default configuration for a pool resource.
      57             :         static constexpr PoolResourceConfig defaultConfig()
      58          60 :         {
      59          60 :             return PoolResourceConfig(static_cast<size_t>(1) << 33, static_cast<size_t>(1) << 22,
      60          60 :                                       1);
      61          60 :         }
      62             : 
      63             :         /// @brief Set the maximum size for chunks allocated by this resource.
      64             :         /// Chunks are regions of memory, from which the blocks that are returned by allocate()
      65             :         /// are suballocated. Hence, this value also limits the size of blocks that are managed
      66             :         /// by this resource. Larger allocations are also accepted, but are serviced by the this
      67             :         /// pool's upstream allocator.
      68             :         /// @param size Maximum size for blocks managed by this pool.
      69             :         /// The pool resource may not use this exact value, as some minimal internal
      70             :         /// alignment requirements are applied to it.
      71             :         /// @return self
      72             :         constexpr PoolResourceConfig& setMaxChunkSize(size_t size)
      73          15 :         {
      74          15 :             maxChunkSize = std::max(detail::alignUp(size, pool_resource::BLOCK_GRANULARITY),
      75          15 :                                     pool_resource::MIN_BLOCK_SIZE);
      76          15 :             return *this;
      77          15 :         }
      78             : 
      79             :         /// @brief Set the size of the chunks requested from the back-end resource. Allocations are
      80             :         /// serviced from these chunks.
      81             :         /// @param size Size of the chunks requested from the back-end resource. Must be at most
      82             :         /// as big as the maximum chunk size. The pool resource may not use this exact value, as
      83             :         /// some minimal internal alignment requirements are applied to it.
      84             :         /// @return self
      85             :         constexpr PoolResourceConfig& setChunkSize(size_t size)
      86          15 :         {
      87          15 :             chunkSize = std::max(detail::alignUp(size, pool_resource::BLOCK_GRANULARITY),
      88          15 :                                  pool_resource::MIN_BLOCK_SIZE);
      89          15 :             return *this;
      90          15 :         }
      91             : 
      92             :         /// @brief  Set the maximum number of empty chunks that are cached. Further empty chunks are
      93             :         /// returned to the upstream allocator.
      94             :         /// @param count Manimum count of empty chunks to cache, before releasing memory to the
      95             :         /// upstream allocator.
      96             :         /// @return self
      97             :         constexpr PoolResourceConfig& setMaxCachedChunks(size_t count)
      98          18 :         {
      99          18 :             maxCachedChunks = count;
     100          18 :             return *this;
     101          18 :         }
     102             : 
     103             :         /// @brief Check if the pool is misconfigured.
     104             :         /// @return true if: the configuration is valid.
     105             :         /// false if: chunkSize > maxChunkSize
     106             :         constexpr bool validate()
     107          60 :         {
     108             :             // chunk must at least by able to accomodate the largest possible block
     109          60 :             return chunkSize <= maxChunkSize && chunkSize >= pool_resource::MIN_BLOCK_SIZE;
     110          60 :         }
     111             :     };
     112             : 
     113             :     namespace pool_resource
     114             :     {
     115             :         struct Block {
     116             :             // size of the block, also storing the free and prevFree flags in its lowest two bits
     117             :             size_t _size;
     118             :             void* _address;
     119             :             union {
     120             :                 // if this block marks the start of a chunk, this field contains its size
     121             :                 size_t _chunkSize;
     122             :                 // if this block is not the beginning of a chunk, this the contains address
     123             :                 // of the block that is prior to this one in contiguous memory
     124             :                 void* _prevAddress;
     125             :             };
     126             :             // next block in the free list
     127             :             Block* _nextFree;
     128             :             // address of the previous block in the free list's next pointer
     129             :             Block** _pprevFree;
     130             : 
     131             :             void markFree();
     132             : 
     133             :             void markAllocated();
     134             : 
     135             :             void markChunkStart();
     136             : 
     137             :             void markPrevFree();
     138             : 
     139             :             void markPrevAllocated();
     140             : 
     141             :             bool isFree();
     142             : 
     143             :             bool isPrevFree();
     144             : 
     145             :             bool isChunkStart();
     146             : 
     147             :             void unlinkFree();
     148             : 
     149             :             void insertAfterFree(Block** pprev);
     150             : 
     151             :             size_t size();
     152             : 
     153             :             void setSize(size_t size);
     154             :         };
     155             : 
     156             :         template <typename FreeListStrategy>
     157             :         class PoolResource : public MemResInterface
     158             :         {
     159             :         private:
     160             :             MemoryResource _upstream;
     161             : 
     162             :             PoolResourceConfig _config;
     163             : 
     164             :             std::unordered_map<void*, std::unique_ptr<pool_resource::Block>> _addressToBlock;
     165             : 
     166             :             std::vector<pool_resource::Block*> _freeLists;
     167             : 
     168             :             uint64_t _freeListNonEmpty{0};
     169             : 
     170             :             size_t _cachedChunkCount{0};
     171             : 
     172             :             std::unique_ptr<std::unique_ptr<pool_resource::Block>[]> _cachedChunks;
     173             : 
     174             :             void insertFreeBlock(std::unique_ptr<pool_resource::Block>&& block);
     175             : 
     176             :             void linkFreeBlock(pool_resource::Block* block);
     177             : 
     178             :             void unlinkFreeBlock(pool_resource::Block* block);
     179             : 
     180             :             size_t freeListIndexForFreeChunk(size_t size);
     181             :             // Returns a registered (in the address map) block that is not in the free list.
     182             :             // The metadata of the block is correctly initialized.
     183             :             pool_resource::Block* expandPool(size_t requestedSize);
     184             : 
     185             :             void shrinkPool(std::unique_ptr<pool_resource::Block> block);
     186             : 
     187             :             void shrinkBlockAtTail(pool_resource::Block& block, void* blockAddress, size_t newSize,
     188             :                                    size_t oldSize);
     189             : 
     190             :             // This function is noexcept, because it makes no allocations. The only potentially
     191             :             // throwing functions it calls, are the find and erase methods on _addressToBlock.
     192             :             // Neither of these should throw, if the inner type raises no exceptions.
     193             :             void doDeallocate(void* ptr) noexcept;
     194             : 
     195             :         public:
     196             :             PoolResource(const PoolResource& other) = delete;
     197             : 
     198             :             PoolResource& operator=(const PoolResource& other) = delete;
     199             : 
     200             :             PoolResource(PoolResource&& other) noexcept = delete;
     201             : 
     202             :             PoolResource& operator=(PoolResource&& other) noexcept = delete;
     203             : 
     204             :         protected:
     205             :             PoolResource(MemoryResource upstream,
     206             :                          PoolResourceConfig config = PoolResourceConfig::defaultConfig());
     207             : 
     208             :             ~PoolResource();
     209             : 
     210             :         public:
     211             :             /// @brief Create a MemoryResource that encapsulates a PoolResource with the given
     212             :             /// config.
     213             :             /// @param upstream The back-end allocator that is called by the pool resource whenever
     214             :             /// it runs out of memory to service allocations.
     215             :             /// @param config The configuration for the created pool resource. It is the caller's
     216             :             /// responsibility to make sure this configuration is valid with
     217             :             /// PoolResourceConfig::validate(). If the configuration is not valid, the default
     218             :             /// configuration is used instead.
     219             :             /// @return A MemoryResource that encapsulates a PoolResource.
     220             :             static MemoryResource
     221             :                 make(MemoryResource upstream = globalResource(),
     222             :                      PoolResourceConfig config = PoolResourceConfig::defaultConfig());
     223             : 
     224             :             void* allocate(size_t size, size_t alignment) override;
     225             : 
     226             :             void deallocate(void* ptr, size_t size, size_t alignment) noexcept override;
     227             : 
     228             :             bool tryResize(void* ptr, size_t size, size_t alignment,
     229             :                            size_t newSize) noexcept override;
     230             :         };
     231             : 
     232             :         struct ConstantTimeFit {
     233             :             static pool_resource::Block*
     234             :                 selectBlock(uint64_t listOccupancy,
     235             :                             const std::vector<pool_resource::Block*>& freeLists, size_t blockSize);
     236             :         };
     237             : 
     238             :         struct FirstFit {
     239             :             static pool_resource::Block*
     240             :                 selectBlock(uint64_t listOccupancy,
     241             :                             const std::vector<pool_resource::Block*>& freeLists, size_t blockSize);
     242             :         };
     243             : 
     244             :         struct HybridFit {
     245             :             static pool_resource::Block*
     246             :                 selectBlock(uint64_t listOccupancy,
     247             :                             const std::vector<pool_resource::Block*>& freeLists, size_t blockSize);
     248             :         };
     249             :     } // namespace pool_resource
     250             : 
     251             :     /// @brief Flexible memory resource for the ContiguousStorage class.
     252             :     /// It allocates large chunks from its upstream allocator and suballocates from these
     253             :     /// chunks to serve requests. It can efficiently handle a wide range of allocation patterns.
     254             :     /// IMPORTANT: THIS RESOURCE IS NOT SYNCHRONIZED!
     255             :     /// Advantage over a plain UniversalResource: allocations/deallocations are faster, memory is
     256             :     /// potentially already mapped from previous use.
     257             :     /// Disadvantage: higher memory usage than a plain UniversalResource. Also, move assignment
     258             :     /// between containers with different memory resources is more costly. Use it only if you are
     259             :     /// sure that memory allocations are your bottle-neck! If your algorithm works on huge data,
     260             :     /// allocations are probably not worth optimizing.
     261             :     using PoolResource = pool_resource::PoolResource<pool_resource::HybridFit>;
     262             : 
     263             :     /// @brief Pool resource able to serve allocations in average constant time (provided the
     264             :     /// upstream allocator also gives this guarantee). May lead to more fragmentation than a first
     265             :     /// fit strategy.
     266             :     using ConstantFitPoolResource = pool_resource::PoolResource<pool_resource::ConstantTimeFit>;
     267             : 
     268             :     /// @brief Pool resource that serves allocations via searching the corresponding seg. free list
     269             :     /// in linear time.
     270             :     using FirstFitPoolResource = pool_resource::PoolResource<pool_resource::FirstFit>;
     271             : 
     272             :     /// @brief Pool resource follows the constant time fit strategy, but falls back to linear time
     273             :     /// search when no block can be found in the larger lists.
     274             :     using HybridFitPoolResource = pool_resource::PoolResource<pool_resource::HybridFit>;
     275             : } // namespace elsa::mr