Vulkan: Use the ResourceMemoryAllocator for all resources

This removes the duplication of the memory allocators in preparation for
using sub-allocation in the Vulkan backend too.

Also renames ResourceMemory to ResourceHeap and MemoryResourceAllocator
to ResourceMemoryAllocator, and fixes a number of unused includes.

BUG=dawn:27

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

BUILD.gn

@@ -430,10 +430,6 @@ source_set("libdawn_native_sources") {
       "src/dawn_native/vulkan/FencedDeleter.cpp",
       "src/dawn_native/vulkan/FencedDeleter.h",
       "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",
@@ -444,8 +440,10 @@ 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/ResourceHeapVk.cpp",
-      "src/dawn_native/vulkan/ResourceMemoryVk.h",
+      "src/dawn_native/vulkan/ResourceHeapVk.h",
+      "src/dawn_native/vulkan/ResourceMemoryAllocatorVk.cpp",
+      "src/dawn_native/vulkan/ResourceMemoryAllocatorVk.h",
       "src/dawn_native/vulkan/SamplerVk.cpp",
       "src/dawn_native/vulkan/SamplerVk.h",
       "src/dawn_native/vulkan/ShaderModuleVk.cpp",

src/dawn_native/vulkan/BufferVk.cpp

@@ -16,8 +16,8 @@
 
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
-#include "dawn_native/vulkan/MemoryResourceAllocatorVk.h"
+#include "dawn_native/vulkan/ResourceHeapVk.h"
-#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/ResourceMemoryAllocatorVk.h"
 #include "dawn_native/vulkan/VulkanError.h"
 
 #include <cstring>

src/dawn_native/vulkan/BufferVk.h

@@ -20,7 +20,6 @@
 #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 {
 

src/dawn_native/vulkan/DeviceVk.cpp

@@ -33,6 +33,7 @@
 #include "dawn_native/vulkan/QueueVk.h"
 #include "dawn_native/vulkan/RenderPassCache.h"
 #include "dawn_native/vulkan/RenderPipelineVk.h"
+#include "dawn_native/vulkan/ResourceMemoryAllocatorVk.h"
 #include "dawn_native/vulkan/SamplerVk.h"
 #include "dawn_native/vulkan/ShaderModuleVk.h"
 #include "dawn_native/vulkan/StagingBufferVk.h"
@@ -68,9 +69,8 @@ namespace dawn_native { namespace vulkan {
         GatherQueueFromDevice();
         mDeleter = std::make_unique<FencedDeleter>(this);
         mMapRequestTracker = std::make_unique<MapRequestTracker>(this);
-        mMemoryAllocator = std::make_unique<MemoryAllocator>(this);
         mRenderPassCache = std::make_unique<RenderPassCache>(this);
-        mResourceAllocator = std::make_unique<MemoryResourceAllocator>(this);
+        mResourceMemoryAllocator = std::make_unique<ResourceMemoryAllocator>(this);
 
         mExternalMemoryService = std::make_unique<external_memory::Service>(this);
         mExternalSemaphoreService = std::make_unique<external_semaphore::Service>(this);
@@ -138,7 +138,6 @@ namespace dawn_native { namespace vulkan {
 
         mDeleter = nullptr;
         mMapRequestTracker = nullptr;
-        mMemoryAllocator = nullptr;
 
         // The VkRenderPasses in the cache can be destroyed immediately since all commands referring
         // to them are guaranteed to be finished executing.
@@ -223,8 +222,6 @@ namespace dawn_native { namespace vulkan {
         // as it enqueues resources to be released.
         mDynamicUploader->Deallocate(mCompletedSerial);
 
-        mMemoryAllocator->Tick(mCompletedSerial);
-
         mDeleter->Tick(mCompletedSerial);
 
         if (mRecordingContext.used) {
@@ -262,10 +259,6 @@ namespace dawn_native { namespace vulkan {
         return mMapRequestTracker.get();
     }
 
-    MemoryAllocator* Device::GetMemoryAllocator() const {
-        return mMemoryAllocator.get();
-    }
-
     FencedDeleter* Device::GetFencedDeleter() const {
         return mDeleter.get();
     }
@@ -689,19 +682,22 @@ namespace dawn_native { namespace vulkan {
         VkMemoryRequirements requirements,
         bool mappable) {
         // TODO(crbug.com/dawn/27): Support sub-allocation.
-        ResourceMemoryAllocation allocation;
+        return mResourceMemoryAllocator->Allocate(requirements, mappable);
-        DAWN_TRY_ASSIGN(allocation, mResourceAllocator->Allocate(requirements, mappable));
-        return allocation;
     }
 
     void Device::DeallocateMemory(ResourceMemoryAllocation& allocation) {
         if (allocation.GetInfo().mMethod == AllocationMethod::kInvalid) {
             return;
         }
-        mResourceAllocator->Deallocate(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 {
+        return mResourceMemoryAllocator.get();
+    }
+
 }}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/DeviceVk.h

@@ -22,7 +22,6 @@
 #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"
 
@@ -39,8 +38,8 @@ namespace dawn_native { namespace vulkan {
     struct ExternalImageDescriptor;
     class FencedDeleter;
     class MapRequestTracker;
-    class MemoryAllocator;
     class RenderPassCache;
+    class ResourceMemoryAllocator;
 
     class Device : public DeviceBase {
       public:
@@ -61,7 +60,6 @@ namespace dawn_native { namespace vulkan {
         BufferUploader* GetBufferUploader() const;
         FencedDeleter* GetFencedDeleter() const;
         MapRequestTracker* GetMapRequestTracker() const;
-        MemoryAllocator* GetMemoryAllocator() const;
         RenderPassCache* GetRenderPassCache() const;
 
         CommandRecordingContext* GetPendingRecordingContext();
@@ -93,9 +91,10 @@ namespace dawn_native { namespace vulkan {
 
         ResultOrError<ResourceMemoryAllocation> AllocateMemory(VkMemoryRequirements requirements,
                                                                bool mappable);
-
         void DeallocateMemory(ResourceMemoryAllocation& allocation);
 
+        ResourceMemoryAllocator* GetResourceMemoryAllocatorForTesting() const;
+
       private:
         ResultOrError<BindGroupBase*> CreateBindGroupImpl(
             const BindGroupDescriptor* descriptor) override;
@@ -133,11 +132,9 @@ 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;
+        std::unique_ptr<ResourceMemoryAllocator> mResourceMemoryAllocator;
         std::unique_ptr<RenderPassCache> mRenderPassCache;
 
         std::unique_ptr<external_memory::Service> mExternalMemoryService;

src/dawn_native/vulkan/Forward.h

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

src/dawn_native/vulkan/MemoryAllocator.cpp

@@ -1,139 +0,0 @@
-// Copyright 2017 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/MemoryAllocator.h"
-
-#include "dawn_native/vulkan/DeviceVk.h"
-#include "dawn_native/vulkan/FencedDeleter.h"
-
-namespace dawn_native { namespace vulkan {
-
-    DeviceMemoryAllocation::~DeviceMemoryAllocation() {
-        ASSERT(mMemory == VK_NULL_HANDLE);
-    }
-
-    VkDeviceMemory DeviceMemoryAllocation::GetMemory() const {
-        return mMemory;
-    }
-
-    size_t DeviceMemoryAllocation::GetMemoryOffset() const {
-        return mOffset;
-    }
-
-    uint8_t* DeviceMemoryAllocation::GetMappedPointer() const {
-        return mMappedPointer;
-    }
-
-    MemoryAllocator::MemoryAllocator(Device* device) : mDevice(device) {
-    }
-
-    MemoryAllocator::~MemoryAllocator() {
-    }
-
-    int MemoryAllocator::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;
-    }
-
-    bool MemoryAllocator::Allocate(VkMemoryRequirements requirements,
-                                   bool mappable,
-                                   DeviceMemoryAllocation* allocation) {
-        int bestType = FindBestTypeIndex(requirements, mappable);
-        ASSERT(bestType >= 0);
-
-        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;
-        if (mDevice->fn.AllocateMemory(mDevice->GetVkDevice(), &allocateInfo, nullptr,
-                                       &allocatedMemory) != VK_SUCCESS) {
-            return false;
-        }
-
-        void* mappedPointer = nullptr;
-        if (mappable) {
-            if (mDevice->fn.MapMemory(mDevice->GetVkDevice(), allocatedMemory, 0, requirements.size,
-                                      0, &mappedPointer) != VK_SUCCESS) {
-                return false;
-            }
-        }
-
-        allocation->mMemory = allocatedMemory;
-        allocation->mOffset = 0;
-        allocation->mMappedPointer = static_cast<uint8_t*>(mappedPointer);
-
-        return true;
-    }
-
-    void MemoryAllocator::Free(DeviceMemoryAllocation* allocation) {
-        mDevice->GetFencedDeleter()->DeleteWhenUnused(allocation->mMemory);
-        allocation->mMemory = VK_NULL_HANDLE;
-        allocation->mOffset = 0;
-        allocation->mMappedPointer = nullptr;
-    }
-
-    void MemoryAllocator::Tick(Serial) {
-    }
-}}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/MemoryAllocator.h

@@ -1,59 +0,0 @@
-// Copyright 2017 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_MEMORYALLOCATOR_H_
-#define DAWNNATIVE_VULKAN_MEMORYALLOCATOR_H_
-
-#include "common/SerialQueue.h"
-#include "common/vulkan_platform.h"
-
-namespace dawn_native { namespace vulkan {
-
-    class Device;
-    class MemoryAllocator;
-
-    class DeviceMemoryAllocation {
-      public:
-        ~DeviceMemoryAllocation();
-        VkDeviceMemory GetMemory() const;
-        size_t GetMemoryOffset() const;
-        uint8_t* GetMappedPointer() const;
-
-      private:
-        friend class MemoryAllocator;
-        VkDeviceMemory mMemory = VK_NULL_HANDLE;
-        size_t mOffset = 0;
-        uint8_t* mMappedPointer = nullptr;
-    };
-
-    class MemoryAllocator {
-      public:
-        MemoryAllocator(Device* device);
-        ~MemoryAllocator();
-
-        int FindBestTypeIndex(VkMemoryRequirements requirements, bool mappable);
-        bool Allocate(VkMemoryRequirements requirements,
-                      bool mappable,
-                      DeviceMemoryAllocation* allocation);
-        void Free(DeviceMemoryAllocation* allocation);
-
-        void Tick(Serial finishedSerial);
-
-      private:
-        Device* mDevice = nullptr;
-    };
-
-}}  // namespace dawn_native::vulkan
-
-#endif  // DAWNNATIVE_VULKAN_MEMORYALLOCATOR_H_

src/dawn_native/vulkan/ResourceHeapVk.cpp

@@ -12,15 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/ResourceHeapVk.h"
 
 namespace dawn_native { namespace vulkan {
 
-    ResourceMemory::ResourceMemory(VkDeviceMemory memory) : mMemory(memory) {
+    ResourceHeap::ResourceHeap(VkDeviceMemory memory) : mMemory(memory) {
     }
 
-    VkDeviceMemory ResourceMemory::GetMemory() const {
+    VkDeviceMemory ResourceHeap::GetMemory() const {
         return mMemory;
     }
 
-}}  // namespace dawn_native::vulkan
+}}  // namespace dawn_native::vulkan

src/dawn_native/vulkan/ResourceHeapVk.h

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef DAWNNATIVE_VULKAN_RESOURCEMEMORYVK_H_
+#ifndef DAWNNATIVE_VULKAN_RESOURCEHEAPVK_H_
-#define DAWNNATIVE_VULKAN_RESOURCEMEMORYVK_H_
+#define DAWNNATIVE_VULKAN_RESOURCEHEAPVK_H_
 
 #include "common/vulkan_platform.h"
 #include "dawn_native/ResourceHeap.h"
@@ -21,16 +21,17 @@
 namespace dawn_native { namespace vulkan {
 
     // Wrapper for physical memory used with or without a resource object.
-    class ResourceMemory : public ResourceHeapBase {
+    class ResourceHeap : public ResourceHeapBase {
       public:
-        ResourceMemory(VkDeviceMemory memory);
+        ResourceHeap(VkDeviceMemory memory);
-        ~ResourceMemory() = default;
+        ~ResourceHeap() = default;
 
         VkDeviceMemory GetMemory() const;
 
       private:
         VkDeviceMemory mMemory = VK_NULL_HANDLE;
     };
+
 }}  // namespace dawn_native::vulkan
 
-#endif  // DAWNNATIVE_VULKAN_RESOURCEMEMORYVK_H_
+#endif  // DAWNNATIVE_VULKAN_RESOURCEHEAPVK_H_

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.cpp

@@ -12,17 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "dawn_native/vulkan/ResourceMemoryAllocatorVk.h"
+
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
-#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/ResourceHeapVk.h"
 #include "dawn_native/vulkan/VulkanError.h"
 
 namespace dawn_native { namespace vulkan {
 
-    MemoryResourceAllocator::MemoryResourceAllocator(Device* device) : mDevice(device) {
+    ResourceMemoryAllocator::ResourceMemoryAllocator(Device* device) : mDevice(device) {
     }
 
-    int MemoryResourceAllocator::FindBestTypeIndex(VkMemoryRequirements requirements,
+    int ResourceMemoryAllocator::FindBestTypeIndex(VkMemoryRequirements requirements,
                                                    bool mappable) {
         const VulkanDeviceInfo& info = mDevice->GetDeviceInfo();
 
@@ -76,7 +78,7 @@ namespace dawn_native { namespace vulkan {
         return bestType;
     }
 
-    ResultOrError<ResourceMemoryAllocation> MemoryResourceAllocator::Allocate(
+    ResultOrError<ResourceMemoryAllocation> ResourceMemoryAllocator::Allocate(
         VkMemoryRequirements requirements,
         bool mappable) {
         int bestType = FindBestTypeIndex(requirements, mappable);
@@ -108,11 +110,11 @@ namespace dawn_native { namespace vulkan {
         AllocationInfo info;
         info.mMethod = AllocationMethod::kDirect;
 
-        return ResourceMemoryAllocation(info, /*offset*/ 0, new ResourceMemory(allocatedMemory),
+        return ResourceMemoryAllocation(info, /*offset*/ 0, new ResourceHeap(allocatedMemory),
                                         static_cast<uint8_t*>(mappedPointer));
     }
 
-    void MemoryResourceAllocator::Deallocate(ResourceMemoryAllocation& allocation) {
+    void ResourceMemoryAllocator::Deallocate(ResourceMemoryAllocation& allocation) {
         mDevice->GetFencedDeleter()->DeleteWhenUnused(
             ToBackend(allocation.GetResourceHeap())->GetMemory());
     }

src/dawn_native/vulkan/ResourceMemoryAllocatorVk.h

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef DAWNNATIVE_VULKAN_MEMORYRESOURCEALLOCATORVK_H_
+#ifndef DAWNNATIVE_VULKAN_RESOURCEMEMORYALLOCATORVK_H_
-#define DAWNNATIVE_VULKAN_MEMORYRESOURCEALLOCATORVK_H_
+#define DAWNNATIVE_VULKAN_RESOURCEMEMORYALLOCATORVK_H_
 
 #include "common/vulkan_platform.h"
 #include "dawn_native/Error.h"
@@ -23,21 +23,21 @@ namespace dawn_native { namespace vulkan {
 
     class Device;
 
-    class MemoryResourceAllocator {
+    class ResourceMemoryAllocator {
       public:
-        MemoryResourceAllocator(Device* device);
+        ResourceMemoryAllocator(Device* device);
-        ~MemoryResourceAllocator() = default;
+        ~ResourceMemoryAllocator() = default;
 
         ResultOrError<ResourceMemoryAllocation> Allocate(VkMemoryRequirements requirements,
                                                          bool mappable);
         void Deallocate(ResourceMemoryAllocation& allocation);
 
-      private:
         int FindBestTypeIndex(VkMemoryRequirements requirements, bool mappable);
 
+      private:
         Device* mDevice;
     };
 
 }}  // namespace dawn_native::vulkan
 
-#endif  // DAWNNATIVE_VULKAN_MEMORYRESOURCEALLOCATORVK_H_
+#endif  // DAWNNATIVE_VULKAN_RESOURCEMEMORYALLOCATORVK_H_

src/dawn_native/vulkan/SamplerVk.h

@@ -19,10 +19,11 @@
 
 #include "common/vulkan_platform.h"
 #include "dawn_native/Error.h"
-#include "dawn_native/vulkan/MemoryAllocator.h"
 
 namespace dawn_native { namespace vulkan {
 
+    class Device;
+
     class Sampler : public SamplerBase {
       public:
         static ResultOrError<Sampler*> Create(Device* device, const SamplerDescriptor* descriptor);

src/dawn_native/vulkan/StagingBufferVk.cpp

@@ -15,7 +15,7 @@
 #include "dawn_native/vulkan/StagingBufferVk.h"
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
-#include "dawn_native/vulkan/ResourceMemoryVk.h"
+#include "dawn_native/vulkan/ResourceHeapVk.h"
 #include "dawn_native/vulkan/VulkanError.h"
 
 namespace dawn_native { namespace vulkan {

src/dawn_native/vulkan/TextureVk.cpp

@@ -24,6 +24,7 @@
 #include "dawn_native/vulkan/CommandRecordingContext.h"
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
+#include "dawn_native/vulkan/ResourceHeapVk.h"
 #include "dawn_native/vulkan/StagingBufferVk.h"
 #include "dawn_native/vulkan/UtilsVulkan.h"
 #include "dawn_native/vulkan/VulkanError.h"
@@ -460,14 +461,13 @@ namespace dawn_native { namespace vulkan {
         VkMemoryRequirements requirements;
         device->fn.GetImageMemoryRequirements(device->GetVkDevice(), mHandle, &requirements);
 
-        if (!device->GetMemoryAllocator()->Allocate(requirements, false, &mMemoryAllocation)) {
+        DAWN_TRY_ASSIGN(mMemoryAllocation, device->AllocateMemory(requirements, false));
-            return DAWN_OUT_OF_MEMORY_ERROR("Failed to allocate texture");
-        }
 
-        DAWN_TRY(CheckVkSuccess(device->fn.BindImageMemory(device->GetVkDevice(), mHandle,
+        DAWN_TRY(CheckVkSuccess(
-                                                           mMemoryAllocation.GetMemory(),
+            device->fn.BindImageMemory(device->GetVkDevice(), mHandle,
-                                                           mMemoryAllocation.GetMemoryOffset()),
+                                       ToBackend(mMemoryAllocation.GetResourceHeap())->GetMemory(),
-                                "BindImageMemory"));
+                                       mMemoryAllocation.GetOffset()),
+            "BindImageMemory"));
 
         if (device->IsToggleEnabled(Toggle::NonzeroClearResourcesOnCreationForTesting)) {
             DAWN_TRY(ClearTexture(ToBackend(GetDevice())->GetPendingRecordingContext(), 0,
@@ -581,13 +581,9 @@ namespace dawn_native { namespace vulkan {
         if (GetTextureState() == TextureState::OwnedInternal) {
             Device* device = ToBackend(GetDevice());
 
-            // If we own the resource, release it.
+            // For textures created from a VkImage, the allocation if kInvalid so the Device knows
-            if (mMemoryAllocation.GetMemory() != VK_NULL_HANDLE) {
+            // to skip the deallocation of the (absence of) VkDeviceMemory.
-                // We need to free both the memory allocation and the container. Memory should be
+            device->DeallocateMemory(mMemoryAllocation);
-                // freed after the VkImage is destroyed and this is taken care of by the
-                // FencedDeleter.
-                device->GetMemoryAllocator()->Free(&mMemoryAllocation);
-            }
 
             if (mHandle != VK_NULL_HANDLE) {
                 device->GetFencedDeleter()->DeleteWhenUnused(mHandle);

src/dawn_native/vulkan/TextureVk.h

@@ -18,12 +18,13 @@
 #include "dawn_native/Texture.h"
 
 #include "common/vulkan_platform.h"
+#include "dawn_native/ResourceMemoryAllocation.h"
 #include "dawn_native/vulkan/ExternalHandle.h"
-#include "dawn_native/vulkan/MemoryAllocator.h"
 
 namespace dawn_native { namespace vulkan {
 
     struct CommandRecordingContext;
+    class Device;
     struct ExternalImageDescriptor;
 
     VkFormat VulkanImageFormat(wgpu::TextureFormat format);
@@ -85,7 +86,7 @@ namespace dawn_native { namespace vulkan {
                                 TextureBase::ClearValue);
 
         VkImage mHandle = VK_NULL_HANDLE;
-        DeviceMemoryAllocation mMemoryAllocation;
+        ResourceMemoryAllocation mMemoryAllocation;
         VkDeviceMemory mExternalAllocation = VK_NULL_HANDLE;
 
         enum class ExternalState {

src/tests/white_box/VulkanImageWrappingTests.cpp

@@ -20,7 +20,7 @@
 #include "dawn_native/vulkan/AdapterVk.h"
 #include "dawn_native/vulkan/DeviceVk.h"
 #include "dawn_native/vulkan/FencedDeleter.h"
-#include "dawn_native/vulkan/MemoryAllocator.h"
+#include "dawn_native/vulkan/ResourceMemoryAllocatorVk.h"
 #include "dawn_native/vulkan/TextureVk.h"
 #include "utils/DawnHelpers.h"
 #include "utils/SystemUtils.h"
@@ -89,7 +89,8 @@ namespace {
             externalInfo.pNext = nullptr;
             externalInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;
 
-            int bestType = deviceVk->GetMemoryAllocator()->FindBestTypeIndex(requirements, false);
+            int bestType = deviceVk->GetResourceMemoryAllocatorForTesting()->FindBestTypeIndex(
+                requirements, false);
             VkMemoryAllocateInfo allocateInfo;
             allocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
             allocateInfo.pNext = &externalInfo;