Vulkan: Implement initial version of the suballocation

This makes the Vulkan backend use the BuddyMemoryAllocator to
sub-allocate small resources inside a larger VkDeviceMemory object.
Right now the heuristic to decide to do suballocation is naive and
should be improved.

BUG=dawn:27

Change-Id: Idcc7b6686c086633c85328a7afb91ee84abf7b8c
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/12662
Reviewed-by: Austin Eng <enga@chromium.org>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>

src/dawn_native/vulkan/BufferVk.cpp

@@ -251,7 +251,7 @@ namespace dawn_native { namespace vulkan {
     }
 
     void Buffer::DestroyImpl() {
-        ToBackend(GetDevice())->DeallocateMemory(mMemoryAllocation);
+        ToBackend(GetDevice())->DeallocateMemory(&mMemoryAllocation);
 
         if (mHandle != VK_NULL_HANDLE) {
             ToBackend(GetDevice())->GetFencedDeleter()->DeleteWhenUnused(mHandle);

src/dawn_native/vulkan/DeviceVk.cpp

@@ -222,6 +222,8 @@ namespace dawn_native { namespace vulkan {
         // as it enqueues resources to be released.
         mDynamicUploader->Deallocate(mCompletedSerial);
 
+        mResourceMemoryAllocator->Tick(mCompletedSerial);
+
         mDeleter->Tick(mCompletedSerial);
 
         if (mRecordingContext.used) {
@@ -681,19 +683,11 @@ namespace dawn_native { namespace vulkan {
     ResultOrError<ResourceMemoryAllocation> Device::AllocateMemory(
         VkMemoryRequirements requirements,
         bool mappable) {
-        // TODO(crbug.com/dawn/27): Support sub-allocation.
         return mResourceMemoryAllocator->Allocate(requirements, mappable);
     }
 
-    void Device::DeallocateMemory(ResourceMemoryAllocation& allocation) {
-        if (allocation.GetInfo().mMethod == AllocationMethod::kInvalid) {
-            return;
-        }
+    void Device::DeallocateMemory(ResourceMemoryAllocation* allocation) {
         mResourceMemoryAllocator->Deallocate(allocation);
-
-        // Invalidate the underlying resource heap in case the client accidentally
-        // calls DeallocateMemory again using the same allocation.
-        allocation.Invalidate();
     }
 
     ResourceMemoryAllocator* Device::GetResourceMemoryAllocatorForTesting() const {

src/dawn_native/vulkan/DeviceVk.h

@@ -91,7 +91,7 @@ namespace dawn_native { namespace vulkan {
 
         ResultOrError<ResourceMemoryAllocation> AllocateMemory(VkMemoryRequirements requirements,
                                                                bool mappable);
-        void DeallocateMemory(ResourceMemoryAllocation& allocation);
+        void DeallocateMemory(ResourceMemoryAllocation* allocation);
 
         ResourceMemoryAllocator* GetResourceMemoryAllocatorForTesting() const;
 

src/dawn_native/vulkan/ResourceHeapVk.cpp

@@ -16,11 +16,16 @@
 
 namespace dawn_native { namespace vulkan {
 
-    ResourceHeap::ResourceHeap(VkDeviceMemory memory) : mMemory(memory) {
+    ResourceHeap::ResourceHeap(VkDeviceMemory memory, size_t memoryType)
+        : mMemory(memory), mMemoryType(memoryType) {
     }
 
     VkDeviceMemory ResourceHeap::GetMemory() const {
         return mMemory;
     }
 
+    size_t ResourceHeap::GetMemoryType() const {
+        return mMemoryType;
+    }
+
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceHeapVk.h

@@ -23,13 +23,15 @@ namespace dawn_native { namespace vulkan {
     // Wrapper for physical memory used with or without a resource object.
     class ResourceHeap : public ResourceHeapBase {
       public:
-        ResourceHeap(VkDeviceMemory memory);
+        ResourceHeap(VkDeviceMemory memory, size_t memoryType);
         ~ResourceHeap() = default;
 
         VkDeviceMemory GetMemory() const;
+        size_t GetMemoryType() const;
 
       private:
         VkDeviceMemory mMemory = VK_NULL_HANDLE;
+        size_t mMemoryType = 0;
     };
 
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.cpp

@@ -14,6 +14,8 @@
 
 #include "dawn_native/vulkan/ResourceMemoryAllocatorVk.h"
 
+#include "dawn_native/BuddyMemoryAllocator.h"
+#include "dawn_native/ResourceHeapAllocator.h"
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
 #include "dawn_native/vulkan/ResourceHeapVk.h"
@@ -21,7 +23,167 @@
 
 namespace dawn_native { namespace vulkan {
 
+    namespace {
+
+        // TODO(cwallez@chromium.org): This is a hardcoded heurstic to choose when to
+        // suballocate but it should ideally depend on the size of the memory heaps and other
+        // factors.
+        constexpr uint64_t kMaxBuddySystemSize = 32ull * 1024ull * 1024ull * 1024ull;  // 32GB
+        constexpr uint64_t kMaxSizeForSubAllocation = 4ull * 1024ull * 1024ull;        // 4MB
+
+        // Have each bucket of the buddy system allocate at least some resource of the maximum
+        // size
+        constexpr uint64_t kBuddyHeapsSize = 2 * kMaxSizeForSubAllocation;
+
+    }  // anonymous namespace
+
+    // SingleTypeAllocator is a combination of a BuddyMemoryAllocator and its client and can
+    // service suballocation requests, but for a single Vulkan memory type.
+
+    class ResourceMemoryAllocator::SingleTypeAllocator : public ResourceHeapAllocator {
+      public:
+        SingleTypeAllocator(Device* device, size_t memoryTypeIndex)
+            : mDevice(device),
+              mMemoryTypeIndex(memoryTypeIndex),
+              mBuddySystem(kMaxBuddySystemSize, kBuddyHeapsSize, this) {
+        }
+        ~SingleTypeAllocator() override = default;
+
+        ResultOrError<ResourceMemoryAllocation> AllocateMemory(
+            const VkMemoryRequirements& requirements) {
+            return mBuddySystem.Allocate(requirements.size, requirements.alignment);
+        }
+
+        void DeallocateMemory(const ResourceMemoryAllocation& allocation) {
+            mBuddySystem.Deallocate(allocation);
+        }
+
+        // Implementation of the MemoryAllocator interface to be a client of BuddyMemoryAllocator
+
+        ResultOrError<std::unique_ptr<ResourceHeapBase>> AllocateResourceHeap(
+            uint64_t size) override {
+            VkMemoryAllocateInfo allocateInfo;
+            allocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
+            allocateInfo.pNext = nullptr;
+            allocateInfo.allocationSize = size;
+            allocateInfo.memoryTypeIndex = mMemoryTypeIndex;
+
+            VkDeviceMemory allocatedMemory = VK_NULL_HANDLE;
+            VkResult allocationResult = mDevice->fn.AllocateMemory(
+                mDevice->GetVkDevice(), &allocateInfo, nullptr, &allocatedMemory);
+
+            // Handle vkAllocateMemory error but differentiate OOM that we want to surface to
+            // the application.
+            if (allocationResult == VK_ERROR_OUT_OF_DEVICE_MEMORY) {
+                return DAWN_OUT_OF_MEMORY_ERROR("OOM while creating the Vkmemory");
+            }
+            DAWN_TRY(CheckVkSuccess(allocationResult, "vkAllocateMemory"));
+
+            ASSERT(allocatedMemory != VK_NULL_HANDLE);
+            return {std::make_unique<ResourceHeap>(allocatedMemory, mMemoryTypeIndex)};
+        }
+
+        void DeallocateResourceHeap(std::unique_ptr<ResourceHeapBase> allocation) override {
+            mDevice->GetFencedDeleter()->DeleteWhenUnused(ToBackend(allocation.get())->GetMemory());
+        }
+
+      private:
+        Device* mDevice;
+        size_t mMemoryTypeIndex;
+        BuddyMemoryAllocator mBuddySystem;
+    };
+
+    // Implementation of ResourceMemoryAllocator
+
     ResourceMemoryAllocator::ResourceMemoryAllocator(Device* device) : mDevice(device) {
+        const VulkanDeviceInfo& info = mDevice->GetDeviceInfo();
+        mAllocatorsPerType.reserve(info.memoryTypes.size());
+
+        for (size_t i = 0; i < info.memoryTypes.size(); i++) {
+            mAllocatorsPerType.emplace_back(std::make_unique<SingleTypeAllocator>(mDevice, i));
+        }
+    }
+
+    ResourceMemoryAllocator::~ResourceMemoryAllocator() = default;
+
+    ResultOrError<ResourceMemoryAllocation> ResourceMemoryAllocator::Allocate(
+        const VkMemoryRequirements& requirements,
+        bool mappable) {
+        // The Vulkan spec guarantees at least on memory type is valid.
+        int memoryType = FindBestTypeIndex(requirements, mappable);
+        ASSERT(memoryType >= 0);
+
+        VkDeviceSize size = requirements.size;
+
+        // If the resource is too big, allocate memory just for it.
+        // Also allocate mappable resources separately because at the moment the mapped pointer
+        // is part of the resource and not the heap, which doesn't match the Vulkan model.
+        // TODO(cwallez@chromium.org): allow sub-allocating mappable resources, maybe.
+        if (requirements.size >= kMaxSizeForSubAllocation || mappable) {
+            std::unique_ptr<ResourceHeapBase> resourceHeap;
+            DAWN_TRY_ASSIGN(resourceHeap,
+                            mAllocatorsPerType[memoryType]->AllocateResourceHeap(size));
+
+            void* mappedPointer = nullptr;
+            if (mappable) {
+                DAWN_TRY(
+                    CheckVkSuccess(mDevice->fn.MapMemory(mDevice->GetVkDevice(),
+                                                         ToBackend(resourceHeap.get())->GetMemory(),
+                                                         0, size, 0, &mappedPointer),
+                                   "vkMapMemory"));
+            }
+
+            AllocationInfo info;
+            info.mMethod = AllocationMethod::kDirect;
+            return ResourceMemoryAllocation(info, /*offset*/ 0, resourceHeap.release(),
+                                            static_cast<uint8_t*>(mappedPointer));
+        } else {
+            return mAllocatorsPerType[memoryType]->AllocateMemory(requirements);
+        }
+    }
+
+    void ResourceMemoryAllocator::Deallocate(ResourceMemoryAllocation* allocation) {
+        switch (allocation->GetInfo().mMethod) {
+            // Some memory allocation can never be initialized, for example when wrapping
+            // swapchain VkImages with a Texture.
+            case AllocationMethod::kInvalid:
+                break;
+
+            // For direct allocation we can put the memory for deletion immediately and the fence
+            // deleter will make sure the resources are freed before the memory.
+            case AllocationMethod::kDirect:
+                mDevice->GetFencedDeleter()->DeleteWhenUnused(
+                    ToBackend(allocation->GetResourceHeap())->GetMemory());
+                break;
+
+            // Suballocations aren't freed immediately, otherwise another resource allocation could
+            // happen just after that aliases the old one and would require a barrier.
+            // TODO(cwallez@chromium.org): Maybe we can produce the correct barriers to reduce the
+            // latency to reclaim memory.
+            case AllocationMethod::kSubAllocated:
+                mSubAllocationsToDelete.Enqueue(*allocation, mDevice->GetPendingCommandSerial());
+                break;
+
+            default:
+                UNREACHABLE();
+                break;
+        }
+
+        // Invalidate the underlying resource heap in case the client accidentally
+        // calls DeallocateMemory again using the same allocation.
+        allocation->Invalidate();
+    }
+
+    void ResourceMemoryAllocator::Tick(Serial completedSerial) {
+        for (const ResourceMemoryAllocation& allocation :
+             mSubAllocationsToDelete.IterateUpTo(completedSerial)) {
+            ASSERT(allocation.GetInfo().mMethod == AllocationMethod::kSubAllocated);
+            size_t memoryType = ToBackend(allocation.GetResourceHeap())->GetMemoryType();
+
+            mAllocatorsPerType[memoryType]->DeallocateMemory(allocation);
+        }
+
+        mSubAllocationsToDelete.ClearUpTo(completedSerial);
     }
 
     int ResourceMemoryAllocator::FindBestTypeIndex(VkMemoryRequirements requirements,
@@ -78,44 +240,4 @@ namespace dawn_native { namespace vulkan {
         return bestType;
     }
 
-    ResultOrError<ResourceMemoryAllocation> ResourceMemoryAllocator::Allocate(
-        VkMemoryRequirements requirements,
-        bool mappable) {
-        int bestType = FindBestTypeIndex(requirements, mappable);
-
-        // TODO(cwallez@chromium.org): I think the Vulkan spec guarantees this should never
-        // happen
-        if (bestType == -1) {
-            return DAWN_DEVICE_LOST_ERROR("Unable to find memory for requirements.");
-        }
-
-        VkMemoryAllocateInfo allocateInfo;
-        allocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
-        allocateInfo.pNext = nullptr;
-        allocateInfo.allocationSize = requirements.size;
-        allocateInfo.memoryTypeIndex = static_cast<uint32_t>(bestType);
-
-        VkDeviceMemory allocatedMemory = VK_NULL_HANDLE;
-        DAWN_TRY(CheckVkSuccess(mDevice->fn.AllocateMemory(mDevice->GetVkDevice(), &allocateInfo,
-                                                           nullptr, &allocatedMemory),
-                                "vkAllocateMemory"));
-
-        void* mappedPointer = nullptr;
-        if (mappable) {
-            DAWN_TRY(CheckVkSuccess(mDevice->fn.MapMemory(mDevice->GetVkDevice(), allocatedMemory,
-                                                          0, requirements.size, 0, &mappedPointer),
-                                    "vkMapMemory"));
-        }
-
-        AllocationInfo info;
-        info.mMethod = AllocationMethod::kDirect;
-
-        return ResourceMemoryAllocation(info, /*offset*/ 0, new ResourceHeap(allocatedMemory),
-                                        static_cast<uint8_t*>(mappedPointer));
-    }
-
-    void ResourceMemoryAllocator::Deallocate(ResourceMemoryAllocation& allocation) {
-        mDevice->GetFencedDeleter()->DeleteWhenUnused(
-            ToBackend(allocation.GetResourceHeap())->GetMemory());
-    }
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.h

@@ -15,10 +15,13 @@
 #ifndef DAWNNATIVE_VULKAN_RESOURCEMEMORYALLOCATORVK_H_
 #define DAWNNATIVE_VULKAN_RESOURCEMEMORYALLOCATORVK_H_
 
+#include "common/SerialQueue.h"
 #include "common/vulkan_platform.h"
 #include "dawn_native/Error.h"
 #include "dawn_native/ResourceMemoryAllocation.h"
 
+#include <vector>
+
 namespace dawn_native { namespace vulkan {
 
     class Device;
@@ -26,16 +29,23 @@ namespace dawn_native { namespace vulkan {
     class ResourceMemoryAllocator {
       public:
         ResourceMemoryAllocator(Device* device);
-        ~ResourceMemoryAllocator() = default;
+        ~ResourceMemoryAllocator();
 
-        ResultOrError<ResourceMemoryAllocation> Allocate(VkMemoryRequirements requirements,
+        ResultOrError<ResourceMemoryAllocation> Allocate(const VkMemoryRequirements& requirements,
                                                          bool mappable);
-        void Deallocate(ResourceMemoryAllocation& allocation);
+        void Deallocate(ResourceMemoryAllocation* allocation);
+
+        void Tick(Serial completedSerial);
 
         int FindBestTypeIndex(VkMemoryRequirements requirements, bool mappable);
 
       private:
         Device* mDevice;
+
+        class SingleTypeAllocator;
+        std::vector<std::unique_ptr<SingleTypeAllocator>> mAllocatorsPerType;
+
+        SerialQueue<ResourceMemoryAllocation> mSubAllocationsToDelete;
     };
 
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/StagingBufferVk.cpp

@@ -61,7 +61,7 @@ namespace dawn_native { namespace vulkan {
     StagingBuffer::~StagingBuffer() {
         mMappedPointer = nullptr;
         mDevice->GetFencedDeleter()->DeleteWhenUnused(mBuffer);
-        mDevice->DeallocateMemory(mAllocation);
+        mDevice->DeallocateMemory(&mAllocation);
     }
 
     VkBuffer StagingBuffer::GetBufferHandle() const {

src/dawn_native/vulkan/TextureVk.cpp

@@ -583,7 +583,7 @@ namespace dawn_native { namespace vulkan {
 
             // For textures created from a VkImage, the allocation if kInvalid so the Device knows
             // to skip the deallocation of the (absence of) VkDeviceMemory.
-            device->DeallocateMemory(mMemoryAllocation);
+            device->DeallocateMemory(&mMemoryAllocation);
 
             if (mHandle != VK_NULL_HANDLE) {
                 device->GetFencedDeleter()->DeleteWhenUnused(mHandle);