Record which queries are reset in each frame & restore them to validity

* It's possible that a frame capture could copy from a query and then reset it,
  without then recording something valid into the query. Then the next replay
  the copy would be unavailable and if WAIT_BIT is set that would lead to a
  device lost.
* We already fill out queries with dummy valid data on create time, so now we
  record any queries that are reset and re-fill them with valid data again.

renderdoc/driver/vulkan/vk_core.cpp

@@ -2696,6 +2696,65 @@ void WrappedVulkan::ApplyInitialContents()
   SubmitCmds();
   FlushQ();
   SubmitAndFlushImageStateBarriers(m_cleanupImageBarriers);
+
+  // reset any queries to a valid copy-able state if they need to be copied.
+  if(!m_ResetQueries.empty())
+  {
+    // sort all pools together
+    std::sort(m_ResetQueries.begin(), m_ResetQueries.end(),
+              [](const ResetQuery &a, const ResetQuery &b) { return a.pool < b.pool; });
+
+    cmd = GetNextCmd();
+
+    vkr = ObjDisp(cmd)->BeginCommandBuffer(Unwrap(cmd), &beginInfo);
+    RDCASSERTEQUAL(vkr, VK_SUCCESS);
+
+    uint32_t i = 0;
+    for(const ResetQuery &r : m_ResetQueries)
+    {
+      ObjDisp(cmd)->CmdResetQueryPool(Unwrap(cmd), Unwrap(r.pool), r.firstQuery, r.queryCount);
+
+      for(uint32_t q = 0; q < r.queryCount; q++)
+      {
+        // Timestamps are easy - we can do these without needing to render
+        if(m_CreationInfo.m_QueryPool[GetResID(r.pool)].queryType == VK_QUERY_TYPE_TIMESTAMP)
+        {
+          ObjDisp(cmd)->CmdWriteTimestamp(Unwrap(cmd), VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
+                                          Unwrap(r.pool), r.firstQuery + q);
+        }
+        else
+        {
+          ObjDisp(cmd)->CmdBeginQuery(Unwrap(cmd), Unwrap(r.pool), r.firstQuery + q, 0);
+          ObjDisp(cmd)->CmdEndQuery(Unwrap(cmd), Unwrap(r.pool), r.firstQuery + q);
+        }
+
+        i++;
+
+        // split the command buffer and flush if the number of queries is massive
+        if(i > 0 && (i % (128 * 1024)) == 0)
+        {
+          vkr = ObjDisp(cmd)->EndCommandBuffer(Unwrap(cmd));
+          RDCASSERTEQUAL(vkr, VK_SUCCESS);
+
+          SubmitCmds();
+          FlushQ();
+
+          cmd = GetNextCmd();
+
+          vkr = ObjDisp(cmd)->BeginCommandBuffer(Unwrap(cmd), &beginInfo);
+          RDCASSERTEQUAL(vkr, VK_SUCCESS);
+        }
+      }
+    }
+
+    vkr = ObjDisp(cmd)->EndCommandBuffer(Unwrap(cmd));
+    RDCASSERTEQUAL(vkr, VK_SUCCESS);
+
+    m_ResetQueries.clear();
+
+    SubmitCmds();
+    FlushQ();
+  }
 }
 
 bool WrappedVulkan::ContextProcessChunk(ReadSerialiser &ser, VulkanChunk chunk)

renderdoc/driver/vulkan/vk_core.h

@@ -551,6 +551,19 @@ private:
   rdcarray<VkEvent> m_CleanupEvents;
   rdcarray<VkEvent> m_PersistentEvents;
 
+  // reset queries must be restored to a valid state on next replay to ensure any query pool
+  // results copies that refer to previous frames have valid data. We initialised all of the
+  // query pool at creation time, but any queries that are reset might not be recorded within
+  // the frame - certainly not if we're doing partial replays. So we record all queries that are
+  // reset to re-initialise them.
+  struct ResetQuery
+  {
+    VkQueryPool pool;
+    uint32_t firstQuery;
+    uint32_t queryCount;
+  };
+  rdcarray<ResetQuery> m_ResetQueries;
+
   const VkFormatProperties &GetFormatProperties(VkFormat f);
 
   struct BakedCmdBufferInfo

renderdoc/driver/vulkan/vk_info.cpp

@@ -1264,6 +1264,14 @@ void VulkanCreationInfo::ShaderModuleReflection::PopulateDisassembly(const rdcsp
     disassembly = spirv.Disassemble(refl.entryPoint.c_str(), instructionLines);
 }
 
+void VulkanCreationInfo::QueryPool::Init(VulkanResourceManager *resourceMan, VulkanCreationInfo &info,
+                                         const VkQueryPoolCreateInfo *pCreateInfo)
+{
+  queryType = pCreateInfo->queryType;
+  queryCount = pCreateInfo->queryCount;
+  pipelineStatistics = pCreateInfo->pipelineStatistics;
+}
+
 void VulkanCreationInfo::DescSetPool::Init(VulkanResourceManager *resourceMan,
                                            VulkanCreationInfo &info,
                                            const VkDescriptorPoolCreateInfo *pCreateInfo)

renderdoc/driver/vulkan/vk_info.h

@@ -609,6 +609,17 @@ struct VulkanCreationInfo
   };
   std::map<ResourceId, DescSetPool> m_DescSetPool;
 
+  struct QueryPool
+  {
+    void Init(VulkanResourceManager *resourceMan, VulkanCreationInfo &info,
+              const VkQueryPoolCreateInfo *pCreateInfo);
+
+    VkQueryType queryType;
+    uint32_t queryCount;
+    VkQueryPipelineStatisticFlags pipelineStatistics;
+  };
+  std::map<ResourceId, QueryPool> m_QueryPool;
+
   std::map<ResourceId, rdcstr> m_Names;
   std::map<ResourceId, SwapchainInfo> m_SwapChain;
   std::map<ResourceId, DescSetLayout> m_DescSetLayout;
@@ -619,6 +630,7 @@ struct VulkanCreationInfo
 
   void erase(ResourceId id)
   {
+    m_QueryPool.erase(id);
     m_Pipeline.erase(id);
     m_PipelineLayout.erase(id);
     m_RenderPass.erase(id);

renderdoc/driver/vulkan/wrappers/vk_cmd_funcs.cpp

@@ -3781,6 +3781,8 @@ bool WrappedVulkan::Serialise_vkCmdResetQueryPool(SerialiserType &ser, VkCommand
     {
       ObjDisp(commandBuffer)
           ->CmdResetQueryPool(Unwrap(commandBuffer), Unwrap(queryPool), firstQuery, queryCount);
+
+      m_ResetQueries.push_back({queryPool, firstQuery, queryCount});
     }
   }
 

renderdoc/driver/vulkan/wrappers/vk_misc_funcs.cpp

@@ -1375,6 +1375,8 @@ bool WrappedVulkan::Serialise_vkCreateQueryPool(SerialiserType &ser, VkDevice de
       ResourceId live = GetResourceManager()->WrapResource(Unwrap(device), pool);
       GetResourceManager()->AddLiveResource(QueryPool, pool);
 
+      m_CreationInfo.m_QueryPool[live].Init(GetResourceManager(), m_CreationInfo, &CreateInfo);
+
       // We fill the query pool with valid but empty data, just so that future copies of query
       // results don't read from invalid data.
 

util/test/demos/vk/vk_query_pool.cpp

@@ -61,10 +61,9 @@ RD_TEST(VK_Query_Pool, VulkanGraphicsTest)
 
     VkPipeline pipe = createGraphicsPipeline(pipeCreateInfo);
 
-    AllocatedBuffer vb(
-        this, vkh::BufferCreateInfo(sizeof(DefaultTri), VK_BUFFER_USAGE_VERTEX_BUFFER_BIT |
-                                                            VK_BUFFER_USAGE_TRANSFER_DST_BIT),
-        VmaAllocationCreateInfo({0, VMA_MEMORY_USAGE_CPU_TO_GPU}));
+    AllocatedBuffer vb(this, vkh::BufferCreateInfo(4096, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT |
+                                                             VK_BUFFER_USAGE_TRANSFER_DST_BIT),
+                       VmaAllocationCreateInfo({0, VMA_MEMORY_USAGE_CPU_TO_GPU}));
 
     vb.upload(DefaultTri);
 
@@ -78,12 +77,44 @@ RD_TEST(VK_Query_Pool, VulkanGraphicsTest)
     VkQueryPool pool;
     vkCreateQueryPool(device, &poolInfo, NULL, &pool);
 
+    // populate the query we'll read the first frame
+    {
+      VkCommandBuffer cmd = GetCommandBuffer();
+
+      vkBeginCommandBuffer(cmd, vkh::CommandBufferBeginInfo());
+
+      vkCmdResetQueryPool(cmd, pool, 100, 4);
+
+      vkCmdBeginQuery(cmd, pool, 100, 0);
+      vkCmdEndQuery(cmd, pool, 100);
+
+      vkEndCommandBuffer(cmd);
+
+      Submit(99, 99, {cmd});
+
+      vkQueueWaitIdle(queue);
+    }
+
     while(Running())
     {
       VkCommandBuffer cmd = GetCommandBuffer();
 
       vkBeginCommandBuffer(cmd, vkh::CommandBufferBeginInfo());
 
+      // use a query from a previous frame then immediately reset it
+      uint32_t prevQueryIdx = (curFrame % 4);
+
+      vkCmdCopyQueryPoolResults(cmd, pool, 100 + prevQueryIdx, 1, vb.buffer,
+                                1024 + sizeof(uint64_t) * prevQueryIdx, sizeof(uint64_t),
+                                VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT);
+      vkCmdResetQueryPool(cmd, pool, 100 + prevQueryIdx, 1);
+
+      uint32_t curQueryIdx = ((curFrame + 1) % 4);
+
+      // fill the next query
+      vkCmdBeginQuery(cmd, pool, 100 + curQueryIdx, 0);
+      vkCmdEndQuery(cmd, pool, 100 + curQueryIdx);
+
       vkCmdResetQueryPool(cmd, pool, 123456, 1);
 
       VkImage swapimg =