From 4001d197e179c25887edaabc54ae911db82e26e2 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Tue, 7 May 2024 15:32:11 +0100
Subject: [PATCH] Serialise handle base for descriptor heaps to GPU unwrap on
 replay

---
 renderdoc/driver/d3d12/d3d12_common.cpp      |  4 ++++
 renderdoc/driver/d3d12/d3d12_device.h        |  2 +-
 renderdoc/driver/d3d12/d3d12_device_wrap.cpp | 15 ++++++++++++++-
 renderdoc/driver/d3d12/d3d12_manager.cpp     |  4 ++--
 renderdoc/driver/d3d12/d3d12_resources.cpp   |  2 ++
 renderdoc/driver/d3d12/d3d12_resources.h     |  8 ++++++++
 6 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/renderdoc/driver/d3d12/d3d12_common.cpp b/renderdoc/driver/d3d12/d3d12_common.cpp
index e9920b2be..b0c7d13d9 100644
--- a/renderdoc/driver/d3d12/d3d12_common.cpp
+++ b/renderdoc/driver/d3d12/d3d12_common.cpp
@@ -419,6 +419,10 @@ bool D3D12InitParams::IsSupportedVersion(uint64_t ver)
   if(ver == 0x10)
     return true;
 
+  // 0x11 -> 0x12 - Descriptor heaps serialise the original pointer to their descriptor array for GPU unwrapping
+  if(ver == 0x11)
+    return true;
+
   return false;
 }
 
diff --git a/renderdoc/driver/d3d12/d3d12_device.h b/renderdoc/driver/d3d12/d3d12_device.h
index d94281785..cdaa4f34d 100644
--- a/renderdoc/driver/d3d12/d3d12_device.h
+++ b/renderdoc/driver/d3d12/d3d12_device.h
@@ -54,7 +54,7 @@ struct D3D12InitParams
   UINT SDKVersion = 0;
 
   // check if a frame capture section version is supported
-  static const uint64_t CurrentVersion = 0x11;
+  static const uint64_t CurrentVersion = 0x12;
 
   static bool IsSupportedVersion(uint64_t ver);
 };
diff --git a/renderdoc/driver/d3d12/d3d12_device_wrap.cpp b/renderdoc/driver/d3d12/d3d12_device_wrap.cpp
index 4a7d7b67f..1c6e15196 100644
--- a/renderdoc/driver/d3d12/d3d12_device_wrap.cpp
+++ b/renderdoc/driver/d3d12/d3d12_device_wrap.cpp
@@ -912,6 +912,14 @@ bool WrappedID3D12Device::Serialise_CreateDescriptorHeap(
   SERIALISE_ELEMENT_LOCAL(pHeap, ((WrappedID3D12DescriptorHeap *)*ppvHeap)->GetResourceID())
       .TypedAs("ID3D12DescriptorHeap *"_lit);
 
+  uint64_t originalGPUBase = 0;
+  if(ser.IsWriting())
+    originalGPUBase = ((WrappedID3D12DescriptorHeap *)*ppvHeap)->GetOriginalGPUBase();
+  if(ser.VersionAtLeast(0x12))
+  {
+    SERIALISE_ELEMENT(originalGPUBase).Hidden();
+  }
+
   SERIALISE_CHECK_READ_ERRORS();
 
   if(IsReplayingAndReading())
@@ -959,7 +967,12 @@ bool WrappedID3D12Device::Serialise_CreateDescriptorHeap(
     }
     else
     {
-      ret = new WrappedID3D12DescriptorHeap(ret, this, PatchedDesc, Descriptor.NumDescriptors);
+      WrappedID3D12DescriptorHeap *wrapped =
+          new WrappedID3D12DescriptorHeap(ret, this, PatchedDesc, Descriptor.NumDescriptors);
+
+      wrapped->SetOriginalGPUBase(originalGPUBase);
+
+      ret = wrapped;
 
       GetResourceManager()->AddLiveResource(pHeap, ret);
 
diff --git a/renderdoc/driver/d3d12/d3d12_manager.cpp b/renderdoc/driver/d3d12/d3d12_manager.cpp
index fd10a2ebc..5e7d3d428 100644
--- a/renderdoc/driver/d3d12/d3d12_manager.cpp
+++ b/renderdoc/driver/d3d12/d3d12_manager.cpp
@@ -1051,14 +1051,14 @@ PatchedRayDispatch D3D12RaytracingResourceAndUtilHandler::PatchRayDispatch(
 
     if(heap->GetDescriptors()->GetType() == D3D12DescriptorType::Sampler)
     {
-      cbufferData.wrapped_sampHeapBase = heap->GetCPUDescriptorHandleForHeapStart().ptr;
+      cbufferData.wrapped_sampHeapBase = heap->GetOriginalGPUBase();
       cbufferData.unwrapped_sampHeapBase = heap->GetGPU(0).ptr;
       cbufferData.wrapped_sampHeapSize = heap->GetNumDescriptors() * sizeof(D3D12Descriptor);
       cbufferData.unwrapped_heapStrides |= uint16_t(heap->GetUnwrappedIncrement());
     }
     else
     {
-      cbufferData.wrapped_srvHeapBase = heap->GetCPUDescriptorHandleForHeapStart().ptr;
+      cbufferData.wrapped_srvHeapBase = heap->GetOriginalGPUBase();
       cbufferData.unwrapped_srvHeapBase = heap->GetGPU(0).ptr;
       cbufferData.wrapped_srvHeapSize = heap->GetNumDescriptors() * sizeof(D3D12Descriptor);
       cbufferData.unwrapped_heapStrides |= uint32_t(heap->GetUnwrappedIncrement()) << 16;
diff --git a/renderdoc/driver/d3d12/d3d12_resources.cpp b/renderdoc/driver/d3d12/d3d12_resources.cpp
index ea8ed951f..981a947f2 100644
--- a/renderdoc/driver/d3d12/d3d12_resources.cpp
+++ b/renderdoc/driver/d3d12/d3d12_resources.cpp
@@ -560,6 +560,8 @@ WrappedID3D12DescriptorHeap::WrappedID3D12DescriptorHeap(ID3D12DescriptorHeap *r
 
   descriptors = new D3D12Descriptor[desc.NumDescriptors];
 
+  m_OriginalWrappedGPUBase = (uint64_t)descriptors;
+
   RDCEraseMem(descriptors, sizeof(D3D12Descriptor) * desc.NumDescriptors);
   for(UINT i = 0; i < desc.NumDescriptors; i++)
     descriptors[i].Setup(this, i);
diff --git a/renderdoc/driver/d3d12/d3d12_resources.h b/renderdoc/driver/d3d12/d3d12_resources.h
index 870a3749d..25b5d3420 100644
--- a/renderdoc/driver/d3d12/d3d12_resources.h
+++ b/renderdoc/driver/d3d12/d3d12_resources.h
@@ -402,6 +402,11 @@ class WrappedID3D12DescriptorHeap : public WrappedDeviceChild12<ID3D12Descriptor
   Descriptor *cachedDescriptors;
   uint64_t *mutableDescriptorBitmask;
 
+  // for GPU handles that sit in GPU memory, this is the base of our descriptors array above during
+  // capture time. When capturing this == descriptors, on replay it is the value that descriptors had
+  // (which applications then queried and passed to the GPU) which is used for GPU-unwrapping handles
+  uint64_t m_OriginalWrappedGPUBase;
+
 public:
   ALLOCATE_WITH_WRAPPED_POOL(WrappedID3D12DescriptorHeap);
 
@@ -442,6 +447,9 @@ public:
     return handle;
   }
 
+  void SetOriginalGPUBase(uint64_t base) { m_OriginalWrappedGPUBase = base; }
+  uint64_t GetOriginalGPUBase() { return m_OriginalWrappedGPUBase; }
+
   D3D12_CPU_DESCRIPTOR_HANDLE GetCPU(uint32_t idx)
   {
     D3D12_CPU_DESCRIPTOR_HANDLE handle = realCPUBase;