Add forced references for ASs and mark all ASBs referenced

renderdoc/driver/d3d12/d3d12_command_list4_wrap.cpp

@@ -1005,7 +1005,6 @@ void WrappedID3D12GraphicsCommandList::BuildRaytracingAccelerationStructure(
       m_ListRecord->AddChunk(scope.Get(m_ListRecord->cmdInfo->alloc));
     }
 
-    D3D12ResourceManager *resManager = m_pDevice->GetResourceManager();
     ResourceId asbWrappedResourceId;
     D3D12BufferOffset asbWrappedResourceBufferOffset;
 
@@ -1015,11 +1014,13 @@ void WrappedID3D12GraphicsCommandList::BuildRaytracingAccelerationStructure(
     D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO preBldInfo;
     m_pDevice->GetRaytracingAccelerationStructurePrebuildInfo(&pDesc->Inputs, &preBldInfo);
 
-    auto PostBldExecute = [resManager, asbWrappedResourceId, asbWrappedResourceBufferOffset,
+    auto PostBldExecute = [this, asbWrappedResourceId, asbWrappedResourceBufferOffset,
                            preBldInfo]() -> bool {
       bool success = false;
       D3D12AccelerationStructure *accStructAtOffset = NULL;
 
+      D3D12ResourceManager *resManager = m_pDevice->GetResourceManager();
+
       WrappedID3D12Resource *asbWrappedResource =
           resManager->GetCurrentAs<WrappedID3D12Resource>(asbWrappedResourceId);
 
@@ -1067,6 +1068,8 @@ void WrappedID3D12GraphicsCommandList::BuildRaytracingAccelerationStructure(
           // register this AS so its resource can be created during replay
           m_pDevice->CreateAS(asbWrappedResource, asbWrappedResourceBufferOffset, preBldInfo,
                               accStructAtOffset);
+
+          m_pDevice->AddForcedReference(record);
         }
         else
         {

renderdoc/driver/d3d12/d3d12_commands.cpp

@@ -923,7 +923,6 @@ bool WrappedID3D12CommandQueue::ProcessChunk(ReadSerialiser &ser, D3D12Chunk chu
 
     case D3D12Chunk::List_ClearState: ret = m_ReplayList->Serialise_ClearState(ser, NULL); break;
 
-    /*-----AMD TODO------*/
     case D3D12Chunk::List_BuildRaytracingAccelerationStructure:
       ret = m_ReplayList->Serialise_BuildRaytracingAccelerationStructure(ser, NULL, 0, NULL);
       break;
@@ -943,8 +942,6 @@ bool WrappedID3D12CommandQueue::ProcessChunk(ReadSerialiser &ser, D3D12Chunk chu
       ret = m_ReplayList->Serialise_SetPipelineState1(ser, NULL);
       break;
 
-    /*-----AMD TODO-------*/
-
     // in order to get a warning if we miss a case, we explicitly handle the device creation chunks
     // here. If we actually encounter one it's an error (we shouldn't see these inside the captured
     // frame itself)

renderdoc/driver/d3d12/d3d12_device.cpp

@@ -2623,6 +2623,11 @@ void WrappedID3D12Device::StartFrameCapture(DeviceOwnedWindow devWnd)
   }
 
   GetResourceManager()->MarkResourceFrameReferenced(m_ResourceID, eFrameRef_Read);
+
+  rdcarray<D3D12ResourceRecord *> forced = GetForcedReferences();
+
+  for(auto it = forced.begin(); it != forced.end(); ++it)
+    GetResourceManager()->MarkResourceFrameReferenced((*it)->GetResourceID(), eFrameRef_Read);
 }
 
 bool WrappedID3D12Device::EndFrameCapture(DeviceOwnedWindow devWnd)
@@ -2668,6 +2673,19 @@ bool WrappedID3D12Device::EndFrameCapture(DeviceOwnedWindow devWnd)
 
   rdcarray<WrappedID3D12CommandQueue *> queues;
 
+  // There is no easy way to mark resource referenced in the AS input and the dispatch tables.
+  // We could be more selective and only force-reference acceleration structure buffers - resources
+  // in the AS state - but this would not include scratch buffers (which could be done by hand) and
+  // build input buffers (which can't).
+  // we don't do this with the forced reference system as we want this to be retroactive - only
+  // after seeing an AS build do we mark all buffers referenced but buffers could be created before
+  // an AS is built.
+  CaptureOptions opts = RenderDoc::Inst().GetCaptureOptions();
+  if(!opts.refAllResources && m_HaveSeenASBuild)
+  {
+    WrappedID3D12Resource::MarkAllBufferResourceFrameReferenced(GetResourceManager());
+  }
+
   // transition back to IDLE and readback initial states atomically
   {
     SCOPED_WRITELOCK(m_CapTransitionLock);
@@ -3188,6 +3206,11 @@ void WrappedID3D12Device::ReleaseResource(ID3D12DeviceChild *res)
     m_ResourceStates.erase(id);
   }
 
+  {
+    SCOPED_LOCK(m_ForcedReferencesLock);
+    m_ForcedReferences.removeOne(GetRecord(res));
+  }
+
   {
     SCOPED_LOCK(m_SparseLock);
     m_SparseResources.erase(id);
@@ -3669,6 +3692,8 @@ void WrappedID3D12Device::CreateAS(ID3D12Resource *pResource, UINT64 resourceOff
   {
     D3D12ResourceRecord *record = as->GetResourceRecord();
 
+    m_HaveSeenASBuild = true;
+
     {
       WriteSerialiser &ser = GetThreadSerialiser();
       SCOPED_SERIALISE_CHUNK(D3D12Chunk::CreateAS);

renderdoc/driver/d3d12/d3d12_device.h

@@ -620,6 +620,24 @@ private:
   rdcarray<WrappedID3D12CommandQueue *> m_RefQueues;
   rdcarray<ID3D12Resource *> m_RefBuffers;
 
+  rdcarray<D3D12ResourceRecord *> m_ForcedReferences;
+  Threading::CriticalSection m_ForcedReferencesLock;
+  bool m_HaveSeenASBuild = false;
+
+  int64_t m_QueueCounter = 0;
+
+  rdcarray<D3D12ResourceRecord *> GetForcedReferences()
+  {
+    rdcarray<D3D12ResourceRecord *> ret;
+
+    {
+      SCOPED_LOCK(m_ForcedReferencesLock);
+      ret = m_ForcedReferences;
+    }
+
+    return ret;
+  }
+
   // the queue we use for all internal work, the first DIRECT queue
   WrappedID3D12CommandQueue *m_Queue;
 
@@ -836,6 +854,12 @@ public:
   const D3D12_FEATURE_DATA_D3D12_OPTIONS16 &GetOpts16() { return m_D3D12Opts16; }
   void RemoveQueue(WrappedID3D12CommandQueue *queue);
 
+  void AddForcedReference(D3D12ResourceRecord *record)
+  {
+    SCOPED_LOCK(m_ForcedReferencesLock);
+    m_ForcedReferences.push_back(record);
+  }
+
   // only valid on replay
   const std::map<ResourceId, WrappedID3D12Resource *> &GetResourceList() { return *m_ResourceList; }
   void AddReplayResource(ResourceId id, WrappedID3D12Resource *res) { (*m_ResourceList)[id] = res; }

renderdoc/driver/d3d12/d3d12_resources.h

@@ -957,6 +957,8 @@ public:
     return WrappedDeviceChild12::IsResident();
   }
 
+  WrappedID3D12Heap *GetHeap() { return m_Heap; }
+
   ID3D12Pageable *UnwrappedResidencyPageable()
   {
     if(m_Heap)