Fix undersized buffer for Vulkan pixel history with many fragments

Before, the buffer could be overrun, which could result in anything from
garbage data to the GPU device being lost to segfaults. Now, correct
data is gathered.

We can't know in advance how many primitives will hit the targetted
pixel, so it's not possible to create the buffer until after our first
pass for each events (fortunately we do know the number of events in
advance). It is possible that we don't need a larger buffer, though,
in which case the original one can be re-used.

renderdoc/driver/vulkan/vk_debug.h

@@ -97,6 +97,8 @@ public:
   bool PixelHistorySetupResources(PixelHistoryResources &resources, VkImage targetImage,
                                   VkExtent3D extent, VkFormat format, VkSampleCountFlagBits samples,
                                   const Subresource &sub, uint32_t numEvents);
+  bool PixelHistorySetupPerFragResources(PixelHistoryResources &resources, uint32_t numEvents,
+                                         uint32_t numFragments);
   bool PixelHistoryDestroyResources(const PixelHistoryResources &resources);
 
   void PixelHistoryCopyPixel(VkCommandBuffer cmd, VkCopyPixelParams &p, size_t offset);

renderdoc/driver/vulkan/vk_pixelhistory.cpp

@@ -3389,9 +3389,9 @@ bool VulkanDebugManager::PixelHistorySetupResources(PixelHistoryResources &resou
                                                     VkFormat format, VkSampleCountFlagBits samples,
                                                     const Subresource &sub, uint32_t numEvents)
 {
-  VkMarkerRegion region(StringFormat::Fmt("PixelHistorySetupResources %ux%ux%u %s %ux MSAA",
-                                          extent.width, extent.height, extent.depth,
-                                          ToStr(format).c_str(), samples));
+  VkMarkerRegion region(
+      StringFormat::Fmt("PixelHistorySetupResources %ux%ux%u %s %ux MSAA, %u events", extent.width,
+                        extent.height, extent.depth, ToStr(format).c_str(), samples, numEvents));
   VulkanCreationInfo::Image targetImageInfo = GetImageInfo(GetResID(targetImage));
 
   VkImage colorImage;
@@ -3501,8 +3501,6 @@ bool VulkanDebugManager::PixelHistorySetupResources(PixelHistoryResources &resou
   CheckVkResult(vkr);
 
   VkBufferCreateInfo bufferInfo = {VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO};
-  // TODO: the size for memory is calculated to fit pre and post modification values and
-  // stencil values. But we might run out of space when getting per fragment data.
   bufferInfo.size = AlignUp((uint32_t)(numEvents * sizeof(EventInfo)), 4096U);
   bufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
 
@@ -3558,6 +3556,75 @@ bool VulkanDebugManager::PixelHistorySetupResources(PixelHistoryResources &resou
   return true;
 }
 
+bool VulkanDebugManager::PixelHistorySetupPerFragResources(PixelHistoryResources &resources,
+                                                           uint32_t numEvents, uint32_t numFrags)
+{
+  const uint32_t existingBufferSize = AlignUp((uint32_t)(numEvents * sizeof(EventInfo)), 4096U);
+  const uint32_t requiredBufferSize = AlignUp((uint32_t)(numFrags * sizeof(PerFragmentInfo)), 4096U);
+
+  // If the existing buffer is big enough for all of the fragments, we can re-use it.
+  const bool canReuseBuffer = existingBufferSize >= requiredBufferSize;
+
+  VkMarkerRegion region(StringFormat::Fmt(
+      "PixelHistorySetupPerFragResources %u events %u frags, buffer size %u -> %u, %s old buffer",
+      numEvents, numFrags, existingBufferSize, requiredBufferSize,
+      canReuseBuffer ? "reusing" : "NOT reusing"));
+
+  if(canReuseBuffer)
+    return true;
+
+  // Otherwise, destroy it and create a new one that's big enough in its place.
+  VkDevice dev = m_pDriver->GetDev();
+
+  if(resources.dstBuffer != VK_NULL_HANDLE)
+    m_pDriver->vkDestroyBuffer(dev, resources.dstBuffer, NULL);
+  if(resources.bufferMemory != VK_NULL_HANDLE)
+    m_pDriver->vkFreeMemory(dev, resources.bufferMemory, NULL);
+  resources.dstBuffer = VK_NULL_HANDLE;
+  resources.bufferMemory = VK_NULL_HANDLE;
+
+  VkBufferCreateInfo bufferInfo = {VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO};
+  bufferInfo.size = requiredBufferSize;
+  bufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
+
+  VkResult vkr = m_pDriver->vkCreateBuffer(m_Device, &bufferInfo, NULL, &resources.dstBuffer);
+  CheckVkResult(vkr);
+
+  // Allocate memory
+  VkMemoryRequirements mrq = {};
+  m_pDriver->vkGetBufferMemoryRequirements(m_Device, resources.dstBuffer, &mrq);
+  VkMemoryAllocateInfo allocInfo = {
+      VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, NULL, mrq.size,
+      m_pDriver->GetReadbackMemoryIndex(mrq.memoryTypeBits),
+  };
+  vkr = m_pDriver->vkAllocateMemory(m_Device, &allocInfo, NULL, &resources.bufferMemory);
+  CheckVkResult(vkr);
+
+  if(vkr != VK_SUCCESS)
+    return false;
+
+  vkr = m_pDriver->vkBindBufferMemory(m_Device, resources.dstBuffer, resources.bufferMemory, 0);
+  CheckVkResult(vkr);
+
+  VkCommandBuffer cmd = m_pDriver->GetNextCmd();
+  VkCommandBufferBeginInfo beginInfo = {VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, NULL,
+                                        VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT};
+
+  if(cmd == VK_NULL_HANDLE)
+    return false;
+
+  vkr = ObjDisp(dev)->BeginCommandBuffer(Unwrap(cmd), &beginInfo);
+  CheckVkResult(vkr);
+  ObjDisp(cmd)->CmdFillBuffer(Unwrap(cmd), Unwrap(resources.dstBuffer), 0, VK_WHOLE_SIZE, 0);
+
+  vkr = ObjDisp(dev)->EndCommandBuffer(Unwrap(cmd));
+  CheckVkResult(vkr);
+  m_pDriver->SubmitCmds();
+  m_pDriver->FlushQ();
+
+  return true;
+}
+
 VkDescriptorSet VulkanReplay::GetPixelHistoryDescriptor()
 {
   VkDescriptorSet descSet;
@@ -4057,6 +4124,17 @@ rdcarray<PixelModification> VulkanReplay::PixelHistory(rdcarray<EventUsage> even
 
   if(eventsWithFrags.size() > 0)
   {
+    uint32_t numFrags = 0;
+    for(auto &item : eventsWithFrags)
+    {
+      numFrags += item.second;
+    }
+
+    GetDebugManager()->PixelHistorySetupPerFragResources(resources, (uint32_t)events.size(),
+                                                         numFrags);
+
+    callbackInfo.dstBuffer = resources.dstBuffer;
+
     // Replay to get shader output value, post modification value and primitive ID for every
     // fragment.
     VulkanPixelHistoryPerFragmentCallback perFragmentCB(m_pDriver, shaderCache, callbackInfo,
@@ -4222,6 +4300,8 @@ rdcarray<PixelModification> VulkanReplay::PixelHistory(rdcarray<EventUsage> even
           history[h].depthBoundsFailed = true;
       }
     }
+
+    m_pDriver->vkUnmapMemory(dev, resources.bufferMemory);
   }
 
   SAFE_DELETE(tfCb);