Simplify MemoryPool interface. (#559)

Removes indirection of wrapping MemoryAllocation in a smart pointer. This allows memory pools to be storable in sequence containers. Support pooled dynamic memory allocation for CPU allocations. #321
intel · Aug 5, 2022 · f71b7c6 · f71b7c6
1 parent 5e12fb6
commit f71b7c6
Show file tree

Hide file tree

Showing 11 changed files with 73 additions and 70 deletions.
diff --git a/src/gpgmm/common/BuddyMemoryAllocator.cpp b/src/gpgmm/common/BuddyMemoryAllocator.cpp
@@ -62,22 +62,24 @@ namespace gpgmm {
                              request.NeverAllocate,
                              [&](const auto& block) -> MemoryBase* {
                                  const uint64_t memoryIndex = GetMemoryIndex(block->Offset);
-                                 std::unique_ptr<MemoryAllocation> memoryAllocation =
+                                 MemoryAllocation memoryAllocation =
                                      mUsedPool.AcquireFromPool(memoryIndex);
 
                                  // No existing, allocate new memory for the block.
-                                 if (memoryAllocation == nullptr) {
+                                 if (memoryAllocation == GPGMM_ERROR_INVALID_ALLOCATION) {
                                      MemoryAllocationRequest newRequest = request;
                                      newRequest.SizeInBytes = mMemorySize;
                                      newRequest.Alignment = mMemoryAlignment;
 
+                                     std::unique_ptr<MemoryAllocation> memoryAllocationPtr;
                                      GPGMM_TRY_ASSIGN(
                                          GetNextInChain()->TryAllocateMemory(newRequest),
-                                         memoryAllocation);
+                                         memoryAllocationPtr);
+                                     memoryAllocation = *memoryAllocationPtr;
                                  }
 
-                                 MemoryBase* memory = memoryAllocation->GetMemory();
-                                 mUsedPool.ReturnToPool(std::move(memoryAllocation), memoryIndex);
+                                 MemoryBase* memory = memoryAllocation.GetMemory();
+                                 mUsedPool.ReturnToPool(memoryAllocation, memoryIndex);
 
                                  return memory;
                              }),
@@ -109,15 +111,16 @@ namespace gpgmm {
 
         mBuddyBlockAllocator.DeallocateBlock(subAllocation->GetBlock());
 
-        std::unique_ptr<MemoryAllocation> memoryAllocation = mUsedPool.AcquireFromPool(memoryIndex);
+        MemoryAllocation memoryAllocation = mUsedPool.AcquireFromPool(memoryIndex);
 
-        MemoryBase* memory = memoryAllocation->GetMemory();
+        MemoryBase* memory = memoryAllocation.GetMemory();
         ASSERT(memory != nullptr);
 
         if (memory->RemoveSubAllocationRef()) {
-            GetNextInChain()->DeallocateMemory(std::move(memoryAllocation));
+            GetNextInChain()->DeallocateMemory(
+                std::make_unique<MemoryAllocation>(memoryAllocation));
         } else {
-            mUsedPool.ReturnToPool(std::move(memoryAllocation), memoryIndex);
+            mUsedPool.ReturnToPool(memoryAllocation, memoryIndex);
         }
     }
 

diff --git a/src/gpgmm/common/Error.h b/src/gpgmm/common/Error.h
@@ -19,6 +19,10 @@
 
 namespace gpgmm {
 
+#define GPGMM_ERROR_INVALID_ALLOCATION \
+    MemoryAllocation {                 \
+    }
+
 #define GPGMM_TRY_ASSIGN(expr, value)            \
     {                                            \
         auto result = expr;                      \

diff --git a/src/gpgmm/common/IndexedMemoryPool.cpp b/src/gpgmm/common/IndexedMemoryPool.cpp
@@ -22,20 +22,19 @@ namespace gpgmm {
     IndexedMemoryPool::IndexedMemoryPool(uint64_t memorySize) : MemoryPool(memorySize) {
     }
 
-    std::unique_ptr<MemoryAllocation> IndexedMemoryPool::AcquireFromPool(uint64_t memoryIndex) {
-        if (memoryIndex >= mPool.size()) {
-            mPool.resize(memoryIndex + 1);
+    MemoryAllocation IndexedMemoryPool::AcquireFromPool(uint64_t indexInPool) {
+        if (indexInPool >= mPool.size()) {
+            mPool.resize(indexInPool + 1);
         }
-        return std::move(mPool[memoryIndex]);
+        MemoryAllocation tmp = mPool[indexInPool];
+        mPool[indexInPool] = {};  // invalidate it
+        return tmp;
     }
 
-    void IndexedMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocation> allocation,
-                                         uint64_t memoryIndex) {
-        ASSERT(allocation != nullptr);
-        ASSERT(memoryIndex < mPool.size());
-        ASSERT(allocation->GetSize() == GetMemorySize());
-
-        mPool[memoryIndex] = std::move(allocation);
+    void IndexedMemoryPool::ReturnToPool(MemoryAllocation allocation, uint64_t indexInPool) {
+        ASSERT(indexInPool < mPool.size());
+        ASSERT(allocation.GetSize() == GetMemorySize());
+        mPool[indexInPool] = std::move(allocation);
     }
 
     uint64_t IndexedMemoryPool::ReleasePool(uint64_t bytesToRelease) {
@@ -45,7 +44,7 @@ namespace gpgmm {
     uint64_t IndexedMemoryPool::GetPoolSize() const {
         uint64_t count = 0;
         for (auto& allocation : mPool) {
-            if (allocation != nullptr) {
+            if (allocation != GPGMM_ERROR_INVALID_ALLOCATION) {
                 count++;
             }
         }

diff --git a/src/gpgmm/common/IndexedMemoryPool.h b/src/gpgmm/common/IndexedMemoryPool.h
@@ -27,15 +27,14 @@ namespace gpgmm {
         ~IndexedMemoryPool() override = default;
 
         // MemoryPool interface
-        std::unique_ptr<MemoryAllocation> AcquireFromPool(uint64_t memoryIndex) override;
-        void ReturnToPool(std::unique_ptr<MemoryAllocation> allocation,
-                          uint64_t memoryIndex) override;
+        MemoryAllocation AcquireFromPool(uint64_t indexInPool) override;
+        void ReturnToPool(MemoryAllocation allocation, uint64_t indexInPool) override;
         uint64_t ReleasePool(uint64_t bytesToRelease) override;
 
         uint64_t GetPoolSize() const override;
 
       private:
-        std::vector<std::unique_ptr<MemoryAllocation>> mPool;
+        std::vector<MemoryAllocation> mPool;
     };
 
 }  // namespace gpgmm

diff --git a/src/gpgmm/common/LIFOMemoryPool.cpp b/src/gpgmm/common/LIFOMemoryPool.cpp
@@ -24,22 +24,20 @@ namespace gpgmm {
     LIFOMemoryPool::LIFOMemoryPool(uint64_t memorySize) : MemoryPool(memorySize) {
     }
 
-    std::unique_ptr<MemoryAllocation> LIFOMemoryPool::AcquireFromPool(uint64_t memoryIndex) {
-        ASSERT(memoryIndex == kInvalidIndex);
+    MemoryAllocation LIFOMemoryPool::AcquireFromPool(uint64_t indexInPool) {
+        ASSERT(indexInPool == kInvalidIndex);
 
-        std::unique_ptr<MemoryAllocation> allocation;
+        MemoryAllocation allocation = {};
         if (!mPool.empty()) {
             allocation = std::move(mPool.front());
             mPool.pop_front();
         }
         return allocation;
     }
 
-    void LIFOMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocation> allocation,
-                                      uint64_t memoryIndex) {
-        ASSERT(allocation != nullptr);
-        ASSERT(memoryIndex == kInvalidIndex);
-        ASSERT(allocation->GetSize() == GetMemorySize());
+    void LIFOMemoryPool::ReturnToPool(MemoryAllocation allocation, uint64_t indexInPool) {
+        ASSERT(indexInPool == kInvalidIndex);
+        ASSERT(allocation.GetSize() == GetMemorySize());
 
         mPool.push_front(std::move(allocation));
     }

diff --git a/src/gpgmm/common/LIFOMemoryPool.h b/src/gpgmm/common/LIFOMemoryPool.h
@@ -28,16 +28,15 @@ namespace gpgmm {
         ~LIFOMemoryPool() override = default;
 
         // MemoryPool interface
-        std::unique_ptr<MemoryAllocation> AcquireFromPool(
-            uint64_t memoryIndex = kInvalidIndex) override;
-        void ReturnToPool(std::unique_ptr<MemoryAllocation> allocation,
-                          uint64_t memoryIndex = kInvalidIndex) override;
+        MemoryAllocation AcquireFromPool(uint64_t indexInPool = kInvalidIndex) override;
+        void ReturnToPool(MemoryAllocation allocation,
+                          uint64_t indexInPool = kInvalidIndex) override;
         uint64_t ReleasePool(uint64_t bytesToFree = kInvalidSize) override;
 
         uint64_t GetPoolSize() const override;
 
       private:
-        std::deque<std::unique_ptr<MemoryAllocation>> mPool;
+        std::deque<MemoryAllocation> mPool;
     };
 
 }  // namespace gpgmm

diff --git a/src/gpgmm/common/MemoryAllocation.h b/src/gpgmm/common/MemoryAllocation.h
@@ -19,8 +19,6 @@
 #include "gpgmm/utils/Limits.h"
 #include "include/gpgmm_export.h"
 
-#include <cstdint>
-
 namespace gpgmm {
 
     /** \enum AllocationMethod

diff --git a/src/gpgmm/common/MemoryPool.h b/src/gpgmm/common/MemoryPool.h
@@ -15,14 +15,13 @@
 #ifndef GPGMM_COMMON_MEMORYPOOL_H_
 #define GPGMM_COMMON_MEMORYPOOL_H_
 
+#include "gpgmm/common/MemoryAllocation.h"
 #include "gpgmm/utils/Limits.h"
 
 #include <memory>
 
 namespace gpgmm {
 
-    class MemoryAllocation;
-
     /** \struct MemoryPoolInfo
     Additional information about the memory pool.
     */
@@ -45,25 +44,24 @@ namespace gpgmm {
       public:
         /** \brief Constructs a pool for memory of the specified size.
 
-        @param memorySize Size, in bytes, of the memory object stored in the pool.
+        @param memorySize Size, in bytes, of the memory allocation stored in the pool.
         */
         explicit MemoryPool(uint64_t memorySize);
         virtual ~MemoryPool();
 
         /** \brief Retrieves a memory allocation from the pool using an optional index.
 
-        @param memoryIndex Optional index of the memory object to retrieve.
+        @param indexInPool Optional index of the memory allocation to retrieve.
         */
-        virtual std::unique_ptr<MemoryAllocation> AcquireFromPool(
-            uint64_t memoryIndex = kInvalidIndex) = 0;
+        virtual MemoryAllocation AcquireFromPool(uint64_t indexInPool = kInvalidIndex) = 0;
 
         /** \brief Returns a memory allocation back to the pool using an optional index.
 
         @param allocation A pointer to MemoryAllocation which will be returned.
-        @param memoryIndex Optional index of the memory object to return.
+        @param indexInPool Optional index of the memory allocation to return.
         */
-        virtual void ReturnToPool(std::unique_ptr<MemoryAllocation> allocation,
-                                  uint64_t memoryIndex = kInvalidIndex) = 0;
+        virtual void ReturnToPool(MemoryAllocation allocation,
+                                  uint64_t indexInPool = kInvalidIndex) = 0;
 
         /** \brief Deallocate or shrink the pool.
 
@@ -105,8 +103,9 @@ namespace gpgmm {
             uint64_t totalBytesReleased = 0;
             uint64_t lastIndex = 0;
             for (auto& allocation : pool) {
-                totalBytesReleased += allocation->GetSize();
-                allocation->GetAllocator()->DeallocateMemory(std::move(allocation));
+                totalBytesReleased += allocation.GetSize();
+                allocation.GetAllocator()->DeallocateMemory(
+                    std::make_unique<MemoryAllocation>(allocation));
                 lastIndex++;
                 if (totalBytesReleased >= bytesToRelease) {
                     break;

diff --git a/src/gpgmm/common/PooledMemoryAllocator.cpp b/src/gpgmm/common/PooledMemoryAllocator.cpp
@@ -42,20 +42,22 @@ namespace gpgmm {
 
         GPGMM_INVALID_IF(!ValidateRequest(request));
 
-        std::unique_ptr<MemoryAllocation> allocation = mPool->AcquireFromPool();
-        if (allocation == nullptr) {
-            GPGMM_TRY_ASSIGN(GetNextInChain()->TryAllocateMemory(request), allocation);
+        MemoryAllocation allocation = mPool->AcquireFromPool();
+        if (allocation == GPGMM_ERROR_INVALID_ALLOCATION) {
+            std::unique_ptr<MemoryAllocation> allocationPtr;
+            GPGMM_TRY_ASSIGN(GetNextInChain()->TryAllocateMemory(request), allocationPtr);
+            allocation = *allocationPtr;
         } else {
-            mInfo.FreeMemoryUsage -= allocation->GetSize();
+            mInfo.FreeMemoryUsage -= allocation.GetSize();
         }
 
         mInfo.UsedMemoryCount++;
-        mInfo.UsedMemoryUsage += allocation->GetSize();
+        mInfo.UsedMemoryUsage += allocation.GetSize();
 
-        MemoryBase* memory = allocation->GetMemory();
+        MemoryBase* memory = allocation.GetMemory();
         ASSERT(memory != nullptr);
 
-        return std::make_unique<MemoryAllocation>(this, memory, allocation->GetRequestSize());
+        return std::make_unique<MemoryAllocation>(this, memory, allocation.GetRequestSize());
     }
 
     void PooledMemoryAllocator::DeallocateMemory(std::unique_ptr<MemoryAllocation> allocation) {
@@ -71,8 +73,8 @@ namespace gpgmm {
         MemoryBase* memory = allocation->GetMemory();
         ASSERT(memory != nullptr);
 
-        mPool->ReturnToPool(std::make_unique<MemoryAllocation>(GetNextInChain(), memory,
-                                                               allocation->GetRequestSize()));
+        mPool->ReturnToPool(
+            MemoryAllocation(GetNextInChain(), memory, allocation->GetRequestSize()));
     }
 
     uint64_t PooledMemoryAllocator::ReleaseMemory(uint64_t bytesToRelease) {

diff --git a/src/gpgmm/common/SegmentedMemoryAllocator.cpp b/src/gpgmm/common/SegmentedMemoryAllocator.cpp
@@ -131,17 +131,19 @@ namespace gpgmm {
         MemorySegment* segment = GetOrCreateFreeSegment(memorySize);
         ASSERT(segment != nullptr);
 
-        std::unique_ptr<MemoryAllocation> allocation = segment->AcquireFromPool();
-        if (allocation == nullptr) {
-            GPGMM_TRY_ASSIGN(GetNextInChain()->TryAllocateMemory(request), allocation);
+        MemoryAllocation allocation = segment->AcquireFromPool();
+        if (allocation == GPGMM_ERROR_INVALID_ALLOCATION) {
+            std::unique_ptr<MemoryAllocation> allocationPtr;
+            GPGMM_TRY_ASSIGN(GetNextInChain()->TryAllocateMemory(request), allocationPtr);
+            allocation = *allocationPtr;
         } else {
-            mInfo.FreeMemoryUsage -= allocation->GetSize();
+            mInfo.FreeMemoryUsage -= allocation.GetSize();
         }
 
         mInfo.UsedMemoryCount++;
-        mInfo.UsedMemoryUsage += allocation->GetSize();
+        mInfo.UsedMemoryUsage += allocation.GetSize();
 
-        MemoryBase* memory = allocation->GetMemory();
+        MemoryBase* memory = allocation.GetMemory();
         ASSERT(memory != nullptr);
 
         memory->SetPool(segment);
@@ -166,8 +168,8 @@ namespace gpgmm {
         MemoryPool* pool = memory->GetPool();
         ASSERT(pool != nullptr);
 
-        pool->ReturnToPool(std::make_unique<MemoryAllocation>(GetNextInChain(), memory,
-                                                              allocation->GetRequestSize()));
+        pool->ReturnToPool(
+            MemoryAllocation(GetNextInChain(), memory, allocation->GetRequestSize()));
     }
 
     uint64_t SegmentedMemoryAllocator::ReleaseMemory(uint64_t bytesToRelease) {

diff --git a/src/tests/unittests/MemoryPoolTests.cpp b/src/tests/unittests/MemoryPoolTests.cpp
@@ -41,7 +41,7 @@ TEST_F(LIFOMemoryPoolTests, SingleAllocation) {
     EXPECT_EQ(pool.ReleasePool(), 0u);
     EXPECT_EQ(pool.GetInfo().SizeInBytes, 0u);
 
-    pool.ReturnToPool(allocator.TryAllocateMemory(CreateBasicRequest(kDefaultMemorySize)));
+    pool.ReturnToPool(*allocator.TryAllocateMemory(CreateBasicRequest(kDefaultMemorySize)));
     EXPECT_EQ(pool.GetInfo().SizeInBytes, kDefaultMemorySize);
     EXPECT_EQ(pool.GetPoolSize(), 1u);
 
@@ -66,7 +66,7 @@ TEST_F(LIFOMemoryPoolTests, MultipleAllocations) {
 
     constexpr uint64_t kPoolSize = 64;
     while (pool.GetPoolSize() < kPoolSize) {
-        pool.ReturnToPool(allocator.TryAllocateMemory(CreateBasicRequest(kDefaultMemorySize)));
+        pool.ReturnToPool(*allocator.TryAllocateMemory(CreateBasicRequest(kDefaultMemorySize)));
     }
 
     EXPECT_EQ(pool.GetInfo().SizeInBytes, kDefaultMemorySize * kPoolSize);