For storage resources be pessimistic about their usage in the frame

* If a storage resource is only referenced in a copy as a destination and gets
  CompleteWrite state, this may be inaccurate because we no longer track
  descriptor set access to images or buffers. That means there could be an
  earlier use of the resource in a read/write fashion but on replay we would
  clear it for initial states.
* For images this is easy as we can fix this up at the end of the frame, for
  memory we need to iterate over all storage buffers and update the memrefs on
  their bindings.

renderdoc/driver/vulkan/imagestate_tests.cpp

@@ -137,9 +137,9 @@ TEST_CASE("Test ImageState type", "[imagestate]")
   VkImage image = (VkImage)123;
   VkFormat format = VK_FORMAT_D16_UNORM_S8_UINT;
   VkExtent3D extent = {100, 100, 13};
-  int levelCount = 11;
+  uint16_t levelCount = 11;
-  int layerCount = 17;
+  uint32_t layerCount = 17;
-  int sampleCount = 1;
+  uint16_t sampleCount = 1;
   ImageInfo imageInfo(format, extent, levelCount, layerCount, sampleCount,
                       VK_IMAGE_LAYOUT_UNDEFINED, VK_SHARING_MODE_EXCLUSIVE);
 

renderdoc/driver/vulkan/vk_common.cpp

@@ -1160,6 +1160,8 @@ void DescriptorSetSlot::AccumulateBindRefs(DescriptorBindRefs &refs, VulkanResou
       AddBindFrameRef(refs, bufView->baseResource, eFrameRef_Read);
     if(bufView->baseResourceMem != ResourceId())
       AddMemFrameRef(refs, bufView->baseResourceMem, bufView->memOffset, bufView->memSize, ref);
+    if(bufView->resInfo->storable)
+      refs.storableRefs.insert(rm->GetResourceRecord(bufView->baseResource));
   }
   if(imgView)
   {
@@ -1176,6 +1178,8 @@ void DescriptorSetSlot::AccumulateBindRefs(DescriptorBindRefs &refs, VulkanResou
       refs.sparseRefs.insert(buffer);
     if(buffer->baseResource != ResourceId())
       AddMemFrameRef(refs, buffer->baseResource, buffer->memOffset, buffer->memSize, ref);
+    if(buffer->storable)
+      refs.storableRefs.insert(buffer);
   }
 }
 

renderdoc/driver/vulkan/vk_core.cpp

@@ -1666,7 +1666,14 @@ template <typename SerialiserType>
 bool WrappedVulkan::Serialise_BeginCaptureFrame(SerialiserType &ser)
 {
   SCOPED_LOCK(m_ImageStatesLock);
+
+  for(auto it = m_ImageStates.begin(); it != m_ImageStates.end(); ++it)
+  {
+    it->second.LockWrite()->FixupStorageReferences();
+  }
+
   GetResourceManager()->SerialiseImageStates(ser, m_ImageStates);
+
   SERIALISE_CHECK_READ_ERRORS();
 
   return true;
@@ -2624,6 +2631,11 @@ ReplayStatus WrappedVulkan::ContextReplayLog(CaptureState readType, uint32_t sta
       ser.SkipCurrentChunk();
 #else
     Serialise_BeginCaptureFrame(ser);
+
+    if(IsLoading(m_State))
+    {
+      AddResourceCurChunk(m_InitParams.InstanceID);
+    }
 #endif
   }
 

renderdoc/driver/vulkan/vk_image_states.cpp

@@ -1502,3 +1502,22 @@ void ImageState::BeginCapture()
     it->SetState(state);
   }
 }
+
+void ImageState::FixupStorageReferences()
+{
+  if(m_Storage)
+  {
+    // storage images we don't track the reference to because they're in descriptor sets, so the
+    // read/write state of them is unknown. We can't allow a 'completewrite' to be used as-is
+    // because
+    // there might be a read before then which we just didn't track at the time.
+    maxRefType = ComposeFrameRefsUnordered(maxRefType, eFrameRef_ReadBeforeWrite);
+
+    for(auto it = subresourceStates.begin(); it != subresourceStates.end(); ++it)
+    {
+      ImageSubresourceState state = it->state();
+      state.refType = ComposeFrameRefsUnordered(state.refType, eFrameRef_ReadBeforeWrite);
+      it->SetState(state);
+    }
+  }
+}

renderdoc/driver/vulkan/vk_manager.cpp

@@ -296,7 +296,7 @@ template <typename SerialiserType>
 void VulkanResourceManager::SerialiseImageStates(SerialiserType &ser,
                                                  std::map<ResourceId, LockingImageState> &states)
 {
-  SERIALISE_ELEMENT_LOCAL(NumImages, (uint32_t)states.size());
+  SERIALISE_ELEMENT_LOCAL(NumImages, (uint32_t)states.size()).Important();
 
   auto srcit = states.begin();
 
@@ -951,6 +951,25 @@ void VulkanResourceManager::MergeReferencedMemory(std::unordered_map<ResourceId,
   }
 }
 
+void VulkanResourceManager::FixupStorageBufferMemory(
+    const std::unordered_set<VkResourceRecord *> &storageBuffers)
+{
+  SCOPED_LOCK_OPTIONAL(m_Lock, m_Capturing);
+
+  for(VkResourceRecord *buf : storageBuffers)
+  {
+    ResourceId id = buf->GetResourceID();
+    MemRefs ref;
+    ref.Update(buf->memOffset, buf->memSize, eFrameRef_ReadBeforeWrite);
+
+    auto i = m_MemFrameRefs.find(id);
+    if(i == m_MemFrameRefs.end())
+      m_MemFrameRefs[id] = ref;
+    else
+      i->second.Merge(ref, ComposeFrameRefsUnordered);
+  }
+}
+
 void VulkanResourceManager::ClearReferencedMemory()
 {
   SCOPED_LOCK_OPTIONAL(m_Lock, m_Capturing);

renderdoc/driver/vulkan/vk_manager.h

@@ -422,6 +422,7 @@ public:
   void RemoveDeviceMemory(ResourceId mem);
 
   void MergeReferencedMemory(std::unordered_map<ResourceId, MemRefs> &memRefs);
+  void FixupStorageBufferMemory(const std::unordered_set<VkResourceRecord *> &storageBuffers);
   void ClearReferencedMemory();
   MemRefs *FindMemRefs(ResourceId mem);
   ImgRefs *FindImgRefs(ResourceId img);

renderdoc/driver/vulkan/vk_resources.h

@@ -851,9 +851,10 @@ struct ImageSubresourceRange;
 
 struct ImageInfo
 {
-  int layerCount = 0;
+  uint32_t layerCount = 0;
-  int levelCount = 0;
+  uint16_t levelCount = 0;
-  int sampleCount = 0;
+  uint16_t sampleCount = 0;
+  bool storage = false;
   VkExtent3D extent = {0, 0, 0};
   VkImageType imageType = VK_IMAGE_TYPE_2D;
   VkFormat format = VK_FORMAT_UNDEFINED;
@@ -861,8 +862,8 @@ struct ImageInfo
   VkSharingMode sharingMode = VK_SHARING_MODE_EXCLUSIVE;
   VkImageAspectFlags aspects = 0;
   ImageInfo() {}
-  ImageInfo(VkFormat format, VkExtent3D extent, int levelCount, int layerCount, int sampleCount,
+  ImageInfo(VkFormat format, VkExtent3D extent, uint16_t levelCount, uint32_t layerCount,
-            VkImageLayout initialLayout, VkSharingMode sharingMode)
+            uint16_t sampleCount, VkImageLayout initialLayout, VkSharingMode sharingMode)
       : format(format),
         extent(extent),
         levelCount(levelCount),
@@ -875,8 +876,8 @@ struct ImageInfo
   }
   ImageInfo(const VkImageCreateInfo &ci)
       : layerCount(ci.arrayLayers),
-        levelCount(ci.mipLevels),
+        levelCount((uint16_t)ci.mipLevels),
-        sampleCount((int)ci.samples),
+        sampleCount((uint16_t)ci.samples),
         extent(ci.extent),
         imageType(ci.imageType),
         format(ci.format),
@@ -896,6 +897,11 @@ struct ImageInfo
       extent.depth = 1;
     }
     aspects = FormatImageAspects(format);
+
+    if(ci.usage & VK_IMAGE_USAGE_STORAGE_BIT)
+    {
+      storage = true;
+    }
   }
   ImageInfo(const VkSwapchainCreateInfoKHR &ci)
       : layerCount(ci.imageArrayLayers),
@@ -972,6 +978,7 @@ struct ResourceInfo
   Sparse::PageTable sparseTable;
   rdcarray<AspectSparseTable> altSparseAspects;
   VkImageAspectFlags sparseAspect;
+  bool storable;
 
   Sparse::PageTable &getSparseTableForAspect(VkImageAspectFlags aspects)
   {
@@ -1111,6 +1118,7 @@ struct DescriptorBindRefs
   std::unordered_map<ResourceId, MemRefs> bindMemRefs;
   rdcflatmap<ResourceId, ImageState> bindImageStates;
   std::unordered_set<VkResourceRecord *> sparseRefs;
+  std::unordered_set<VkResourceRecord *> storableRefs;
 };
 
 struct PipelineLayoutData
@@ -1701,6 +1709,7 @@ struct ImageState
   rdcarray<VkImageMemoryBarrier> newQueueFamilyTransfers;
   bool isMemoryBound = false;
   bool m_Overlay = false;
+  bool m_Storage = false;
   ResourceId boundMemory = ResourceId();
   VkDeviceSize boundMemoryOffset = 0ull;
   VkDeviceSize boundMemorySize = 0ull;
@@ -1710,7 +1719,10 @@ struct ImageState
   inline const ImageInfo &GetImageInfo() const { return subresourceStates.GetImageInfo(); }
   inline ImageState() {}
   inline ImageState(VkImage wrappedHandle, const ImageInfo &imageInfo, FrameRefType refType)
-      : wrappedHandle(wrappedHandle), subresourceStates(imageInfo, refType), maxRefType(refType)
+      : wrappedHandle(wrappedHandle),
+        subresourceStates(imageInfo, refType),
+        maxRefType(refType),
+        m_Storage(imageInfo.storage)
   {
   }
   void SetOverlay() { m_Overlay = true; }
@@ -1778,6 +1790,7 @@ struct ImageState
                                uint32_t arrayLayer) const;
 
   void BeginCapture();
+  void FixupStorageReferences();
 };
 
 DECLARE_REFLECTION_STRUCT(ImageState);

renderdoc/driver/vulkan/wrappers/vk_queue_funcs.cpp

@@ -981,6 +981,15 @@ void WrappedVulkan::CaptureQueueSubmit(VkQueue queue,
 
       UpdateImageStates(refs.bindImageStates);
       GetResourceManager()->MergeReferencedMemory(refs.bindMemRefs);
+
+      // for storage buffers we have to pessimise memory references because the order matters - if
+      // the first recorded reference is a complete write then a later readbeforewrite won't
+      // properly mark it as needing initial states preserved. So we do that here. Images are
+      // handled separately
+      for(auto refit = refs.storableRefs.begin(); refit != refs.storableRefs.end(); ++refit)
+      {
+        GetResourceManager()->FixupStorageBufferMemory(refs.storableRefs);
+      }
     }
 
     GetResourceManager()->MarkResourceFrameReferenced(GetResID(queue), eFrameRef_Read);

renderdoc/driver/vulkan/wrappers/vk_resource_funcs.cpp

@@ -1700,6 +1700,8 @@ VkResult WrappedVulkan::vkCreateBuffer(VkDevice device, const VkBufferCreateInfo
       {
         record->resInfo = new ResourceInfo();
 
+        record->resInfo->storable = record->storable;
+
         // pre-populate memory requirements
         ObjDisp(device)->GetBufferMemoryRequirements(Unwrap(device), Unwrap(*pBuffer),
                                                      &record->resInfo->memreqs);
@@ -2254,7 +2256,8 @@ VkResult WrappedVulkan::vkCreateImage(VkDevice device, const VkImageCreateInfo *
       ResourceInfo &resInfo = *record->resInfo;
       resInfo.imageInfo = ImageInfo(*pCreateInfo);
 
-      record->storable = (pCreateInfo->usage & VK_IMAGE_USAGE_STORAGE_BIT) != 0;
+      record->resInfo->storable = record->storable =
+          (pCreateInfo->usage & VK_IMAGE_USAGE_STORAGE_BIT) != 0;
 
       // pre-populate memory requirements
       ObjDisp(device)->GetImageMemoryRequirements(Unwrap(device), Unwrap(*pImage), &resInfo.memreqs);