Resource Management 6: VK support for resource allocation.

Refactor existing memory allocators by using a common
memory type and handle.

BUG=dawn:27

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

BUILD.gn

@@ -405,6 +405,8 @@ source_set("libdawn_native_sources") {
       "src/dawn_native/vulkan/Forward.h",
       "src/dawn_native/vulkan/MemoryAllocator.cpp",
       "src/dawn_native/vulkan/MemoryAllocator.h",
+      "src/dawn_native/vulkan/MemoryResourceAllocatorVk.cpp",
+      "src/dawn_native/vulkan/MemoryResourceAllocatorVk.h",
       "src/dawn_native/vulkan/NativeSwapChainImplVk.cpp",
       "src/dawn_native/vulkan/NativeSwapChainImplVk.h",
       "src/dawn_native/vulkan/PipelineLayoutVk.cpp",
@@ -415,6 +417,8 @@ source_set("libdawn_native_sources") {
       "src/dawn_native/vulkan/RenderPassCache.h",
       "src/dawn_native/vulkan/RenderPipelineVk.cpp",
       "src/dawn_native/vulkan/RenderPipelineVk.h",
+      "src/dawn_native/vulkan/ResourceMemoryVk.cpp",
+      "src/dawn_native/vulkan/ResourceMemoryVk.h",
       "src/dawn_native/vulkan/SamplerVk.cpp",
       "src/dawn_native/vulkan/SamplerVk.h",
       "src/dawn_native/vulkan/ShaderModuleVk.cpp",

src/dawn_native/ResourceMemoryAllocation.cpp

@@ -19,16 +19,23 @@
 
 namespace dawn_native {
 
-    static constexpr uint64_t INVALID_OFFSET = std::numeric_limits<uint64_t>::max();
+    static constexpr uint64_t kInvalidOffset = std::numeric_limits<uint64_t>::max();
 
     ResourceMemoryAllocation::ResourceMemoryAllocation()
-        : mMethod(AllocationMethod::kInvalid), mOffset(INVALID_OFFSET), mResourceHeap(nullptr) {
+        : mMethod(AllocationMethod::kInvalid),
+          mOffset(0),
+          mResourceHeap(nullptr),
+          mMappedPointer(nullptr) {
     }
 
     ResourceMemoryAllocation::ResourceMemoryAllocation(uint64_t offset,
                                                        ResourceHeapBase* resourceHeap,
-                                                       AllocationMethod method)
+                                                       AllocationMethod method,
-        : mMethod(method), mOffset(offset), mResourceHeap(resourceHeap) {
+                                                       uint8_t* mappedPointer)
+        : mMethod(method),
+          mOffset(offset),
+          mResourceHeap(resourceHeap),
+          mMappedPointer(mappedPointer) {
     }
 
     ResourceHeapBase* ResourceMemoryAllocation::GetResourceHeap() const {
@@ -46,8 +53,13 @@ namespace dawn_native {
         return mMethod;
     }
 
+    uint8_t* ResourceMemoryAllocation::GetMappedPointer() const {
+        return mMappedPointer;
+    }
+
     void ResourceMemoryAllocation::Invalidate() {
         mResourceHeap = nullptr;
         mMethod = AllocationMethod::kInvalid;
+        mOffset = kInvalidOffset;
     }
 }  // namespace dawn_native

src/dawn_native/ResourceMemoryAllocation.h

@@ -41,12 +41,14 @@ namespace dawn_native {
         ResourceMemoryAllocation();
         ResourceMemoryAllocation(uint64_t offset,
                                  ResourceHeapBase* resourceHeap,
-                                 AllocationMethod method);
+                                 AllocationMethod method,
+                                 uint8_t* mappedPointer = nullptr);
         ~ResourceMemoryAllocation() = default;
 
         ResourceHeapBase* GetResourceHeap() const;
         uint64_t GetOffset() const;
         AllocationMethod GetAllocationMethod() const;
+        uint8_t* GetMappedPointer() const;
 
         void Invalidate();
 
@@ -54,6 +56,7 @@ namespace dawn_native {
         AllocationMethod mMethod;
         uint64_t mOffset;
         ResourceHeapBase* mResourceHeap;
+        uint8_t* mMappedPointer;
     };
 }  // namespace dawn_native
 

src/dawn_native/vulkan/BufferVk.cpp

@@ -16,6 +16,9 @@
 
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
+#include "dawn_native/vulkan/MemoryResourceAllocatorVk.h"
+#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/VulkanError.h"
 
 #include <cstring>
 
@@ -113,6 +116,9 @@ namespace dawn_native { namespace vulkan {
 
     Buffer::Buffer(Device* device, const BufferDescriptor* descriptor)
         : BufferBase(device, descriptor) {
+    }
+
+    MaybeError Buffer::Initialize() {
         VkBufferCreateInfo createInfo;
         createInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
         createInfo.pNext = nullptr;
@@ -125,26 +131,25 @@ namespace dawn_native { namespace vulkan {
         createInfo.queueFamilyIndexCount = 0;
         createInfo.pQueueFamilyIndices = 0;
 
-        if (device->fn.CreateBuffer(device->GetVkDevice(), &createInfo, nullptr, &mHandle) !=
+        Device* device = ToBackend(GetDevice());
-            VK_SUCCESS) {
+        DAWN_TRY(CheckVkSuccess(
-            ASSERT(false);
+            device->fn.CreateBuffer(device->GetVkDevice(), &createInfo, nullptr, &mHandle),
-        }
+            "vkCreateBuffer"));
 
         VkMemoryRequirements requirements;
         device->fn.GetBufferMemoryRequirements(device->GetVkDevice(), mHandle, &requirements);
 
         bool requestMappable =
             (GetUsage() & (dawn::BufferUsage::MapRead | dawn::BufferUsage::MapWrite)) != 0;
-        if (!device->GetMemoryAllocator()->Allocate(requirements, requestMappable,
+        DAWN_TRY_ASSIGN(mMemoryAllocation, device->AllocateMemory(requirements, requestMappable));
-                                                    &mMemoryAllocation)) {
-            ASSERT(false);
-        }
 
-        if (device->fn.BindBufferMemory(device->GetVkDevice(), mHandle,
+        DAWN_TRY(CheckVkSuccess(
-                                        mMemoryAllocation.GetMemory(),
+            device->fn.BindBufferMemory(device->GetVkDevice(), mHandle,
-                                        mMemoryAllocation.GetMemoryOffset()) != VK_SUCCESS) {
+                                        ToBackend(mMemoryAllocation.GetResourceHeap())->GetMemory(),
-            ASSERT(false);
+                                        mMemoryAllocation.GetOffset()),
-        }
+            "vkBindBufferMemory"));
+
+        return {};
     }
 
     Buffer::~Buffer() {
@@ -243,7 +248,7 @@ namespace dawn_native { namespace vulkan {
     }
 
     void Buffer::DestroyImpl() {
-        ToBackend(GetDevice())->GetMemoryAllocator()->Free(&mMemoryAllocation);
+        ToBackend(GetDevice())->DeallocateMemory(mMemoryAllocation);
 
         if (mHandle != VK_NULL_HANDLE) {
             ToBackend(GetDevice())->GetFencedDeleter()->DeleteWhenUnused(mHandle);

src/dawn_native/vulkan/BufferVk.h

@@ -19,6 +19,7 @@
 
 #include "common/SerialQueue.h"
 #include "common/vulkan_platform.h"
+#include "dawn_native/ResourceMemoryAllocation.h"
 #include "dawn_native/vulkan/MemoryAllocator.h"
 
 namespace dawn_native { namespace vulkan {
@@ -31,6 +32,8 @@ namespace dawn_native { namespace vulkan {
         Buffer(Device* device, const BufferDescriptor* descriptor);
         ~Buffer();
 
+        MaybeError Initialize();
+
         void OnMapReadCommandSerialFinished(uint32_t mapSerial, const void* data);
         void OnMapWriteCommandSerialFinished(uint32_t mapSerial, void* data);
 
@@ -52,7 +55,7 @@ namespace dawn_native { namespace vulkan {
         MaybeError MapAtCreationImpl(uint8_t** mappedPointer) override;
 
         VkBuffer mHandle = VK_NULL_HANDLE;
-        DeviceMemoryAllocation mMemoryAllocation;
+        ResourceMemoryAllocation mMemoryAllocation;
 
         dawn::BufferUsage mLastUsage = dawn::BufferUsage::None;
     };

src/dawn_native/vulkan/DeviceVk.cpp

@@ -70,6 +70,7 @@ namespace dawn_native { namespace vulkan {
         mMapRequestTracker = std::make_unique<MapRequestTracker>(this);
         mMemoryAllocator = std::make_unique<MemoryAllocator>(this);
         mRenderPassCache = std::make_unique<RenderPassCache>(this);
+        mResourceAllocator = std::make_unique<MemoryResourceAllocator>(this);
 
         mExternalMemoryService = std::make_unique<external_memory::Service>(this);
         mExternalSemaphoreService = std::make_unique<external_semaphore::Service>(this);
@@ -153,7 +154,9 @@ namespace dawn_native { namespace vulkan {
         return new BindGroupLayout(this, descriptor);
     }
     ResultOrError<BufferBase*> Device::CreateBufferImpl(const BufferDescriptor* descriptor) {
-        return new Buffer(this, descriptor);
+        std::unique_ptr<Buffer> buffer = std::make_unique<Buffer>(this, descriptor);
+        DAWN_TRY(buffer->Initialize());
+        return buffer.release();
     }
     CommandBufferBase* Device::CreateCommandBuffer(CommandEncoderBase* encoder,
                                                    const CommandBufferDescriptor* descriptor) {
@@ -675,4 +678,21 @@ namespace dawn_native { namespace vulkan {
         return new Texture(this, descriptor, textureDescriptor, signalSemaphore, allocation,
                            waitSemaphores);
     }
+
+    ResultOrError<ResourceMemoryAllocation> Device::AllocateMemory(
+        VkMemoryRequirements requirements,
+        bool mappable) {
+        // TODO(crbug.com/dawn/27): Support sub-allocation.
+        ResourceMemoryAllocation allocation;
+        DAWN_TRY_ASSIGN(allocation, mResourceAllocator->Allocate(requirements, mappable));
+        return allocation;
+    }
+
+    void Device::DeallocateMemory(ResourceMemoryAllocation& allocation) {
+        mResourceAllocator->Deallocate(allocation);
+
+        // Invalidate the underlying resource heap in case the client accidentally
+        // calls DeallocateMemory again using the same allocation.
+        allocation.Invalidate();
+    }
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/DeviceVk.h

@@ -22,6 +22,7 @@
 #include "dawn_native/Device.h"
 #include "dawn_native/vulkan/CommandRecordingContext.h"
 #include "dawn_native/vulkan/Forward.h"
+#include "dawn_native/vulkan/MemoryResourceAllocatorVk.h"
 #include "dawn_native/vulkan/VulkanFunctions.h"
 #include "dawn_native/vulkan/VulkanInfo.h"
 
@@ -91,6 +92,11 @@ namespace dawn_native { namespace vulkan {
                                            uint64_t destinationOffset,
                                            uint64_t size) override;
 
+        ResultOrError<ResourceMemoryAllocation> AllocateMemory(VkMemoryRequirements requirements,
+                                                               bool mappable);
+
+        void DeallocateMemory(ResourceMemoryAllocation& allocation);
+
       private:
         ResultOrError<BindGroupBase*> CreateBindGroupImpl(
             const BindGroupDescriptor* descriptor) override;
@@ -128,6 +134,8 @@ namespace dawn_native { namespace vulkan {
         uint32_t mQueueFamily = 0;
         VkQueue mQueue = VK_NULL_HANDLE;
 
+        std::unique_ptr<MemoryResourceAllocator> mResourceAllocator;
+
         std::unique_ptr<FencedDeleter> mDeleter;
         std::unique_ptr<MapRequestTracker> mMapRequestTracker;
         std::unique_ptr<MemoryAllocator> mMemoryAllocator;

src/dawn_native/vulkan/Forward.h

@@ -29,6 +29,7 @@ namespace dawn_native { namespace vulkan {
     class PipelineLayout;
     class Queue;
     class RenderPipeline;
+    class ResourceMemory;
     class Sampler;
     class ShaderModule;
     class StagingBuffer;
@@ -47,6 +48,7 @@ namespace dawn_native { namespace vulkan {
         using PipelineLayoutType = PipelineLayout;
         using QueueType = Queue;
         using RenderPipelineType = RenderPipeline;
+        using ResourceHeapType = ResourceMemory;
         using SamplerType = Sampler;
         using ShaderModuleType = ShaderModule;
         using StagingBufferType = StagingBuffer;

src/dawn_native/vulkan/MemoryResourceAllocatorVk.cpp

@@ -0,0 +1,117 @@
+// Copyright 2019 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/vulkan/DeviceVk.h"
+#include "dawn_native/vulkan/FencedDeleter.h"
+#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/VulkanError.h"
+
+namespace dawn_native { namespace vulkan {
+
+    MemoryResourceAllocator::MemoryResourceAllocator(Device* device) : mDevice(device) {
+    }
+
+    int MemoryResourceAllocator::FindBestTypeIndex(VkMemoryRequirements requirements,
+                                                   bool mappable) {
+        const VulkanDeviceInfo& info = mDevice->GetDeviceInfo();
+
+        // Find a suitable memory type for this allocation
+        int bestType = -1;
+        for (size_t i = 0; i < info.memoryTypes.size(); ++i) {
+            // Resource must support this memory type
+            if ((requirements.memoryTypeBits & (1 << i)) == 0) {
+                continue;
+            }
+
+            // Mappable resource must be host visible
+            if (mappable &&
+                (info.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) == 0) {
+                continue;
+            }
+
+            // Mappable must also be host coherent.
+            if (mappable &&
+                (info.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) == 0) {
+                continue;
+            }
+
+            // Found the first candidate memory type
+            if (bestType == -1) {
+                bestType = static_cast<int>(i);
+                continue;
+            }
+
+            // For non-mappable resources, favor device local memory.
+            if (!mappable) {
+                if ((info.memoryTypes[bestType].propertyFlags &
+                     VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) == 0 &&
+                    (info.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) !=
+                        0) {
+                    bestType = static_cast<int>(i);
+                    continue;
+                }
+            }
+
+            // All things equal favor the memory in the biggest heap
+            VkDeviceSize bestTypeHeapSize =
+                info.memoryHeaps[info.memoryTypes[bestType].heapIndex].size;
+            VkDeviceSize candidateHeapSize = info.memoryHeaps[info.memoryTypes[i].heapIndex].size;
+            if (candidateHeapSize > bestTypeHeapSize) {
+                bestType = static_cast<int>(i);
+                continue;
+            }
+        }
+
+        return bestType;
+    }
+
+    ResultOrError<ResourceMemoryAllocation> MemoryResourceAllocator::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"));
+        }
+
+        return ResourceMemoryAllocation(/*offset*/ 0, new ResourceMemory(allocatedMemory),
+                                        AllocationMethod::kDirect,
+                                        static_cast<uint8_t*>(mappedPointer));
+    }
+
+    void MemoryResourceAllocator::Deallocate(ResourceMemoryAllocation& allocation) {
+        mDevice->GetFencedDeleter()->DeleteWhenUnused(
+            ToBackend(allocation.GetResourceHeap())->GetMemory());
+    }
+}}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/MemoryResourceAllocatorVk.h

@@ -0,0 +1,43 @@
+// Copyright 2019 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_VULKAN_MEMORYRESOURCEALLOCATORVK_H_
+#define DAWNNATIVE_VULKAN_MEMORYRESOURCEALLOCATORVK_H_
+
+#include "common/vulkan_platform.h"
+#include "dawn_native/Error.h"
+#include "dawn_native/ResourceMemoryAllocation.h"
+
+namespace dawn_native { namespace vulkan {
+
+    class Device;
+
+    class MemoryResourceAllocator {
+      public:
+        MemoryResourceAllocator(Device* device);
+        ~MemoryResourceAllocator() = default;
+
+        ResultOrError<ResourceMemoryAllocation> Allocate(VkMemoryRequirements requirements,
+                                                         bool mappable);
+        void Deallocate(ResourceMemoryAllocation& allocation);
+
+      private:
+        int FindBestTypeIndex(VkMemoryRequirements requirements, bool mappable);
+
+        Device* mDevice;
+    };
+
+}}  // namespace dawn_native::vulkan
+
+#endif  // DAWNNATIVE_VULKAN_MEMORYRESOURCEALLOCATORVK_H_

src/dawn_native/vulkan/ResourceMemoryVk.cpp

@@ -0,0 +1,26 @@
+// Copyright 2019 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/vulkan/ResourceMemoryVk.h"
+
+namespace dawn_native { namespace vulkan {
+
+    ResourceMemory::ResourceMemory(VkDeviceMemory memory) : mMemory(memory) {
+    }
+
+    VkDeviceMemory ResourceMemory::GetMemory() const {
+        return mMemory;
+    }
+
+}}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceMemoryVk.h

@@ -0,0 +1,36 @@
+// Copyright 2019 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_VULKAN_RESOURCEMEMORYVK_H_
+#define DAWNNATIVE_VULKAN_RESOURCEMEMORYVK_H_
+
+#include "common/vulkan_platform.h"
+#include "dawn_native/ResourceHeap.h"
+
+namespace dawn_native { namespace vulkan {
+
+    // Wrapper for physical memory used with or without a resource object.
+    class ResourceMemory : public ResourceHeapBase {
+      public:
+        ResourceMemory(VkDeviceMemory memory);
+        ~ResourceMemory() = default;
+
+        VkDeviceMemory GetMemory() const;
+
+      private:
+        VkDeviceMemory mMemory = VK_NULL_HANDLE;
+    };
+}}  // namespace dawn_native::vulkan
+
+#endif  // DAWNNATIVE_VULKAN_RESOURCEMEMORYVK_H_

src/dawn_native/vulkan/StagingBufferVk.cpp

@@ -15,7 +15,8 @@
 #include "dawn_native/vulkan/StagingBufferVk.h"
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
-#include "dawn_native/vulkan/MemoryAllocator.h"
+#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/VulkanError.h"
 
 namespace dawn_native { namespace vulkan {
 
@@ -34,22 +35,20 @@ namespace dawn_native { namespace vulkan {
         createInfo.queueFamilyIndexCount = 0;
         createInfo.pQueueFamilyIndices = 0;
 
-        if (mDevice->fn.CreateBuffer(mDevice->GetVkDevice(), &createInfo, nullptr, &mBuffer) !=
+        DAWN_TRY(CheckVkSuccess(
-            VK_SUCCESS) {
+            mDevice->fn.CreateBuffer(mDevice->GetVkDevice(), &createInfo, nullptr, &mBuffer),
-            return DAWN_DEVICE_LOST_ERROR("Unable to create staging buffer.");
+            "vkCreateBuffer"));
-        }
 
         VkMemoryRequirements requirements;
         mDevice->fn.GetBufferMemoryRequirements(mDevice->GetVkDevice(), mBuffer, &requirements);
 
-        if (!mDevice->GetMemoryAllocator()->Allocate(requirements, true, &mAllocation)) {
+        DAWN_TRY_ASSIGN(mAllocation, mDevice->AllocateMemory(requirements, true));
-            return DAWN_DEVICE_LOST_ERROR("Unable to allocate memory for staging buffer.");
-        }
 
-        if (mDevice->fn.BindBufferMemory(mDevice->GetVkDevice(), mBuffer, mAllocation.GetMemory(),
+        DAWN_TRY(CheckVkSuccess(
-                                         mAllocation.GetMemoryOffset()) != VK_SUCCESS) {
+            mDevice->fn.BindBufferMemory(mDevice->GetVkDevice(), mBuffer,
-            return DAWN_DEVICE_LOST_ERROR("Unable to attach memory to the staging buffer.");
+                                         ToBackend(mAllocation.GetResourceHeap())->GetMemory(),
-        }
+                                         mAllocation.GetOffset()),
+            "vkBindBufferMemory"));
 
         mMappedPointer = mAllocation.GetMappedPointer();
         if (mMappedPointer == nullptr) {
@@ -62,7 +61,7 @@ namespace dawn_native { namespace vulkan {
     StagingBuffer::~StagingBuffer() {
         mMappedPointer = nullptr;
         mDevice->GetFencedDeleter()->DeleteWhenUnused(mBuffer);
-        mDevice->GetMemoryAllocator()->Free(&mAllocation);
+        mDevice->DeallocateMemory(mAllocation);
     }
 
     VkBuffer StagingBuffer::GetBufferHandle() const {

src/dawn_native/vulkan/StagingBufferVk.h

@@ -15,8 +15,9 @@
 #ifndef DAWNNATIVE_STAGINGBUFFERVK_H_
 #define DAWNNATIVE_STAGINGBUFFERVK_H_
 
+#include "common/vulkan_platform.h"
+#include "dawn_native/ResourceMemoryAllocation.h"
 #include "dawn_native/StagingBuffer.h"
-#include "dawn_native/vulkan/MemoryAllocator.h"
 
 namespace dawn_native { namespace vulkan {
 
@@ -34,7 +35,7 @@ namespace dawn_native { namespace vulkan {
       private:
         Device* mDevice;
         VkBuffer mBuffer;
-        DeviceMemoryAllocation mAllocation;
+        ResourceMemoryAllocation mAllocation;
     };
 }}  // namespace dawn_native::vulkan