Vulkan AS rebuild-on-replay: Serialise and Replay

renderdoc/driver/vulkan/vk_acceleration_structure.h

@@ -39,7 +39,6 @@ struct VkAccelerationStructureInfo
       VkDeviceSize vertexStride;
       uint32_t maxVertex;
       VkIndexType indexType;
-      bool hasTransformData;
     };
 
     struct Aabbs
@@ -50,28 +49,36 @@ struct VkAccelerationStructureInfo
     VkGeometryTypeKHR geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
     VkGeometryFlagsKHR flags;
 
-    VkDeviceMemory readbackMem;
-    VkDeviceSize memSize;
-
     Triangles tris;
     Aabbs aabbs;
 
     VkAccelerationStructureBuildRangeInfoKHR buildRangeInfo;
+    VkDeviceSize memOffset;
   };
 
   ~VkAccelerationStructureInfo();
 
   void AddRef() { Atomic::Inc32(&refCount); }
   void Release();
 
-  VkDevice device = VK_NULL_HANDLE;
+  uint64_t GetSerialisedSize() const;
+
+  void convertGeometryData(rdcarray<VkAccelerationStructureGeometryKHR> &geometry) const;
+  rdcarray<VkAccelerationStructureBuildRangeInfoKHR> getBuildRanges() const;
 
   VkAccelerationStructureTypeKHR type =
       VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_GENERIC_KHR;
   VkBuildAccelerationStructureFlagsKHR flags = 0;
 
   rdcarray<GeometryData> geometryData;
 
+  GPUBuffer readbackMem;
+  VkDeviceSize memSize = 0;
+
+  MemoryAllocation uploadAlloc;
+  VkBuffer uploadBuf = VK_NULL_HANDLE;
+  VkAccelerationStructureKHR replayAS = VK_NULL_HANDLE;
+
   bool accelerationStructureBuilt = false;
 
 private:
@@ -81,26 +88,6 @@ struct VkAccelerationStructureInfo
 class VulkanAccelerationStructureManager
 {
 public:
-  struct ASMemory
-  {
-    MemoryAllocation alloc;
-    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
@@ -115,28 +102,45 @@ class VulkanAccelerationStructureManager
   void CopyAccelerationStructure(VkCommandBuffer commandBuffer,
                                  const VkCopyAccelerationStructureInfoKHR &pInfo);
 
-  // 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,
-               ASMemory &result);
+  uint64_t GetSize_InitialState(ResourceId id, const VkInitialContents &initial);
 
   template <typename SerialiserType>
   bool Serialise(SerialiserType &ser, ResourceId id, const VkInitialContents *initial,
                  CaptureState state);
 
-  // Called when the initial state is applied.  The AS data is deserialised from the upload buffer
-  // into the acceleration structure
-  void Apply(ResourceId id, const VkInitialContents &initial);
+  // Called when the initial state is applied.
+  void Apply(ResourceId id, VkInitialContents &initial);
 
 private:
-  Allocation CreateReadBackMemory(VkDevice device, VkDeviceSize size, VkDeviceSize alignment = 0);
+  struct Allocation
+  {
+    MemoryAllocation memAlloc;
+    VkBuffer buf = VK_NULL_HANDLE;
+  };
+
+  struct RecordAndOffset
+  {
+    VkResourceRecord *record = NULL;
+    VkDeviceAddress address = 0x0;
+    VkDeviceSize offset = 0;
+  };
+
+  GPUBuffer CreateTempReadBackBuffer(VkDevice device, VkDeviceSize size);
+  Allocation CreateTempReplayBuffer(MemoryType memType, VkDeviceSize size, VkDeviceSize alignment,
+                                    VkBufferUsageFlags extraUsageFlags = 0);
+
+  bool FixUpReplayBDAs(VkAccelerationStructureInfo *asInfo, VkBuffer buf,
+                       rdcarray<VkAccelerationStructureGeometryKHR> &geoms);
+
+  void UpdateScratch(VkDeviceSize requiredSize);
 
   RecordAndOffset GetDeviceAddressData(VkDeviceAddress address) const;
 
   template <typename T>
   void DeletePreviousInfo(VkCommandBuffer commandBuffer, T *info);
 
-  VkDeviceSize SerialisedASSize(VkAccelerationStructureKHR unwrappedAs);
-
   WrappedVulkan *m_pDriver;
+
+  Allocation scratch;
+  VkDeviceOrHostAddressKHR scratchAddressUnion;
 };

renderdoc/driver/vulkan/vk_common.h

@@ -454,6 +454,8 @@ enum class MemoryScope : uint8_t
   // allocated the same way
   ImmutableReplayDebug = InitialContents,
   IndirectReadback,
+  // Same as initial contents but freed after first Serialise/Apply cycle
+  InitialContentsFirstApplyOnly,
   Count,
 };
 

renderdoc/driver/vulkan/vk_core.cpp

@@ -2305,6 +2305,8 @@ void WrappedVulkan::StartFrameCapture(DeviceOwnedWindow devWnd)
   GetResourceManager()->ClearReferencedResources();
   GetResourceManager()->ClearReferencedMemory();
 
+  CheckPendingCommandBufferCallbacks();
+
   // need to do all this atomically so that no other commands
   // will check to see if they need to markdirty or markpendingdirty
   // and go into the frame record.
@@ -2346,7 +2348,6 @@ void WrappedVulkan::StartFrameCapture(DeviceOwnedWindow devWnd)
     }
 
     m_PreparedNotSerialisedInitStates.clear();
-    CheckPendingCommandBufferCallbacks();
     GetResourceManager()->PrepareInitialContents();
 
     {
@@ -3264,7 +3265,7 @@ RDResult WrappedVulkan::ReadLogInitialisation(RDCFile *rdc, bool storeStructured
   GetReplay()->WriteFrameRecord().frameInfo.initDataSize =
       chunkInfos[(VulkanChunk)SystemChunk::InitialContents].totalsize;
 
-  RDCDEBUG("Allocating %llu persistant bytes of memory for the log.",
+  RDCDEBUG("Allocating %llu persistent bytes of memory for the log.",
            GetReplay()->WriteFrameRecord().frameInfo.persistentSize);
 
   // ensure the capture at least created a device and fetched a queue.
@@ -3780,6 +3781,8 @@ void WrappedVulkan::ApplyInitialContents()
     SubmitCmds();
     FlushQ();
   }
+
+  FreeAllMemory(MemoryScope::InitialContentsFirstApplyOnly);
 }
 
 bool WrappedVulkan::ContextProcessChunk(ReadSerialiser &ser, VulkanChunk chunk)

renderdoc/driver/vulkan/vk_initstate.cpp

@@ -92,6 +92,12 @@ bool WrappedVulkan::Prepare_InitialState(WrappedVkRes *res)
     if(imageInfo.levelCount > 1)
       estimatedSize *= 2;
   }
+  else if(type == eResAccelerationStructureKHR)
+  {
+    VkResourceRecord *record = GetResourceManager()->GetResourceRecord(id);
+    if(record && record->accelerationStructureInfo)
+      estimatedSize += record->accelerationStructureInfo->memSize;
+  }
 
   uint32_t softMemoryLimit = RenderDoc::Inst().GetCaptureOptions().softMemoryLimit;
   if(softMemoryLimit > 0 && !m_PreparedNotSerialisedInitStates.empty() &&
@@ -118,8 +124,6 @@ bool WrappedVulkan::Prepare_InitialState(WrappedVkRes *res)
     {
       VkInitialContents initData = GetResourceManager()->GetInitialContents(flushId);
 
-      GetResourceManager()->SetInitialContents(flushId, VkInitialContents());
-
       uint64_t start = ser.GetWriter()->GetOffset();
       {
         uint64_t size = GetSize_InitialState(flushId, initData);
@@ -128,6 +132,7 @@ bool WrappedVulkan::Prepare_InitialState(WrappedVkRes *res)
 
         // record is not needed on vulkan
         Serialise_InitialState(ser, flushId, NULL, &initData);
+        GetResourceManager()->SetInitialContents(flushId, VkInitialContents());
       }
       uint64_t end = ser.GetWriter()->GetOffset();
 
@@ -577,25 +582,18 @@ bool WrappedVulkan::Prepare_InitialState(WrappedVkRes *res)
   else if(type == eResAccelerationStructureKHR)
   {
     VkResourceRecord *record = GetResourceManager()->GetResourceRecord(id);
-
     if(!record->accelerationStructureInfo->accelerationStructureBuilt)
     {
       RDCDEBUG("Skipping AS %s as it has not been built", ToStr(id).c_str());
       return true;
     }
 
-    VulkanAccelerationStructureManager::ASMemory result;
-    VkAccelerationStructureKHR as = ToUnwrappedHandle<VkAccelerationStructureKHR>(res);
-    if(!GetAccelerationStructureManager()->Prepare(as, m_QueueFamilyIndices, result))
-    {
-      SET_ERROR_RESULT(m_LastCaptureError, ResultCode::OutOfMemory,
-                       "Couldn't allocate readback memory");
-      m_CaptureFailure = true;
-      return false;
-    }
-
-    VkInitialContents ic = VkInitialContents(type, result.alloc);
-    ic.isTLAS = result.isTLAS;
+    // The input buffers and metadata have all been created by this point, so we just need to
+    // assemble a VkInitialContents
+    VkInitialContents ic;
+    ic.type = type;
+    ic.accelerationStructureInfo = record->accelerationStructureInfo;
+    ic.accelerationStructureInfo->AddRef();
 
     GetResourceManager()->SetInitialContents(id, ic);
     m_PreparedNotSerialisedInitStates.push_back(id);
@@ -639,12 +637,15 @@ uint64_t WrappedVulkan::GetSize_InitialState(ResourceId id, const VkInitialConte
     // buffers only have initial states when they're sparse
     return ret;
   }
-  else if(initial.type == eResImage || initial.type == eResDeviceMemory ||
-          initial.type == eResAccelerationStructureKHR)
+  else if(initial.type == eResImage || initial.type == eResDeviceMemory)
   {
     // the size primarily comes from the buffer, the size of which we conveniently have stored.
     return ret + uint64_t(128 + initial.mem.size + WriteSerialiser::GetChunkAlignment());
   }
+  else if(initial.type == eResAccelerationStructureKHR)
+  {
+    return GetAccelerationStructureManager()->GetSize_InitialState(id, initial);
+  }
 
   RDCERR("Unhandled resource type %s", ToStr(initial.type).c_str());
   return 128;

renderdoc/driver/vulkan/vk_manager.cpp

@@ -1042,9 +1042,17 @@ rdcarray<ResourceId> VulkanResourceManager::InitialContentResources()
     const InitialContentData &bData = m_InitialContents[b].data;
 
     // Always sort BLASs before TLASs, as a TLAS holds device addresses for it's BLASs
-    // and we make sure those addresses are valid
-    if(!aData.isTLAS && bData.isTLAS)
-      return true;
+    // and we make sure those addresses are valid.  There's no good handling for the generic types,
+    // so we just assume it is a TLAS
+    if(aData.accelerationStructureInfo && bData.accelerationStructureInfo)
+    {
+      const VkAccelerationStructureTypeKHR aType = aData.accelerationStructureInfo->type;
+      const VkAccelerationStructureTypeKHR bType = bData.accelerationStructureInfo->type;
+      if(aType == VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR &&
+         (bType == VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR ||
+          bType == VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_GENERIC_KHR))
+        return true;
+    }
 
     return aData.type < bData.type;
   });

renderdoc/driver/vulkan/vk_manager.h

@@ -25,6 +25,7 @@
 #pragma once
 
 #include "core/resource_manager.h"
+#include "vk_acceleration_structure.h"
 #include "vk_resources.h"
 
 class WrappedVulkan;
@@ -113,7 +114,9 @@ struct VkInitialContents
     SAFE_DELETE(sparseTables);
     SAFE_DELETE(sparseBind);
 
-    // MemoryAllocation and serialised ASes are not free'd here
+    SAFE_RELEASE(accelerationStructureInfo);
+
+    // MemoryAllocation ise not free'd here
   }
 
   // for descriptor heaps, when capturing we save the slots, when replaying we store direct writes
@@ -139,7 +142,7 @@ struct VkInitialContents
   rdcarray<AspectSparseTable> *sparseTables;
   SparseBinding *sparseBind;
 
-  bool isTLAS;    // If the contents are an AS, this determines if it is a TLAS or BLAS
+  VkAccelerationStructureInfo *accelerationStructureInfo;
 };
 
 struct VulkanResourceManagerConfiguration

renderdoc/driver/vulkan/vk_resources.cpp

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

renderdoc/driver/vulkan/vk_stringise.cpp

@@ -298,6 +298,7 @@ rdcstr DoStringise(const MemoryScope &el)
   {
     STRINGISE_ENUM_CLASS(InitialContents);
     STRINGISE_ENUM_CLASS(IndirectReadback);
+    STRINGISE_ENUM_CLASS(InitialContentsFirstApplyOnly);
   }
   END_ENUM_STRINGISE()
 }

renderdoc/driver/vulkan/wrappers/vk_device_funcs.cpp

@@ -1014,6 +1014,7 @@ void WrappedVulkan::Shutdown()
   }
 
   FreeAllMemory(MemoryScope::InitialContents);
+  FreeAllMemory(MemoryScope::InitialContentsFirstApplyOnly);
 
   if(m_MemoryFreeThread)
   {