Add sparse image and buffer support

The current trim handling for buffers and images does not support sparse
buffers or sparse images. It only allows one image or buffer binding to
a single memory range of a single memory object. As a result, it cannot
manage sparse resources, which has led to missing asset issues in trim
trace for some titles that utilize sparse resources. This commit
introduces support for sparse buffers and sparse images (using opaque
memory binding only), effectively resolving the missing asset issue.

Originally by Ming Zheng <[email protected]>

Author: bradgrantham-lunarg

framework/encode/custom_vulkan_encoder_commands.h

@@ -505,6 +505,26 @@ struct CustomEncoderPostCall<format::ApiCallId::ApiCall_vkBindImageMemory2KHR>
     }
 };
 
+template <>
+struct CustomEncoderPostCall<format::ApiCallId::ApiCall_vkCreateBuffer>
+{
+    template <typename... Args>
+    static void Dispatch(VulkanCaptureManager* manager, VkResult result, Args... args)
+    {
+        manager->PostProcess_vkCreateBuffer(result, args...);
+    }
+};
+
+template <>
+struct CustomEncoderPostCall<format::ApiCallId::ApiCall_vkCreateImage>
+{
+    template <typename... Args>
+    static void Dispatch(VulkanCaptureManager* manager, VkResult result, Args... args)
+    {
+        manager->PostProcess_vkCreateImage(result, args...);
+    }
+};
+
 template <>
 struct CustomEncoderPostCall<format::ApiCallId::ApiCall_vkCmdBeginRenderPass>
 {

framework/encode/vulkan_capture_manager.h

@@ -551,7 +551,7 @@ class VulkanCaptureManager : public ApiCaptureManager
     }
 
     void PostProcess_vkQueueBindSparse(
-        VkResult result, VkQueue, uint32_t bindInfoCount, const VkBindSparseInfo* pBindInfo, VkFence)
+        VkResult result, VkQueue queue, uint32_t bindInfoCount, const VkBindSparseInfo* pBindInfo, VkFence)
     {
         if (IsCaptureModeTrack() && (result == VK_SUCCESS))
         {
@@ -563,6 +563,110 @@ class VulkanCaptureManager : public ApiCaptureManager
                                                           pBindInfo[i].signalSemaphoreCount,
                                                           pBindInfo[i].pSignalSemaphores);
             }
+
+            // In default mode, the capture manager uses a shared mutex to capture every API function. As a result,
+            // multiple threads may access the sparse resource maps concurrently. Therefore, we use a dedicated mutex
+            // for write access to these maps.
+            const std::lock_guard<std::mutex> lock(sparse_resource_mutex);
+            for (uint32_t bind_info_index = 0; bind_info_index < bindInfoCount; bind_info_index++)
+            {
+                auto& bind_info = pBindInfo[bind_info_index];
+
+                // TODO: add device group support. In the following handling, we assume that the system only has one
+                // physical device or that resourceDeviceIndex and memoryDeviceIndex of VkDeviceGroupBindSparseInfo in
+                // the pnext chain are zero.
+
+                if (bind_info.pBufferBinds != nullptr)
+                {
+                    // The title binds sparse buffers to memory ranges, so we need to track the buffer binding
+                    // information. The following updates will reflect the latest binding states for all buffers in this
+                    // vkQueueBindSparse command, covering both fully-resident and partially-resident buffers.
+                    for (uint32_t buffer_bind_index = 0; buffer_bind_index < bind_info.bufferBindCount;
+                         buffer_bind_index++)
+                    {
+                        auto& buffer_bind   = bind_info.pBufferBinds[buffer_bind_index];
+                        auto  sparse_buffer = buffer_bind.buffer;
+                        auto  wrapper = vulkan_wrappers::GetWrapper<vulkan_wrappers::BufferWrapper>(sparse_buffer);
+
+                        if (wrapper != nullptr)
+                        {
+                            wrapper->sparse_bind_queue = queue;
+                            for (uint32_t bind_memory_range_index = 0; bind_memory_range_index < buffer_bind.bindCount;
+                                 bind_memory_range_index++)
+                            {
+                                auto& bind_memory_range = buffer_bind.pBinds[bind_memory_range_index];
+                                graphics::UpdateSparseMemoryBindMap(wrapper->sparse_memory_bind_map, bind_memory_range);
+                            }
+                        }
+                    }
+                }
+
+                if (bind_info.pImageOpaqueBinds != nullptr)
+                {
+                    // The title binds sparse images to opaque memory ranges, so we need to track the image binding
+                    // information. The following handling will update the latest binding states for all images in this
+                    // vkQueueBindSparse command, which utilizes opaque memory binding. There are two cases covered by
+                    // the tracking. In the first case, the sparse image exclusively uses opaque memory binding. For
+                    // this case, the target title treats the binding memory ranges as a linear unified region. This
+                    // should represent a fully-resident binding because this linear region is entirely opaque, meaning
+                    // there is no application-visible mapping between texel locations and memory offsets. In another
+                    // case, the image utilizes subresource sparse memory binding, just binding only its mip tail region
+                    // to an opaque memory range. For this situation, we use the sparse_opaque_memory_bind_map and
+                    // sparse_subresource_memory_bind_map of the image wrapper to track the subresource bindings and
+                    // opaque bindings separately.
+                    for (uint32_t image_opaque_bind_index = 0; image_opaque_bind_index < bind_info.imageOpaqueBindCount;
+                         image_opaque_bind_index++)
+                    {
+                        auto& image_opaque_bind = bind_info.pImageOpaqueBinds[image_opaque_bind_index];
+                        auto  sparse_image      = image_opaque_bind.image;
+                        auto  wrapper = vulkan_wrappers::GetWrapper<vulkan_wrappers::ImageWrapper>(sparse_image);
+
+                        if (wrapper != nullptr)
+                        {
+                            wrapper->sparse_bind_queue = queue;
+
+                            for (uint32_t bind_memory_range_index = 0;
+                                 bind_memory_range_index < image_opaque_bind.bindCount;
+                                 bind_memory_range_index++)
+                            {
+                                auto& bind_memory_range = image_opaque_bind.pBinds[bind_memory_range_index];
+                                graphics::UpdateSparseMemoryBindMap(wrapper->sparse_opaque_memory_bind_map,
+                                                                    bind_memory_range);
+                            }
+                        }
+                    }
+                }
+
+                if (bind_info.pImageBinds != nullptr)
+                {
+                    // The title binds subresources of a sparse image to memory ranges, which requires us to keep track
+                    // of the sparse image subresource binding information. It's important to note that while the image
+                    // mainly use subresource sparse memory binding, its mip tail region must be bound to an opaque
+                    // memory range. Therefore, we use the sparse_opaque_memory_bind_map and
+                    // sparse_subresource_memory_bind_map of the image wrapper to separately track both the
+                    // subresource bindings and the opaque bindings.
+                    for (uint32_t image_bind_index = 0; image_bind_index < bind_info.imageBindCount; image_bind_index++)
+                    {
+                        auto& image_bind   = bind_info.pImageBinds[image_bind_index];
+                        auto  sparse_image = image_bind.image;
+                        auto  wrapper      = vulkan_wrappers::GetWrapper<vulkan_wrappers::ImageWrapper>(sparse_image);
+
+                        if (wrapper != nullptr)
+                        {
+                            wrapper->sparse_bind_queue = queue;
+
+                            for (uint32_t bind_memory_range_index = 0; bind_memory_range_index < image_bind.bindCount;
+                                 bind_memory_range_index++)
+                            {
+                                auto& bind_memory_range = image_bind.pBinds[bind_memory_range_index];
+                                // TODO: Implement handling for tracking binding information of sparse image
+                                // subresources.
+                                GFXRECON_LOG_ERROR_ONCE("Binding of sparse image blocks is not supported!");
+                            }
+                        }
+                    }
+                }
+            }
         }
     }
 
@@ -842,6 +946,50 @@ class VulkanCaptureManager : public ApiCaptureManager
         }
     }
 
+    void PostProcess_vkCreateBuffer(VkResult                     result,
+                                    VkDevice                     device,
+                                    const VkBufferCreateInfo*    pCreateInfo,
+                                    const VkAllocationCallbacks* pAllocator,
+                                    VkBuffer*                    pBuffer)
+    {
+        if (IsCaptureModeTrack() && (result == VK_SUCCESS) && (pCreateInfo != nullptr))
+        {
+            assert(state_tracker_ != nullptr);
+
+            auto buffer_wrapper = vulkan_wrappers::GetWrapper<vulkan_wrappers::BufferWrapper>(*pBuffer);
+
+            if (buffer_wrapper->is_sparse_buffer)
+            {
+                // We will need to set the bind_device for handling sparse buffers. There will be no subsequent
+                // vkBindBufferMemory, vkBindBufferMemory2 or vkBindBufferMemory2KHR calls for sparse buffer, so we
+                // assign bind_device to the device that created the buffer.
+                buffer_wrapper->bind_device = vulkan_wrappers::GetWrapper<vulkan_wrappers::DeviceWrapper>(device);
+            }
+        }
+    }
+
+    void PostProcess_vkCreateImage(VkResult                     result,
+                                   VkDevice                     device,
+                                   const VkImageCreateInfo*     pCreateInfo,
+                                   const VkAllocationCallbacks* pAllocator,
+                                   VkImage*                     pImage)
+    {
+        if (IsCaptureModeTrack() && (result == VK_SUCCESS) && (pCreateInfo != nullptr))
+        {
+            assert(state_tracker_ != nullptr);
+
+            auto image_wrapper = vulkan_wrappers::GetWrapper<vulkan_wrappers::ImageWrapper>(*pImage);
+
+            if (image_wrapper->is_sparse_image)
+            {
+                // We will need to set the bind_device for handling sparse images. There will be no subsequent
+                // vkBindImageMemory, vkBindImageMemory2, or vkBindImageMemory2KHR calls for sparse image, so we assign
+                // bind_device to the device that created the image.
+                image_wrapper->bind_device = vulkan_wrappers::GetWrapper<vulkan_wrappers::DeviceWrapper>(device);
+            }
+        }
+    }
+
     void PostProcess_vkCmdBeginRenderPass(VkCommandBuffer              commandBuffer,
                                           const VkRenderPassBeginInfo* pRenderPassBegin,
                                           VkSubpassContents)
@@ -1662,6 +1810,7 @@ class VulkanCaptureManager : public ApiCaptureManager
     std::unique_ptr<VulkanStateTracker>             state_tracker_;
     HardwareBufferMap                               hardware_buffers_;
     std::mutex                                      deferred_operation_mutex;
+    std::mutex                                      sparse_resource_mutex;
 };
 
 GFXRECON_END_NAMESPACE(encode)

framework/encode/vulkan_handle_wrappers.h

@@ -30,6 +30,7 @@
 #include "format/format.h"
 #include "generated/generated_vulkan_dispatch_table.h"
 #include "graphics/vulkan_device_util.h"
+#include "graphics/vulkan_resources_util.h"
 #include "util/defines.h"
 #include "util/memory_output_stream.h"
 #include "util/page_guard_manager.h"
@@ -210,6 +211,19 @@ struct BufferWrapper : public HandleWrapper<VkBuffer>, AssetWrapperBase
     VkBufferUsageFlags usage{ 0 };
 
     std::set<BufferViewWrapper*> buffer_views;
+    
+    DeviceWrapper*             bind_device{ nullptr };
+    const void*                bind_pnext{ nullptr };
+    std::unique_ptr<uint8_t[]> bind_pnext_memory;
+
+    bool                                       is_sparse_buffer{ false };
+    std::map<VkDeviceSize, VkSparseMemoryBind> sparse_memory_bind_map;
+    VkQueue                                    sparse_bind_queue;
+
+    format::HandleId bind_memory_id{ format::kNullHandleId };
+    VkDeviceSize     bind_offset{ 0 };
+    uint32_t         queue_family_index{ 0 };
+    VkDeviceSize     created_size{ 0 };
 };
 
 struct ImageViewWrapper;
@@ -226,6 +240,20 @@ struct ImageWrapper : public HandleWrapper<VkImage>, AssetWrapperBase
     bool                     is_swapchain_image{ false };
 
     std::set<ImageViewWrapper*> image_views;
+
+    DeviceWrapper*             bind_device{ nullptr };
+    const void*                bind_pnext{ nullptr };
+    std::unique_ptr<uint8_t[]> bind_pnext_memory;
+
+    bool                                                is_sparse_image{ false };
+    std::map<VkDeviceSize, VkSparseMemoryBind>          sparse_opaque_memory_bind_map;
+    graphics::VulkanSubresourceSparseImageMemoryBindMap sparse_subresource_memory_bind_map;
+    VkQueue                                             sparse_bind_queue;
+
+    format::HandleId         bind_memory_id{ format::kNullHandleId };
+    VkDeviceSize             bind_offset{ 0 };
+    uint32_t                 queue_family_index{ 0 };
+    std::set<VkSwapchainKHR> parent_swapchains;
 };
 
 struct SamplerWrapper : public HandleWrapper<VkSampler>

framework/encode/vulkan_state_tracker_initializers.h

@@ -608,6 +608,13 @@ inline void InitializeState<VkDevice, vulkan_wrappers::BufferWrapper, VkBufferCr
     wrapper->create_call_id    = create_call_id;
     wrapper->create_parameters = std::move(create_parameters);
 
+    wrapper->created_size = create_info->size;
+
+    if ((create_info->flags & VK_BUFFER_CREATE_SPARSE_BINDING_BIT) != 0)
+    {
+        wrapper->is_sparse_buffer = true;
+    }
+
     // TODO: Do we need to track the queue family that the buffer is actually used with?
     if ((create_info->queueFamilyIndexCount > 0) && (create_info->pQueueFamilyIndices != nullptr))
     {
@@ -641,6 +648,11 @@ inline void InitializeState<VkDevice, vulkan_wrappers::ImageWrapper, VkImageCrea
     wrapper->samples      = create_info->samples;
     wrapper->tiling       = create_info->tiling;
 
+    if ((create_info->flags & VK_IMAGE_CREATE_SPARSE_BINDING_BIT) != 0)
+    {
+        wrapper->is_sparse_image = true;
+    }
+
     // TODO: Do we need to track the queue family that the image is actually used with?
     if ((create_info->queueFamilyIndexCount > 0) && (create_info->pQueueFamilyIndices != nullptr))
     {