2 files changed, 91 insertions, 55 deletions

diff --git a/src/video_core/vulkan_common/vulkan_memory_allocator.cpp b/src/video_core/vulkan_common/vulkan_memory_allocator.cpp
index 8bb15794d..f15061d0c 100644
--- a/src/video_core/vulkan_common/vulkan_memory_allocator.cpp
+++ b/src/video_core/vulkan_common/vulkan_memory_allocator.cpp
@@ -5,7 +5,6 @@
 #include <algorithm>
 #include <bit>
 #include <optional>
-#include <tuple>
 #include <vector>
 
 #include "common/alignment.h"
@@ -27,7 +26,7 @@ struct Range {
     }
 };
 
-[[nodiscard]] u64 GetAllocationChunkSize(u64 required_size) {
+[[nodiscard]] u64 AllocationChunkSize(u64 required_size) {
     static constexpr std::array sizes{
         0x1000ULL << 10,  0x1400ULL << 10,  0x1800ULL << 10,  0x1c00ULL << 10, 0x2000ULL << 10,
         0x3200ULL << 10,  0x4000ULL << 10,  0x6000ULL << 10,  0x8000ULL << 10, 0xA000ULL << 10,
@@ -38,14 +37,28 @@ struct Range {
     const auto it = std::ranges::lower_bound(sizes, required_size);
     return it != sizes.end() ? *it : Common::AlignUp(required_size, 4ULL << 20);
 }
+
+[[nodiscard]] VkMemoryPropertyFlags MemoryUsagePropertyFlags(MemoryUsage usage) {
+    switch (usage) {
+    case MemoryUsage::DeviceLocal:
+        return VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
+    case MemoryUsage::Upload:
+        return VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
+    case MemoryUsage::Download:
+        return VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT |
+               VK_MEMORY_PROPERTY_HOST_CACHED_BIT;
+    }
+    UNREACHABLE_MSG("Invalid memory usage={}", usage);
+    return VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
+}
 } // Anonymous namespace
 
 class MemoryAllocation {
 public:
     explicit MemoryAllocation(const Device& device_, vk::DeviceMemory memory_,
-                              VkMemoryPropertyFlags properties_, u64 allocation_size_, u32 type_)
-        : device{device_}, memory{std::move(memory_)}, properties{properties_},
-          allocation_size{allocation_size_}, shifted_type{ShiftType(type_)} {}
+                              VkMemoryPropertyFlags properties, u64 allocation_size_, u32 type)
+        : device{device_}, memory{std::move(memory_)}, allocation_size{allocation_size_},
+          property_flags{properties}, shifted_memory_type{1U << type} {}
 
     [[nodiscard]] std::optional<MemoryCommit> Commit(VkDeviceSize size, VkDeviceSize alignment) {
         const std::optional<u64> alloc = FindFreeRegion(size, alignment);
@@ -68,17 +81,16 @@ public:
     }
 
     [[nodiscard]] std::span<u8> Map() {
-        if (!memory_mapped_span.empty()) {
-            return memory_mapped_span;
+        if (memory_mapped_span.empty()) {
+            u8* const raw_pointer = memory.Map(0, allocation_size);
+            memory_mapped_span = std::span<u8>(raw_pointer, allocation_size);
         }
-        u8* const raw_pointer = memory.Map(0, allocation_size);
-        memory_mapped_span = std::span<u8>(raw_pointer, allocation_size);
         return memory_mapped_span;
     }
 
     /// Returns whether this allocation is compatible with the arguments.
-    [[nodiscard]] bool IsCompatible(VkMemoryPropertyFlags wanted_properties, u32 type_mask) const {
-        return (wanted_properties & properties) && (type_mask & shifted_type) != 0;
+    [[nodiscard]] bool IsCompatible(VkMemoryPropertyFlags flags, u32 type_mask) const {
+        return (flags & property_flags) && (type_mask & shifted_memory_type) != 0;
     }
 
 private:
@@ -106,13 +118,13 @@ private:
         return candidate;
     }
 
-    const Device& device;                   ///< Vulkan device.
-    const vk::DeviceMemory memory;          ///< Vulkan memory allocation handler.
-    const VkMemoryPropertyFlags properties; ///< Vulkan properties.
-    const u64 allocation_size;              ///< Size of this allocation.
-    const u32 shifted_type;                 ///< Stored Vulkan type of this allocation, shifted.
-    std::vector<Range> commits;             ///< All commit ranges done from this allocation.
-    std::span<u8> memory_mapped_span;       ///< Memory mapped span. Empty if not queried before.
+    const Device& device;                       ///< Vulkan device.
+    const vk::DeviceMemory memory;              ///< Vulkan memory allocation handler.
+    const u64 allocation_size;                  ///< Size of this allocation.
+    const VkMemoryPropertyFlags property_flags; ///< Vulkan memory property flags.
+    const u32 shifted_memory_type;              ///< Shifted Vulkan memory type.
+    std::vector<Range> commits;                 ///< All commit ranges done from this allocation.
+    std::span<u8> memory_mapped_span; ///< Memory mapped span. Empty if not queried before.
 };
 
 MemoryCommit::MemoryCommit(const Device& device_, MemoryAllocation* allocation_,
@@ -138,10 +150,9 @@ MemoryCommit::MemoryCommit(MemoryCommit&& rhs) noexcept
       interval{rhs.interval}, span{std::exchange(rhs.span, std::span<u8>{})} {}
 
 std::span<u8> MemoryCommit::Map() {
-    if (!span.empty()) {
-        return span;
+    if (span.empty()) {
+        span = allocation->Map().subspan(interval.first, interval.second - interval.first);
     }
-    span = allocation->Map().subspan(interval.first, interval.second - interval.first);
     return span;
 }
 
@@ -157,25 +168,18 @@ MemoryAllocator::MemoryAllocator(const Device& device_)
 MemoryAllocator::~MemoryAllocator() = default;
 
 MemoryCommit MemoryAllocator::Commit(const VkMemoryRequirements& requirements, MemoryUsage usage) {
-    const u64 chunk_size = GetAllocationChunkSize(requirements.size);
-
-    // When a host visible commit is asked, search for host visible and coherent, otherwise search
-    // for a fast device local type.
-    // TODO: Deduce memory types from usage in a better way
-    const bool host_visible = IsHostVisible(usage);
-    const VkMemoryPropertyFlags wanted_properties =
-        host_visible ? VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
-                     : VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
-    if (std::optional<MemoryCommit> commit = TryAllocCommit(requirements, wanted_properties)) {
+    // Find the fastest memory flags we can afford with the current requirements
+    const VkMemoryPropertyFlags flags = MemoryPropertyFlags(requirements.memoryTypeBits, usage);
+    if (std::optional<MemoryCommit> commit = TryCommit(requirements, flags)) {
         return std::move(*commit);
     }
     // Commit has failed, allocate more memory.
     // TODO(Rodrigo): Handle out of memory situations in some way like flushing to guest memory.
-    AllocMemory(wanted_properties, requirements.memoryTypeBits, chunk_size);
+    AllocMemory(flags, requirements.memoryTypeBits, AllocationChunkSize(requirements.size));
 
     // Commit again, this time it won't fail since there's a fresh allocation above.
     // If it does, there's a bug.
-    return TryAllocCommit(requirements, wanted_properties).value();
+    return TryCommit(requirements, flags).value();
 }
 
 MemoryCommit MemoryAllocator::Commit(const vk::Buffer& buffer, MemoryUsage usage) {
@@ -190,33 +194,22 @@ MemoryCommit MemoryAllocator::Commit(const vk::Image& image, MemoryUsage usage)
     return commit;
 }
 
-void MemoryAllocator::AllocMemory(VkMemoryPropertyFlags wanted_properties, u32 type_mask,
-                                  u64 size) {
-    const u32 type = [&] {
-        for (u32 type_index = 0; type_index < properties.memoryTypeCount; ++type_index) {
-            const auto flags = properties.memoryTypes[type_index].propertyFlags;
-            if ((type_mask & (1U << type_index)) && (flags & wanted_properties)) {
-                // The type matches in type and in the wanted properties.
-                return type_index;
-            }
-        }
-        UNREACHABLE_MSG("Couldn't find a compatible memory type!");
-        return 0U;
-    }();
+void MemoryAllocator::AllocMemory(VkMemoryPropertyFlags flags, u32 type_mask, u64 size) {
+    const u32 type = FindType(flags, type_mask).value();
     vk::DeviceMemory memory = device.GetLogical().AllocateMemory({
         .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
         .pNext = nullptr,
         .allocationSize = size,
         .memoryTypeIndex = type,
     });
-    allocations.push_back(std::make_unique<MemoryAllocation>(device, std::move(memory),
-                                                             wanted_properties, size, type));
+    allocations.push_back(
+        std::make_unique<MemoryAllocation>(device, std::move(memory), flags, size, type));
 }
 
-std::optional<MemoryCommit> MemoryAllocator::TryAllocCommit(
-    const VkMemoryRequirements& requirements, VkMemoryPropertyFlags wanted_properties) {
+std::optional<MemoryCommit> MemoryAllocator::TryCommit(const VkMemoryRequirements& requirements,
+                                                       VkMemoryPropertyFlags flags) {
     for (auto& allocation : allocations) {
-        if (!allocation->IsCompatible(wanted_properties, requirements.memoryTypeBits)) {
+        if (!allocation->IsCompatible(flags, requirements.memoryTypeBits)) {
             continue;
         }
         if (auto commit = allocation->Commit(requirements.size, requirements.alignment)) {
@@ -226,6 +219,40 @@ std::optional<MemoryCommit> MemoryAllocator::TryAllocCommit(
     return std::nullopt;
 }
 
+VkMemoryPropertyFlags MemoryAllocator::MemoryPropertyFlags(u32 type_mask, MemoryUsage usage) const {
+    return MemoryPropertyFlags(type_mask, MemoryUsagePropertyFlags(usage));
+}
+
+VkMemoryPropertyFlags MemoryAllocator::MemoryPropertyFlags(u32 type_mask,
+                                                           VkMemoryPropertyFlags flags) const {
+    if (FindType(flags, type_mask)) {
+        // Found a memory type with those requirements
+        return flags;
+    }
+    if (flags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT) {
+        // Remove host cached bit in case it's not supported
+        return MemoryPropertyFlags(type_mask, flags & ~VK_MEMORY_PROPERTY_HOST_CACHED_BIT);
+    }
+    if (flags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) {
+        // Remove device local, if it's not supported by the requested resource
+        return MemoryPropertyFlags(type_mask, flags & ~VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
+    }
+    UNREACHABLE_MSG("No compatible memory types found");
+    return 0;
+}
+
+std::optional<u32> MemoryAllocator::FindType(VkMemoryPropertyFlags flags, u32 type_mask) const {
+    for (u32 type_index = 0; type_index < properties.memoryTypeCount; ++type_index) {
+        const VkMemoryPropertyFlags type_flags = properties.memoryTypes[type_index].propertyFlags;
+        if ((type_mask & (1U << type_index)) && (type_flags & flags)) {
+            // The type matches in type and in the wanted properties.
+            return type_index;
+        }
+    }
+    // Failed to find index
+    return std::nullopt;
+}
+
 bool IsHostVisible(MemoryUsage usage) noexcept {
     switch (usage) {
     case MemoryUsage::DeviceLocal:
diff --git a/src/video_core/vulkan_common/vulkan_memory_allocator.h b/src/video_core/vulkan_common/vulkan_memory_allocator.h
index efb32167a..d4e34c843 100644
--- a/src/video_core/vulkan_common/vulkan_memory_allocator.h
+++ b/src/video_core/vulkan_common/vulkan_memory_allocator.h
@@ -92,13 +92,22 @@ public:
 
 private:
     /// Allocates a chunk of memory.
-    void AllocMemory(VkMemoryPropertyFlags wanted_properties, u32 type_mask, u64 size);
+    void AllocMemory(VkMemoryPropertyFlags flags, u32 type_mask, u64 size);
 
     /// Tries to allocate a memory commit.
-    std::optional<MemoryCommit> TryAllocCommit(const VkMemoryRequirements& requirements,
-                                               VkMemoryPropertyFlags wanted_properties);
+    std::optional<MemoryCommit> TryCommit(const VkMemoryRequirements& requirements,
+                                          VkMemoryPropertyFlags flags);
 
-    const Device& device;                                       ///< Device handler.
+    /// Returns the fastest compatible memory property flags from a wanted usage.
+    VkMemoryPropertyFlags MemoryPropertyFlags(u32 type_mask, MemoryUsage usage) const;
+
+    /// Returns the fastest compatible memory property flags from the wanted flags.
+    VkMemoryPropertyFlags MemoryPropertyFlags(u32 type_mask, VkMemoryPropertyFlags flags) const;
+
+    /// Returns index to the fastest memory type compatible with the passed requirements.
+    std::optional<u32> FindType(VkMemoryPropertyFlags flags, u32 type_mask) const;
+
+    const Device& device;                                       ///< Device handle.
     const VkPhysicalDeviceMemoryProperties properties;          ///< Physical device properties.
     std::vector<std::unique_ptr<MemoryAllocation>> allocations; ///< Current allocations.
 };