Vulkan AS rebuild-on-replay: On queue submission completion events

If an AS build command is called on an already built AS (i.e. an update build) then we need to safely destroy the previously copied input buffers and their memory.  However at the point this happens, we do not own the command buffer so we don't know if the original data has actually been copied yet.

Using the mechanism in this patch we associate function objects with the command buffer that will be executed when the queue the command buffer is submitted to, has finished executing on the GPU.

This works by carying the function objects in command buffer and once it is submitted, we poll on a VkEvent we inserted.  Once signalled, we call the function objects in the order they were inserted.

renderdoc/driver/vulkan/vk_core.cpp

@@ -705,6 +705,75 @@ void WrappedVulkan::SetDebugMessageSink(WrappedVulkan::ScopedDebugMessageSink *s
   Threading::SetTLSValue(debugMessageSinkTLSSlot, (void *)sink);
 }
 
+void WrappedVulkan::MarkPendingCommandBufferAsDeleted(VkCommandBuffer commandBuffer)
+{
+  SCOPED_LOCK(m_PendingCmdBufferCallbacksLock);
+
+  for(PendingCommandBufferCallbacks &pending : m_PendingCmdBufferCallbacks)
+  {
+    if(pending.commandBuffer == commandBuffer)
+    {
+      pending.commandBuffer = VK_NULL_HANDLE;
+      return;
+    }
+  }
+}
+
+void WrappedVulkan::AddPendingCommandBufferCallbacks(VkCommandBuffer commandBuffer)
+{
+  SCOPED_LOCK(m_PendingCmdBufferCallbacksLock);
+
+  VkResourceRecord *cmdRecord = GetRecord(commandBuffer);
+  rdcarray<std::function<void()>> &callbacks = cmdRecord->cmdInfo->pendingSubmissionCompleteCallbacks;
+
+  if(callbacks.empty())
+    return;
+
+  const VkEventCreateInfo info = {VK_STRUCTURE_TYPE_EVENT_CREATE_INFO};
+  VkEvent event;
+  const VkResult vkr = ObjDisp(m_Device)->CreateEvent(Unwrap(m_Device), &info, NULL, &event);
+  CheckVkResult(vkr);
+
+  ObjDisp(commandBuffer)->CmdSetEvent(Unwrap(commandBuffer), event, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT);
+
+  m_PendingCmdBufferCallbacks.push_back({event, commandBuffer, std::move(callbacks)});
+}
+
+void WrappedVulkan::CheckPendingCommandBufferCallbacks()
+{
+  SCOPED_LOCK(m_PendingCmdBufferCallbacksLock);
+
+  // Delete any events whose command buffers have been deleted and their callbacks called
+  m_PendingCmdBufferCallbacks.removeIf([d = m_Device](const PendingCommandBufferCallbacks &pending) {
+    const bool destroyEvent = pending.commandBuffer == VK_NULL_HANDLE && pending.callbacks.empty();
+    if(destroyEvent)
+      ObjDisp(d)->DestroyEvent(Unwrap(d), pending.event, NULL);
+
+    return destroyEvent;
+  });
+
+  for(PendingCommandBufferCallbacks &pending : m_PendingCmdBufferCallbacks)
+  {
+    // We've already executed the callbacks, we're now just waiting for the command buffer to be
+    // destroyed/reused so we can delete the event
+    if(pending.callbacks.empty())
+      continue;
+
+    const VkResult vkr = ObjDisp(m_Device)->GetEventStatus(Unwrap(m_Device), pending.event);
+    if(vkr == VK_EVENT_SET)
+    {
+      for(std::function<void()> &f : pending.callbacks)
+        f();
+
+      pending.callbacks.clear();
+    }
+    else if(vkr != VK_EVENT_RESET)
+    {
+      CheckVkResult(vkr);
+    }
+  }
+}
+
 byte *WrappedVulkan::GetRingTempMemory(size_t s)
 {
   TempMem *mem = (TempMem *)Threading::GetTLSValue(tempMemoryTLSSlot);
@@ -2268,6 +2337,7 @@ void WrappedVulkan::StartFrameCapture(DeviceOwnedWindow devWnd)
     }
 
     m_PreparedNotSerialisedInitStates.clear();
+    CheckPendingCommandBufferCallbacks();
     GetResourceManager()->PrepareInitialContents();
 
     {

renderdoc/driver/vulkan/vk_core.h

@@ -979,6 +979,19 @@ private:
 
   bytebuf m_MaskedMapData;
 
+  struct PendingCommandBufferCallbacks
+  {
+    VkEvent event;
+    VkCommandBuffer commandBuffer;
+    rdcarray<std::function<void()>> callbacks;
+  };
+
+  Threading::CriticalSection m_PendingCmdBufferCallbacksLock;
+  rdcarray<PendingCommandBufferCallbacks> m_PendingCmdBufferCallbacks;
+  void MarkPendingCommandBufferAsDeleted(VkCommandBuffer commandBuffer);
+  void AddPendingCommandBufferCallbacks(VkCommandBuffer commandBuffer);
+  void CheckPendingCommandBufferCallbacks();
+
   GPUAddressRangeTracker m_AddressTracker;
   GPUAddressRange CreateAddressRange(VkDevice device, VkBuffer buffer);
 

renderdoc/driver/vulkan/vk_resources.h

@@ -1144,6 +1144,9 @@ struct CmdBufferRecordingInfo
   // A list of acceleration structures that this command buffer will build or copy
   rdcarray<VkResourceRecord *> accelerationStructures;
 
+  // A list of callbacks to be executed once the command buffer execution has been completed
+  rdcarray<std::function<void()>> pendingSubmissionCompleteCallbacks;
+
   // AdvanceFrame/Present should be called after this buffer is submitted
   bool present;
   // BeginFrameCapture should be called *before* this buffer is submitted.
@@ -2247,6 +2250,8 @@ public:
     cmdInfo->imageStates.swap(bakedCommands->cmdInfo->imageStates);
     cmdInfo->memFrameRefs.swap(bakedCommands->cmdInfo->memFrameRefs);
     cmdInfo->accelerationStructures.swap(bakedCommands->cmdInfo->accelerationStructures);
+    cmdInfo->pendingSubmissionCompleteCallbacks.swap(
+        bakedCommands->cmdInfo->pendingSubmissionCompleteCallbacks);
   }
 
   // we have a lot of 'cold' data in the resource record, as it can be accessed

renderdoc/driver/vulkan/wrappers/vk_cmd_funcs.cpp

@@ -1378,11 +1378,14 @@ VkResult WrappedVulkan::vkBeginCommandBuffer(VkCommandBuffer commandBuffer,
 
   if(record)
   {
-    // If a command bfufer was already recorded (ie we have some baked commands),
+    // If a command buffer was already recorded (ie we have some baked commands),
     // then begin is spec'd to implicitly reset. That means we need to tidy up
     // any existing baked commands before creating a new set.
     if(record->bakedCommands)
+    {
+      MarkPendingCommandBufferAsDeleted(commandBuffer);
       record->bakedCommands->Delete(GetResourceManager());
+    }
 
     record->bakedCommands = GetResourceManager()->AddResourceRecord(ResourceIDGen::GetNewUniqueID());
     record->bakedCommands->resType = eResCommandBuffer;
@@ -1670,6 +1673,9 @@ VkResult WrappedVulkan::vkEndCommandBuffer(VkCommandBuffer commandBuffer)
   VkResourceRecord *record = GetRecord(commandBuffer);
   RDCASSERT(record);
 
+  if(IsCaptureMode(m_State))
+    AddPendingCommandBufferCallbacks(commandBuffer);
+
   VkResult ret;
   SERIALISE_TIME_CALL(ret = ObjDisp(commandBuffer)->EndCommandBuffer(Unwrap(commandBuffer)));
 

renderdoc/driver/vulkan/wrappers/vk_misc_funcs.cpp

@@ -345,6 +345,8 @@ void WrappedVulkan::vkFreeCommandBuffers(VkDevice device, VkCommandPool commandP
     if(pCommandBuffers[c] == VK_NULL_HANDLE)
       continue;
 
+    MarkPendingCommandBufferAsDeleted(pCommandBuffers[c]);
+
     WrappedVkDispRes *wrapped = (WrappedVkDispRes *)GetWrapped(pCommandBuffers[c]);
 
 #if ENABLED(VERBOSE_PARTIAL_REPLAY)

renderdoc/driver/vulkan/wrappers/vk_queue_funcs.cpp

@@ -1337,6 +1337,8 @@ void WrappedVulkan::CaptureQueueSubmit(VkQueue queue,
 
   for(VkResourceRecord *asRecord : accelerationStructures)
     asRecord->accelerationStructureInfo->accelerationStructureBuilt = true;
+
+  CheckPendingCommandBufferCallbacks();
 }
 
 template <typename SerialiserType>