Fix memory reference tracking of forced references

renderdoc/driver/vulkan/vk_core.cpp

@@ -1552,15 +1552,17 @@ void WrappedVulkan::StartFrameCapture(void *dev, void *wnd)
   GetResourceManager()->MarkResourceFrameReferenced(GetResID(m_Device), eFrameRef_Read);
   GetResourceManager()->MarkResourceFrameReferenced(GetResID(m_Queue), eFrameRef_Read);
 
-  std::map<ResourceId, FrameRefType> forced = GetForcedReferences();
+  rdcarray<VkResourceRecord *> forced = GetForcedReferences();
 
   // Note we force read-before-write because this resource is implicitly untracked so we have no
   // way of knowing how it's used
   for(auto it = forced.begin(); it != forced.end(); ++it)
   {
-    GetResourceManager()->MarkResourceFrameReferenced(it->first, eFrameRef_Read);
-    if(it->second != eFrameRef_Read)
-      GetResourceManager()->MarkResourceFrameReferenced(it->first, it->second);
+    // reference the buffer
+    GetResourceManager()->MarkResourceFrameReferenced((*it)->GetResourceID(), eFrameRef_Read);
+    // and its backing memory
+    GetResourceManager()->MarkMemoryFrameReferenced((*it)->baseResource, (*it)->memOffset,
+                                                    (*it)->memSize, eFrameRef_ReadBeforeWrite);
   }
 
   RDCLOG("Starting capture, frame %u", m_FrameCounter);

renderdoc/driver/vulkan/vk_core.h

@@ -719,12 +719,12 @@ private:
   std::vector<VkResourceRecord *> m_CoherentMaps;
   Threading::CriticalSection m_CoherentMapsLock;
 
-  std::map<ResourceId, FrameRefType> m_ForcedReferences;
+  rdcarray<VkResourceRecord *> m_ForcedReferences;
   Threading::CriticalSection m_ForcedReferencesLock;
 
-  std::map<ResourceId, FrameRefType> GetForcedReferences()
+  rdcarray<VkResourceRecord *> GetForcedReferences()
   {
-    std::map<ResourceId, FrameRefType> ret;
+    rdcarray<VkResourceRecord *> ret;
 
     {
       SCOPED_LOCK(m_ForcedReferencesLock);
@@ -734,30 +734,22 @@ private:
     return ret;
   }
 
-  bool IsForcedReference(ResourceId id)
+  bool IsForcedReference(VkResourceRecord *record)
   {
     bool ret = false;
 
     {
       SCOPED_LOCK(m_ForcedReferencesLock);
-      ret = (m_ForcedReferences.find(id) != m_ForcedReferences.end());
+      ret = (m_ForcedReferences.indexOf(record) != -1);
     }
 
     return ret;
   }
 
-  void AddForcedReference(ResourceId id, FrameRefType ref)
+  void AddForcedReference(VkResourceRecord *record)
   {
     SCOPED_LOCK(m_ForcedReferencesLock);
-    m_ForcedReferences[id] = ref;
-
-    // also add it immediately in case we're mid-way through a frame, and the forced references have
-    // already been processed for this frame.
-    // Note we force read-before-write because this resource is implicitly untracked so we have no
-    // way of knowing how it's used
-    GetResourceManager()->MarkResourceFrameReferenced(id, eFrameRef_Read);
-    if(ref != eFrameRef_Read)
-      GetResourceManager()->MarkResourceFrameReferenced(id, ref);
+    m_ForcedReferences.push_back(record);
   }
 
   // used both on capture and replay side to track image layouts. Only locked

renderdoc/driver/vulkan/wrappers/vk_misc_funcs.cpp

@@ -98,7 +98,7 @@ VkFramebufferCreateInfo WrappedVulkan::UnwrapInfo(const VkFramebufferCreateInfo
     if(obj == VK_NULL_HANDLE)                                                                      \
       return;                                                                                      \
     type unwrappedObj = Unwrap(obj);                                                               \
-    m_ForcedReferences.erase(GetResID(obj));                                                       \
+    m_ForcedReferences.removeOne(GetRecord(obj));                                                  \
     if(IsReplayMode(m_State))                                                                      \
       m_CreationInfo.erase(GetResID(obj));                                                         \
     GetResourceManager()->ReleaseWrappedResource(obj, true);                                       \

renderdoc/driver/vulkan/wrappers/vk_resource_funcs.cpp

@@ -511,7 +511,6 @@ void WrappedVulkan::vkFreeMemory(VkDevice device, VkDeviceMemory memory,
     }
   }
 
-  m_ForcedReferences.erase(GetResID(memory));
   m_CreationInfo.erase(GetResID(memory));
 
   GetResourceManager()->ReleaseWrappedResource(memory);
@@ -919,9 +918,14 @@ VkResult WrappedVulkan::vkBindBufferMemory(VkDevice device, VkBuffer buffer, VkD
     record->memOffset = memoryOffset;
 
     // if the buffer was force-referenced, do the same with the memory
-    if(IsForcedReference(GetResID(buffer)))
+    if(IsForcedReference(record))
     {
-      AddForcedReference(GetResID(memory), eFrameRef_ReadBeforeWrite);
+      // in case we're currently capturing, immediately consider the buffer and backing memory as
+      // read-before-write referenced
+      GetResourceManager()->MarkResourceFrameReferenced(record->GetResourceID(),
+                                                        eFrameRef_ReadBeforeWrite);
+      GetResourceManager()->MarkMemoryFrameReferenced(GetResID(memory), memoryOffset,
+                                                      record->memSize, eFrameRef_ReadBeforeWrite);
 
       // the memory is immediately dirty because we have no way of tracking writes to it
       GetResourceManager()->MarkDirtyResource(GetResID(memory));
@@ -1135,6 +1139,9 @@ VkResult WrappedVulkan::vkCreateBuffer(VkDevice device, const VkBufferCreateInfo
           VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_CREATE_INFO_EXT,
       };
 
+      VkResourceRecord *record = GetResourceManager()->AddResourceRecord(*pBuffer);
+      record->memSize = pCreateInfo->size;
+
       // if we're using VK_EXT_buffer_device_address, we fetch the device address that's been
       // allocated and insert it into the next chain and patch the flags so that it replays
       // naturally.
@@ -1159,7 +1166,7 @@ VkResult WrappedVulkan::vkCreateBuffer(VkDevice device, const VkBufferCreateInfo
 
         // this buffer must be forced to be in any captures, since we can't track when it's used by
         // address
-        AddForcedReference(GetResID(*pBuffer), eFrameRef_Read);
+        AddForcedReference(record);
       }
 
       {
@@ -1171,9 +1178,7 @@ VkResult WrappedVulkan::vkCreateBuffer(VkDevice device, const VkBufferCreateInfo
         chunk = scope.Get();
       }
 
-      VkResourceRecord *record = GetResourceManager()->AddResourceRecord(*pBuffer);
       record->AddChunk(chunk);
-      record->memSize = pCreateInfo->size;
 
       bool isSparse = (pCreateInfo->flags & (VK_BUFFER_CREATE_SPARSE_BINDING_BIT |
                                              VK_BUFFER_CREATE_SPARSE_RESIDENCY_BIT)) != 0;
@@ -1994,9 +1999,15 @@ VkResult WrappedVulkan::vkBindBufferMemory2(VkDevice device, uint32_t bindInfoCo
       bufrecord->memOffset = pBindInfos[i].memoryOffset;
 
       // if the buffer was force-referenced, do the same with the memory
-      if(IsForcedReference(GetResID(pBindInfos[i].buffer)))
+      if(IsForcedReference(bufrecord))
       {
-        AddForcedReference(GetResID(pBindInfos[i].memory), eFrameRef_ReadBeforeWrite);
+        // in case we're currently capturing, immediately consider the buffer and backing memory as
+        // read-before-write referenced
+        GetResourceManager()->MarkResourceFrameReferenced(bufrecord->GetResourceID(),
+                                                          eFrameRef_ReadBeforeWrite);
+        GetResourceManager()->MarkMemoryFrameReferenced(
+            GetResID(pBindInfos[i].memory), pBindInfos[i].memoryOffset, bufrecord->memSize,
+            eFrameRef_ReadBeforeWrite);
 
         // the memory is immediately dirty because we have no way of tracking writes to it
         GetResourceManager()->MarkDirtyResource(GetResID(pBindInfos[i].memory));