Vulkan AS rebuild-on-replay: Copy AS input buffers on build

When vkCmdBuildAccelerationStructuresKHR is called, the input buffers used for the build are copied into host-accessible memory, so when initial state serialisation is triggers it can be copied into the capture file.

To unify the behaviour of the different geometry types, there is only one readback buffer per AS, and in the case of triangle data it can contain the vertex, index, and transform data concatenated into a single contiguous block.

The buffer data along with build arguments are stored in a ref-counted VkAccelerationStructureInfo object held by the AS's VkResourceRecord.  Once the info object is destroyed, then the copied buffer mem is also freed.

Serialisation of the initial state will follow in later patches.

Caveats:
* vkCmdBuildAccelerationStructuresIndirectKHR is not handled
* VkAccelerationStructureGeometryInstancesDataKHR::arrayOfPointers mode is not handled

renderdoc/driver/vulkan/vk_acceleration_structure.h

@@ -24,14 +24,58 @@
 
 #pragma once
 
-#include "vk_manager.h"
+#include "vk_resources.h"
 
 class WrappedVulkan;
+struct VkInitialContents;
 
-// Just holds the built flag, will eventually hold all the AS build data
 struct VkAccelerationStructureInfo
 {
+  struct GeometryData
+  {
+    struct Triangles
+    {
+      VkFormat vertexFormat;
+      VkDeviceSize vertexStride;
+      uint32_t maxVertex;
+      VkIndexType indexType;
+      bool hasTransformData;
+    };
+
+    struct Aabbs
+    {
+      VkDeviceSize stride;
+    };
+
+    VkGeometryTypeKHR geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
+    VkGeometryFlagsKHR flags;
+
+    VkDeviceMemory readbackMem;
+    VkDeviceSize memSize;
+
+    Triangles tris;
+    Aabbs aabbs;
+
+    VkAccelerationStructureBuildRangeInfoKHR buildRangeInfo;
+  };
+
+  ~VkAccelerationStructureInfo();
+
+  void AddRef() { Atomic::Inc32(&refCount); }
+  void Release();
+
+  VkDevice device = VK_NULL_HANDLE;
+
+  VkAccelerationStructureTypeKHR type =
+      VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_GENERIC_KHR;
+  VkBuildAccelerationStructureFlagsKHR flags = 0;
+
+  rdcarray<GeometryData> geometryData;
+
   bool accelerationStructureBuilt = false;
+
+private:
+  int32_t refCount = 1;
 };
 
 class VulkanAccelerationStructureManager
@@ -43,8 +87,35 @@ class VulkanAccelerationStructureManager
     bool isTLAS;
   };
 
+  struct Allocation
+  {
+    VkDeviceMemory mem = VK_NULL_HANDLE;
+    VkDeviceSize size = 0;
+    VkBuffer buf = VK_NULL_HANDLE;
+  };
+
+  struct RecordAndOffset
+  {
+    VkResourceRecord *record = NULL;
+    VkDeviceAddress address = 0x0;
+    VkDeviceSize offset = 0;
+  };
+
   explicit VulkanAccelerationStructureManager(WrappedVulkan *driver);
 
+  // Allocates readback mem and injects commands into the command buffer so that the input buffers
+  // are copied.
+  void CopyInputBuffers(VkCommandBuffer commandBuffer,
+                        const VkAccelerationStructureBuildGeometryInfoKHR &info,
+                        const VkAccelerationStructureBuildRangeInfoKHR *buildRange,
+                        CaptureState state);
+
+  // Copies the metadata from src to dst, the input buffers are identical so don't need to be
+  // duplicated.  Compaction is ignored but the copy is still performed so the dst handle is valid
+  // on replay
+  void CopyAccelerationStructure(VkCommandBuffer commandBuffer,
+                                 const VkCopyAccelerationStructureInfoKHR &pInfo, CaptureState state);
+
   // Called when the initial state is prepared.  Any TLAS and BLAS data is copied into temporary
   // buffers and the handles for that memory and the buffers is stored in the init state
   bool Prepare(VkAccelerationStructureKHR unwrappedAs, const rdcarray<uint32_t> &queueFamilyIndices,
@@ -59,6 +130,13 @@ class VulkanAccelerationStructureManager
   void Apply(ResourceId id, const VkInitialContents &initial);
 
 private:
+  Allocation CreateReadBackMemory(VkDevice device, VkDeviceSize size, VkDeviceSize alignment = 0);
+
+  RecordAndOffset GetDeviceAddressData(VkDeviceAddress address) const;
+
+  template <typename T>
+  void DeletePreviousInfo(VkCommandBuffer commandBuffer, T *info, CaptureState state);
+
   VkDeviceSize SerialisedASSize(VkAccelerationStructureKHR unwrappedAs);
 
   WrappedVulkan *m_pDriver;

renderdoc/driver/vulkan/vk_core.h

@@ -1225,6 +1225,7 @@ class WrappedVulkan : public IFrameCapturer
   void RemapQueueFamilyIndices(uint32_t &srcQueueFamily, uint32_t &dstQueueFamily);
   uint32_t GetQueueFamilyIndex() const { return m_QueueFamilyIdx; }
   bool ReleaseResource(WrappedVkRes *res);
+  const rdcarray<uint32_t> &GetQueueFamilyIndices() const { return m_QueueFamilyIndices; }
 
   void AddDebugMessage(MessageCategory c, MessageSeverity sv, MessageSource src, rdcstr d);
 
@@ -1275,6 +1276,7 @@ class WrappedVulkan : public IFrameCapturer
 
   void TrackBufferAddress(VkDevice device, VkBuffer buffer);
   void UntrackBufferAddress(VkDevice device, VkBuffer buffer);
+  void GetResIDFromAddr(GPUAddressRange::Address addr, ResourceId &id, uint64_t &offs);
 
   EventFlags GetEventFlags(uint32_t eid) { return m_EventFlags[eid]; }
   rdcarray<EventUsage> GetUsage(ResourceId id) { return m_ResourceUses[id]; }

renderdoc/driver/vulkan/vk_memory.cpp

@@ -77,6 +77,11 @@ void WrappedVulkan::UntrackBufferAddress(VkDevice device, VkBuffer buffer)
   m_AddressTracker.RemoveFrom(rng);
 }
 
+void WrappedVulkan::GetResIDFromAddr(GPUAddressRange::Address addr, ResourceId &id, uint64_t &offs)
+{
+  m_AddressTracker.GetResIDFromAddr(addr, id, offs);
+}
+
 void WrappedVulkan::ChooseMemoryIndices()
 {
   // we need to do this little dance because Get*MemoryIndex checks to see if the existing

renderdoc/driver/vulkan/vk_resources.cpp

@@ -3996,8 +3996,8 @@ VkResourceRecord::~VkResourceRecord()
   if(resType == eResCommandPool)
     SAFE_DELETE(cmdPoolInfo);
 
-  if(resType == eResAccelerationStructureKHR)
-    SAFE_DELETE(accelerationStructureInfo);
+  if(resType == eResAccelerationStructureKHR && accelerationStructureInfo)
+    accelerationStructureInfo->Release();
 }
 
 void VkResourceRecord::MarkImageFrameReferenced(VkResourceRecord *img, const ImageRange &range,

renderdoc/driver/vulkan/wrappers/vk_cmd_funcs.cpp

@@ -7873,6 +7873,9 @@ void WrappedVulkan::vkCmdBuildAccelerationStructuresKHR(
 
       // Add to the command buffer metadata, so we can know when it has been submitted
       record->cmdInfo->accelerationStructures.push_back(GetRecord(geomInfo.dstAccelerationStructure));
+
+      GetAccelerationStructureManager()->CopyInputBuffers(commandBuffer, geomInfo,
+                                                          ppBuildRangeInfos[i], m_State);
     }
   }
 }

renderdoc/driver/vulkan/wrappers/vk_get_funcs.cpp

@@ -860,8 +860,9 @@ void WrappedVulkan::vkGetPhysicalDeviceFeatures2(VkPhysicalDevice physicalDevice
 
   if(accStruct && accStruct->accelerationStructureHostCommands)
   {
-    RDCWARN("Disabling support for acceleration structure host commands");
+    RDCWARN("Disabling support for acceleration structure host commands and indirect builds");
     accStruct->accelerationStructureHostCommands = VK_FALSE;
+    accStruct->accelerationStructureIndirectBuild = VK_FALSE;
   }
 }
 

renderdoc/driver/vulkan/wrappers/vk_resource_funcs.cpp

@@ -3389,7 +3389,6 @@ VkResult WrappedVulkan::vkCreateAccelerationStructureKHR(
       record->baseResource = bufferRecord->GetResourceID();
       record->baseResourceMem = bufferRecord->baseResource;
       record->dedicated = bufferRecord->dedicated;
-      record->resInfo = bufferRecord->resInfo;
       record->storable = bufferRecord->storable;
       record->memOffset = bufferRecord->memOffset + pCreateInfo->offset;
       record->memSize = pCreateInfo->size;