Pool sub-allocated resource heaps.

Allow resource heaps to be recycled when
no longer used.

BUG=dawn:496

Change-Id: I36518f8b0c0b26d2cceecc4e7b05e00a5fd5bd25
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/26126
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Austin Eng <enga@chromium.org>
Commit-Queue: Bryan Bernhart <bryan.bernhart@intel.com>

src/dawn_native/BUILD.gn

@@ -228,6 +228,8 @@ source_set("dawn_native_sources") {
     "Pipeline.h",
     "PipelineLayout.cpp",
     "PipelineLayout.h",
+    "PooledResourceMemoryAllocator.cpp",
+    "PooledResourceMemoryAllocator.h",
     "ProgrammablePassEncoder.cpp",
     "ProgrammablePassEncoder.h",
     "QuerySet.cpp",

src/dawn_native/CMakeLists.txt

@@ -103,6 +103,8 @@ target_sources(dawn_native PRIVATE
     "Pipeline.h"
     "PipelineLayout.cpp"
     "PipelineLayout.h"
+    "PooledResourceMemoryAllocator.cpp"
+    "PooledResourceMemoryAllocator.h"
     "ProgrammablePassEncoder.cpp"
     "ProgrammablePassEncoder.h"
     "QuerySet.cpp"

src/dawn_native/PooledResourceMemoryAllocator.cpp

@@ -0,0 +1,60 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "dawn_native/PooledResourceMemoryAllocator.h"
+#include "dawn_native/Device.h"
+
+namespace dawn_native {
+
+    PooledResourceMemoryAllocator::PooledResourceMemoryAllocator(
+        ResourceHeapAllocator* heapAllocator)
+        : mHeapAllocator(heapAllocator) {
+    }
+
+    void PooledResourceMemoryAllocator::DestroyPool() {
+        for (auto& resourceHeap : mPool) {
+            ASSERT(resourceHeap != nullptr);
+            mHeapAllocator->DeallocateResourceHeap(std::move(resourceHeap));
+        }
+
+        mPool.clear();
+    }
+
+    ResultOrError<std::unique_ptr<ResourceHeapBase>>
+    PooledResourceMemoryAllocator::AllocateResourceHeap(uint64_t size) {
+        // Pooled memory is LIFO because memory can be evicted by LRU. However, this means
+        // pooling is disabled in-frame when the memory is still pending. For high in-frame
+        // memory users, FIFO might be preferable when memory consumption is a higher priority.
+        std::unique_ptr<ResourceHeapBase> memory;
+        if (!mPool.empty()) {
+            memory = std::move(mPool.front());
+            mPool.pop_front();
+        }
+
+        if (memory == nullptr) {
+            DAWN_TRY_ASSIGN(memory, mHeapAllocator->AllocateResourceHeap(size));
+        }
+
+        return std::move(memory);
+    }
+
+    void PooledResourceMemoryAllocator::DeallocateResourceHeap(
+        std::unique_ptr<ResourceHeapBase> allocation) {
+        mPool.push_front(std::move(allocation));
+    }
+
+    uint64_t PooledResourceMemoryAllocator::GetPoolSizeForTesting() const {
+        return mPool.size();
+    }
+}  // namespace dawn_native

src/dawn_native/PooledResourceMemoryAllocator.h

@@ -0,0 +1,53 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DAWNNATIVE_POOLEDRESOURCEMEMORYALLOCATOR_H_
+#define DAWNNATIVE_POOLEDRESOURCEMEMORYALLOCATOR_H_
+
+#include "common/SerialQueue.h"
+#include "dawn_native/ResourceHeapAllocator.h"
+
+#include <deque>
+
+namespace dawn_native {
+
+    class DeviceBase;
+
+    // |PooledResourceMemoryAllocator| allocates a fixed-size resource memory from a resource memory
+    // pool. Internally, it manages a list of heaps using LIFO (newest heaps are recycled first).
+    // The heap is in one of two states: AVAILABLE or not. Upon de-allocate, the heap is returned
+    // the pool and made AVAILABLE.
+    class PooledResourceMemoryAllocator : public ResourceHeapAllocator {
+      public:
+        PooledResourceMemoryAllocator(ResourceHeapAllocator* heapAllocator);
+        ~PooledResourceMemoryAllocator() override = default;
+
+        ResultOrError<std::unique_ptr<ResourceHeapBase>> AllocateResourceHeap(
+            uint64_t size) override;
+        void DeallocateResourceHeap(std::unique_ptr<ResourceHeapBase> allocation) override;
+
+        void DestroyPool();
+
+        // For testing purposes.
+        uint64_t GetPoolSizeForTesting() const;
+
+      private:
+        ResourceHeapAllocator* mHeapAllocator = nullptr;
+
+        std::deque<std::unique_ptr<ResourceHeapBase>> mPool;
+    };
+
+}  // namespace dawn_native
+
+#endif  // DAWNNATIVE_POOLEDRESOURCEMEMORYALLOCATOR_H_

src/dawn_native/d3d12/DeviceD3D12.cpp

@@ -570,8 +570,11 @@ namespace dawn_native { namespace d3d12 {
             ::CloseHandle(mFenceEvent);
         }
 
+        // Release recycled resource heaps.
+        mResourceAllocatorManager->DestroyPool();
+
         // We need to handle clearing up com object refs that were enqeued after TickImpl
-        mUsedComObjectRefs.ClearUpTo(GetCompletedCommandSerial());
+        mUsedComObjectRefs.ClearUpTo(std::numeric_limits<Serial>::max());
 
         ASSERT(mUsedComObjectRefs.Empty());
         ASSERT(!mPendingCommands.IsOpen());

src/dawn_native/d3d12/ResourceAllocatorManagerD3D12.cpp

@@ -172,8 +172,10 @@ namespace dawn_native { namespace d3d12 {
             mHeapAllocators[i] = std::make_unique<HeapAllocator>(
                 mDevice, GetD3D12HeapType(resourceHeapKind), GetD3D12HeapFlags(resourceHeapKind),
                 GetMemorySegment(device, GetD3D12HeapType(resourceHeapKind)));
+            mPooledHeapAllocators[i] =
+                std::make_unique<PooledResourceMemoryAllocator>(mHeapAllocators[i].get());
             mSubAllocatedResourceAllocators[i] = std::make_unique<BuddyMemoryAllocator>(
-                kMaxHeapSize, kMinHeapSize, mHeapAllocators[i].get());
+                kMaxHeapSize, kMinHeapSize, mPooledHeapAllocators[i].get());
         }
     }
 
@@ -396,4 +398,10 @@ namespace dawn_native { namespace d3d12 {
                                       /*offset*/ 0, std::move(committedResource), heap};
     }
 
+    void ResourceAllocatorManager::DestroyPool() {
+        for (auto& alloc : mPooledHeapAllocators) {
+            alloc->DestroyPool();
+        }
+    }
+
 }}  // namespace dawn_native::d3d12

src/dawn_native/d3d12/ResourceAllocatorManagerD3D12.h

@@ -17,6 +17,7 @@
 
 #include "common/SerialQueue.h"
 #include "dawn_native/BuddyMemoryAllocator.h"
+#include "dawn_native/PooledResourceMemoryAllocator.h"
 #include "dawn_native/d3d12/HeapAllocatorD3D12.h"
 #include "dawn_native/d3d12/ResourceHeapAllocationD3D12.h"
 
@@ -67,6 +68,8 @@ namespace dawn_native { namespace d3d12 {
 
         void Tick(Serial lastCompletedSerial);
 
+        void DestroyPool();
+
       private:
         void FreeMemory(ResourceHeapAllocation& allocation);
 
@@ -92,6 +95,9 @@ namespace dawn_native { namespace d3d12 {
             mSubAllocatedResourceAllocators;
         std::array<std::unique_ptr<HeapAllocator>, ResourceHeapKind::EnumCount> mHeapAllocators;
 
+        std::array<std::unique_ptr<PooledResourceMemoryAllocator>, ResourceHeapKind::EnumCount>
+            mPooledHeapAllocators;
+
         SerialQueue<ResourceHeapAllocation> mAllocationsToDelete;
     };
 

src/dawn_native/vulkan/DeviceVk.cpp

@@ -888,6 +888,9 @@ namespace dawn_native { namespace vulkan {
         mResourceMemoryAllocator->Tick(GetCompletedCommandSerial());
         mDeleter->Tick(GetCompletedCommandSerial());
 
+        // Allow recycled memory to be deleted.
+        mResourceMemoryAllocator->DestroyPool();
+
         // The VkRenderPasses in the cache can be destroyed immediately since all commands referring
         // to them are guaranteed to be finished executing.
         mRenderPassCache = nullptr;

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.cpp

@@ -46,17 +46,22 @@ namespace dawn_native { namespace vulkan {
             : mDevice(device),
               mMemoryTypeIndex(memoryTypeIndex),
               mMemoryHeapSize(memoryHeapSize),
+              mPooledMemoryAllocator(this),
               mBuddySystem(
                   // Round down to a power of 2 that's <= mMemoryHeapSize. This will always
                   // be a multiple of kBuddyHeapsSize because kBuddyHeapsSize is a power of 2.
                   uint64_t(1) << Log2(mMemoryHeapSize),
                   // Take the min in the very unlikely case the memory heap is tiny.
                   std::min(uint64_t(1) << Log2(mMemoryHeapSize), kBuddyHeapsSize),
-                  this) {
+                  &mPooledMemoryAllocator) {
             ASSERT(IsPowerOfTwo(kBuddyHeapsSize));
         }
         ~SingleTypeAllocator() override = default;
 
+        void DestroyPool() {
+            mPooledMemoryAllocator.DestroyPool();
+        }
+
         ResultOrError<ResourceMemoryAllocation> AllocateMemory(
             const VkMemoryRequirements& requirements) {
             return mBuddySystem.Allocate(requirements.size, requirements.alignment);
@@ -100,6 +105,7 @@ namespace dawn_native { namespace vulkan {
         Device* mDevice;
         size_t mMemoryTypeIndex;
         VkDeviceSize mMemoryHeapSize;
+        PooledResourceMemoryAllocator mPooledMemoryAllocator;
         BuddyMemoryAllocator mBuddySystem;
     };
 
@@ -258,4 +264,10 @@ namespace dawn_native { namespace vulkan {
         return bestType;
     }
 
+    void ResourceMemoryAllocator::DestroyPool() {
+        for (auto& alloc : mAllocatorsPerType) {
+            alloc->DestroyPool();
+        }
+    }
+
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.h

@@ -18,6 +18,7 @@
 #include "common/SerialQueue.h"
 #include "common/vulkan_platform.h"
 #include "dawn_native/Error.h"
+#include "dawn_native/PooledResourceMemoryAllocator.h"
 #include "dawn_native/ResourceMemoryAllocation.h"
 
 #include <memory>
@@ -36,6 +37,8 @@ namespace dawn_native { namespace vulkan {
                                                          bool mappable);
         void Deallocate(ResourceMemoryAllocation* allocation);
 
+        void DestroyPool();
+
         void Tick(Serial completedSerial);
 
         int FindBestTypeIndex(VkMemoryRequirements requirements, bool mappable);

src/tests/unittests/BuddyMemoryAllocatorTests.cpp

@@ -15,8 +15,12 @@
 #include <gtest/gtest.h>
 
 #include "dawn_native/BuddyMemoryAllocator.h"
+#include "dawn_native/PooledResourceMemoryAllocator.h"
 #include "dawn_native/ResourceHeapAllocator.h"
 
+#include <set>
+#include <vector>
+
 using namespace dawn_native;
 
 class DummyResourceHeapAllocator : public ResourceHeapAllocator {
@@ -34,6 +38,12 @@ class DummyBuddyResourceAllocator {
         : mAllocator(maxBlockSize, memorySize, &mHeapAllocator) {
     }
 
+    DummyBuddyResourceAllocator(uint64_t maxBlockSize,
+                                uint64_t memorySize,
+                                ResourceHeapAllocator* heapAllocator)
+        : mAllocator(maxBlockSize, memorySize, heapAllocator) {
+    }
+
     ResourceMemoryAllocation Allocate(uint64_t allocationSize, uint64_t alignment = 1) {
         ResultOrError<ResourceMemoryAllocation> result =
             mAllocator.Allocate(allocationSize, alignment);
@@ -120,6 +130,7 @@ TEST(BuddyMemoryAllocatorTests, MultipleHeaps) {
 
     // Second allocation creates second heap.
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 2u);
+    ASSERT_NE(allocation1.GetResourceHeap(), allocation2.GetResourceHeap());
 
     // Deallocate both allocations
     allocator.Deallocate(allocation1);
@@ -159,6 +170,7 @@ TEST(BuddyMemoryAllocatorTests, MultipleSplitHeaps) {
 
     // Second allocation re-uses first heap.
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 1u);
+    ASSERT_EQ(allocation1.GetResourceHeap(), allocation2.GetResourceHeap());
 
     ResourceMemoryAllocation allocation3 = allocator.Allocate(heapSize / 2);
     ASSERT_EQ(allocation3.GetInfo().mBlockOffset, heapSize);
@@ -166,6 +178,7 @@ TEST(BuddyMemoryAllocatorTests, MultipleSplitHeaps) {
 
     // Third allocation creates second heap.
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 2u);
+    ASSERT_NE(allocation1.GetResourceHeap(), allocation3.GetResourceHeap());
 
     // Deallocate all allocations in reverse order.
     allocator.Deallocate(allocation1);
@@ -210,6 +223,7 @@ TEST(BuddyMemoryAllocatorTests, MultiplSplitHeapsVariableSizes) {
 
     // A1 and A2 share H0
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 1u);
+    ASSERT_EQ(allocation1.GetResourceHeap(), allocation2.GetResourceHeap());
 
     ResourceMemoryAllocation allocation3 = allocator.Allocate(128);
     ASSERT_EQ(allocation3.GetInfo().mBlockOffset, 128u);
@@ -218,6 +232,7 @@ TEST(BuddyMemoryAllocatorTests, MultiplSplitHeapsVariableSizes) {
 
     // A3 creates and fully occupies a new heap.
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 2u);
+    ASSERT_NE(allocation2.GetResourceHeap(), allocation3.GetResourceHeap());
 
     ResourceMemoryAllocation allocation4 = allocator.Allocate(64);
     ASSERT_EQ(allocation4.GetInfo().mBlockOffset, 256u);
@@ -225,6 +240,7 @@ TEST(BuddyMemoryAllocatorTests, MultiplSplitHeapsVariableSizes) {
     ASSERT_EQ(allocation4.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 3u);
+    ASSERT_NE(allocation3.GetResourceHeap(), allocation4.GetResourceHeap());
 
     // R5 size forms 64 byte hole after R4.
     ResourceMemoryAllocation allocation5 = allocator.Allocate(128);
@@ -233,6 +249,7 @@ TEST(BuddyMemoryAllocatorTests, MultiplSplitHeapsVariableSizes) {
     ASSERT_EQ(allocation5.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 4u);
+    ASSERT_NE(allocation4.GetResourceHeap(), allocation5.GetResourceHeap());
 
     // Deallocate allocations in staggered order.
     allocator.Deallocate(allocation1);
@@ -282,6 +299,7 @@ TEST(BuddyMemoryAllocatorTests, SameSizeVariousAlignment) {
     ASSERT_EQ(allocation2.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 2u);
+    ASSERT_NE(allocation1.GetResourceHeap(), allocation2.GetResourceHeap());
 
     ResourceMemoryAllocation allocation3 = allocator.Allocate(64, 128);
     ASSERT_EQ(allocation3.GetInfo().mBlockOffset, 256u);
@@ -289,6 +307,7 @@ TEST(BuddyMemoryAllocatorTests, SameSizeVariousAlignment) {
     ASSERT_EQ(allocation3.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 3u);
+    ASSERT_NE(allocation2.GetResourceHeap(), allocation3.GetResourceHeap());
 
     ResourceMemoryAllocation allocation4 = allocator.Allocate(64, 64);
     ASSERT_EQ(allocation4.GetInfo().mBlockOffset, 320u);
@@ -296,6 +315,7 @@ TEST(BuddyMemoryAllocatorTests, SameSizeVariousAlignment) {
     ASSERT_EQ(allocation4.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 3u);
+    ASSERT_EQ(allocation3.GetResourceHeap(), allocation4.GetResourceHeap());
 }
 
 // Verify resource sub-allocation of various sizes with same alignments.
@@ -330,6 +350,7 @@ TEST(BuddyMemoryAllocatorTests, VariousSizeSameAlignment) {
     ASSERT_EQ(allocation2.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 1u);  // Reuses H0
+    ASSERT_EQ(allocation1.GetResourceHeap(), allocation2.GetResourceHeap());
 
     ResourceMemoryAllocation allocation3 = allocator.Allocate(128, alignment);
     ASSERT_EQ(allocation3.GetInfo().mBlockOffset, 128u);
@@ -337,6 +358,7 @@ TEST(BuddyMemoryAllocatorTests, VariousSizeSameAlignment) {
     ASSERT_EQ(allocation3.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 2u);
+    ASSERT_NE(allocation2.GetResourceHeap(), allocation3.GetResourceHeap());
 
     ResourceMemoryAllocation allocation4 = allocator.Allocate(128, alignment);
     ASSERT_EQ(allocation4.GetInfo().mBlockOffset, 256u);
@@ -344,6 +366,7 @@ TEST(BuddyMemoryAllocatorTests, VariousSizeSameAlignment) {
     ASSERT_EQ(allocation4.GetInfo().mMethod, AllocationMethod::kSubAllocated);
 
     ASSERT_EQ(allocator.ComputeTotalNumOfHeapsForTesting(), 3u);
+    ASSERT_NE(allocation3.GetResourceHeap(), allocation4.GetResourceHeap());
 }
 
 // Verify allocating a very large resource does not overflow.
@@ -356,3 +379,82 @@ TEST(BuddyMemoryAllocatorTests, AllocationOverflow) {
     ResourceMemoryAllocation invalidAllocation = allocator.Allocate(largeBlock);
     ASSERT_EQ(invalidAllocation.GetInfo().mMethod, AllocationMethod::kInvalid);
 }
+
+// Verify resource heaps will be reused from a pool.
+TEST(BuddyMemoryAllocatorTests, ReuseFreedHeaps) {
+    constexpr uint64_t kHeapSize = 128;
+    constexpr uint64_t kMaxBlockSize = 4096;
+
+    DummyResourceHeapAllocator heapAllocator;
+    PooledResourceMemoryAllocator poolAllocator(&heapAllocator);
+    DummyBuddyResourceAllocator allocator(kMaxBlockSize, kHeapSize, &poolAllocator);
+
+    std::set<ResourceHeapBase*> heaps = {};
+    std::vector<ResourceMemoryAllocation> allocations = {};
+
+    constexpr uint32_t kNumOfAllocations = 100;
+
+    // Allocate |kNumOfAllocations|.
+    for (uint32_t i = 0; i < kNumOfAllocations; i++) {
+        ResourceMemoryAllocation allocation = allocator.Allocate(4);
+        ASSERT_EQ(allocation.GetInfo().mMethod, AllocationMethod::kSubAllocated);
+        heaps.insert(allocation.GetResourceHeap());
+        allocations.push_back(std::move(allocation));
+    }
+
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), 0u);
+
+    // Return the allocations to the pool.
+    for (ResourceMemoryAllocation& allocation : allocations) {
+        allocator.Deallocate(allocation);
+    }
+
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), heaps.size());
+
+    // Allocate again reusing the same heaps.
+    for (uint32_t i = 0; i < kNumOfAllocations; i++) {
+        ResourceMemoryAllocation allocation = allocator.Allocate(4);
+        ASSERT_EQ(allocation.GetInfo().mMethod, AllocationMethod::kSubAllocated);
+        ASSERT_FALSE(heaps.insert(allocation.GetResourceHeap()).second);
+    }
+
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), 0u);
+}
+
+// Verify resource heaps that were reused from a pool can be destroyed.
+TEST(BuddyMemoryAllocatorTests, DestroyHeaps) {
+    constexpr uint64_t kHeapSize = 128;
+    constexpr uint64_t kMaxBlockSize = 4096;
+
+    DummyResourceHeapAllocator heapAllocator;
+    PooledResourceMemoryAllocator poolAllocator(&heapAllocator);
+    DummyBuddyResourceAllocator allocator(kMaxBlockSize, kHeapSize, &poolAllocator);
+
+    std::set<ResourceHeapBase*> heaps = {};
+    std::vector<ResourceMemoryAllocation> allocations = {};
+
+    // Count by heap (vs number of allocations) to ensure there are exactly |kNumOfHeaps| worth of
+    // buffers. Otherwise, the heap may be reused if not full.
+    constexpr uint32_t kNumOfHeaps = 10;
+
+    // Allocate |kNumOfHeaps| worth.
+    while (heaps.size() < kNumOfHeaps) {
+        ResourceMemoryAllocation allocation = allocator.Allocate(4);
+        ASSERT_EQ(allocation.GetInfo().mMethod, AllocationMethod::kSubAllocated);
+        heaps.insert(allocation.GetResourceHeap());
+        allocations.push_back(std::move(allocation));
+    }
+
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), 0u);
+
+    // Return the allocations to the pool.
+    for (ResourceMemoryAllocation& allocation : allocations) {
+        allocator.Deallocate(allocation);
+    }
+
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), kNumOfHeaps);
+
+    // Make sure we can destroy the remaining heaps.
+    poolAllocator.DestroyPool();
+    ASSERT_EQ(poolAllocator.GetPoolSizeForTesting(), 0u);
+}