Track ASs by address when descriptor buffers are enabled

* They are written into descriptors via address not object with descriptor
  buffer

renderdoc/driver/vulkan/vk_acceleration_structure.h

@@ -66,6 +66,8 @@ struct VkAccelerationStructureInfo
   void convertGeometryData(rdcarray<VkAccelerationStructureGeometryKHR> &geometry) const;
   rdcarray<VkAccelerationStructureBuildRangeInfoKHR> getBuildRanges() const;
 
+  VkDeviceAddress address;
+
   VkAccelerationStructureTypeKHR type =
       VkAccelerationStructureTypeKHR::VK_ACCELERATION_STRUCTURE_TYPE_GENERIC_KHR;
   VkBuildAccelerationStructureFlagsKHR flags = 0;

renderdoc/driver/vulkan/vk_core.h

@@ -447,6 +447,9 @@ private:
     VkQueueFamilyProperties queueProps[16] = {};
   };
 
+  Threading::CriticalSection m_ASLookupByAddrLock;
+  rdcflatmap<VkDeviceAddress, ResourceId> m_ASLookupByAddr;
+
   bool m_SeparateDepthStencil = false;
   bool m_NULLDescriptorsAllowed = false;
   bool m_ExtendedDynState = false;

renderdoc/driver/vulkan/wrappers/vk_misc_funcs.cpp

@@ -156,18 +156,6 @@ static void MakeSubpassLoadRP(RPCreateInfo &info, const RPCreateInfo *origInfo,
   }
 }
 
-template <typename type>
-bool RemoveForcedRef()
-{
-  return false;
-}
-
-template <>
-bool RemoveForcedRef<VkAccelerationStructureKHR>()
-{
-  return true;
-}
-
 // note, for threading reasons we ensure to release the wrappers before
 // releasing the underlying object. Otherwise after releasing the vulkan object
 // that same handle could be returned by create on another thread, and we
@@ -178,11 +166,6 @@ bool RemoveForcedRef<VkAccelerationStructureKHR>()
     if(obj == VK_NULL_HANDLE)                                                            \
       return;                                                                            \
     type unwrappedObj = Unwrap(obj);                                                     \
-    if(RemoveForcedRef<type>())                                                          \
-    {                                                                                    \
-      SCOPED_LOCK(m_ForcedReferencesLock);                                               \
-      m_ForcedReferences.removeOne(GetRecord(obj));                                      \
-    }                                                                                    \
     if(IsReplayMode(m_State))                                                            \
       m_CreationInfo.erase(GetResID(obj));                                               \
     GetResourceManager()->ReleaseWrappedResource(obj, true);                             \
@@ -205,11 +188,30 @@ DESTROY_IMPL(VkCommandPool, DestroyCommandPool)
 DESTROY_IMPL(VkQueryPool, DestroyQueryPool)
 DESTROY_IMPL(VkDescriptorUpdateTemplate, DestroyDescriptorUpdateTemplate)
 DESTROY_IMPL(VkSamplerYcbcrConversion, DestroySamplerYcbcrConversion)
-DESTROY_IMPL(VkAccelerationStructureKHR, DestroyAccelerationStructureKHR)
 DESTROY_IMPL(VkShaderEXT, DestroyShaderEXT)
 
 #undef DESTROY_IMPL
 
+void WrappedVulkan::vkDestroyAccelerationStructureKHR(VkDevice device, VkAccelerationStructureKHR obj,
+                                                      const VkAllocationCallbacks *)
+{
+  if(obj == VK_NULL_HANDLE)
+    return;
+  VkAccelerationStructureKHR unwrappedObj = Unwrap(obj);
+  {
+    SCOPED_LOCK(m_ASLookupByAddrLock);
+    m_ASLookupByAddr.erase(GetRecord(obj)->accelerationStructureInfo->address);
+  }
+  {
+    SCOPED_LOCK(m_ForcedReferencesLock);
+    m_ForcedReferences.removeOne(GetRecord(obj));
+  }
+  if(IsReplayMode(m_State))
+    m_CreationInfo.erase(GetResID(obj));
+  GetResourceManager()->ReleaseWrappedResource(obj, true);
+  ObjDisp(device)->DestroyAccelerationStructureKHR(Unwrap(device), unwrappedObj, NULL);
+}
+
 void WrappedVulkan::vkDestroyFramebuffer(VkDevice device, VkFramebuffer obj,
                                          const VkAllocationCallbacks *)
 {

renderdoc/driver/vulkan/wrappers/vk_resource_funcs.cpp

@@ -3567,6 +3567,16 @@ bool WrappedVulkan::Serialise_vkCreateAccelerationStructureKHR(
 
         m_CreationInfo.m_AccelerationStructure[live].Init(GetResourceManager(), m_CreationInfo,
                                                           &CreateInfo);
+
+        const VkAccelerationStructureDeviceAddressInfoKHR getInfo = {
+            VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
+            NULL,
+            Unwrap(acc),
+        };
+        const VkDeviceAddress addr =
+            ObjDisp(device)->GetAccelerationStructureDeviceAddressKHR(Unwrap(device), &getInfo);
+
+        m_ASLookupByAddr[addr] = live;
       }
     }
 
@@ -3638,6 +3648,7 @@ VkResult WrappedVulkan::vkCreateAccelerationStructureKHR(
       record->AddParent(bufferRecord);
 
       record->accelerationStructureInfo = new VkAccelerationStructureInfo();
+      record->accelerationStructureInfo->address = addr;
 
       // store the base resource
       record->baseResource = bufferRecord->GetResourceID();
@@ -3659,6 +3670,14 @@ VkResult WrappedVulkan::vkCreateAccelerationStructureKHR(
         // in case we're currently capturing, immediately consider the AS as referenced
         GetResourceManager()->MarkResourceFrameReferenced(record->GetResourceID(), eFrameRef_Read);
       }
+
+      if(m_DescriptorBuffers)
+      {
+        {
+          SCOPED_LOCK(m_ASLookupByAddrLock);
+          m_ASLookupByAddr[addr] = id;
+        }
+      }
     }
     else
     {