From 19701205efe3ff103f3fcebba5f4814516a4754a Mon Sep 17 00:00:00 2001
From: Amit Prakash <amit.prakash@amd.com>
Date: Mon, 25 Jul 2022 19:40:14 -0400
Subject: [PATCH] Add helper to reference all buffer resources

* There is no easy way to track the referenced resources in
  acceleration structure builds, especially BLAS resource, as there is
  no way to track how and when they will change.
  Now with this change, we will track all the buffer resource
  irrespective they are used in AS or not after the first call to AS
  build.
---
 renderdoc/driver/d3d12/d3d12_resources.cpp |  3 +++
 renderdoc/driver/d3d12/d3d12_resources.h   | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/renderdoc/driver/d3d12/d3d12_resources.cpp b/renderdoc/driver/d3d12/d3d12_resources.cpp
index 5d706784d..37c484a1d 100644
--- a/renderdoc/driver/d3d12/d3d12_resources.cpp
+++ b/renderdoc/driver/d3d12/d3d12_resources.cpp
@@ -28,6 +28,7 @@
 #include "d3d12_command_queue.h"
 
 GPUAddressRangeTracker WrappedID3D12Resource::m_Addresses;
+rdcarray<ResourceId> WrappedID3D12Resource::m_bufferResources;
 std::map<WrappedID3D12PipelineState::DXBCKey, WrappedID3D12Shader *> WrappedID3D12Shader::m_Shaders;
 bool WrappedID3D12Shader::m_InternalResources = false;
 int32_t WrappedID3D12CommandAllocator::m_ResetEnabled = 1;
@@ -170,6 +171,8 @@ WrappedID3D12Resource::~WrappedID3D12Resource()
     range.id = GetResourceID();
 
     m_Addresses.RemoveFrom(range);
+
+    m_bufferResources.removeOne(GetResourceID());
   }
 
   Shutdown();
diff --git a/renderdoc/driver/d3d12/d3d12_resources.h b/renderdoc/driver/d3d12/d3d12_resources.h
index 25305ac64..baa4d4924 100644
--- a/renderdoc/driver/d3d12/d3d12_resources.h
+++ b/renderdoc/driver/d3d12/d3d12_resources.h
@@ -934,6 +934,7 @@ class WrappedID3D12Resource
     : public WrappedDeviceChild12<ID3D12Resource, ID3D12Resource1, ID3D12Resource2>
 {
   static GPUAddressRangeTracker m_Addresses;
+  static rdcarray<ResourceId> m_bufferResources;
 
   WriteSerialiser &GetThreadSerialiser();
   size_t DeleteOverlappingAccStructsInRangeAtOffset(D3D12BufferOffset bufferOffset);
@@ -979,6 +980,13 @@ public:
 
   bool IsAccelerationStructureResource() const { return m_isAccelerationStructureResource; }
   void MarkAsAccelerationStructureResource() { m_isAccelerationStructureResource = true; }
+  static void MarkAllBufferResourceFrameReferenced(D3D12ResourceManager *rm)
+  {
+    for(ResourceId id : m_bufferResources)
+    {
+      rm->MarkResourceFrameReferenced(id, eFrameRef_Read);
+    }
+  }
 
   static void RefBuffers(D3D12ResourceManager *rm);
   static void GetMappableIDs(D3D12ResourceManager *rm, const std::unordered_set<ResourceId> &refdIDs,
@@ -1047,6 +1055,8 @@ public:
       range.id = GetResourceID();
 
       m_Addresses.AddTo(range);
+
+      m_bufferResources.push_back(GetResourceID());
     }
   }
   virtual ~WrappedID3D12Resource();