Add sparse image and buffer support

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,107 @@ class VulkanCaptureManager : public ApiCaptureManager
                                                           pBindInfo[i].signalSemaphoreCount,
                                                           pBindInfo[i].pSignalSemaphores);
             }
+
+            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!");
+                            }
+                        }
+                    }
+                }
+            }
         }
     }
 
@@ -804,6 +905,50 @@ class VulkanCaptureManager : public ApiCaptureManager
                                         uint32_t                     bindInfoCount,
                                         const VkBindImageMemoryInfo* pBindInfos);
 
+    void PostProcess_vkCreateBuffer(VkResult                     result,
+                                    VkDevice                     device,
+                                    const VkBufferCreateInfo*    pCreateInfo,
+                                    const VkAllocationCallbacks* pAllocator,
+                                    VkBuffer*                    pBuffer)
+    {
+        if (IsCaptureModeTrack() && (result == VK_SUCCESS) && (pCreateInfo != nullptr))
+        {
+            GFXRECON_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))
+        {
+            GFXRECON_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)
@@ -1624,6 +1769,11 @@ class VulkanCaptureManager : public ApiCaptureManager
     std::unique_ptr<VulkanStateTracker>             state_tracker_;
     HardwareBufferMap                               hardware_buffers_;
     std::mutex                                      deferred_operation_mutex;
+
+    // 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.
+    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;
@@ -227,6 +241,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))
     {