From 84dc1208719e5ddadff731a3306f29bf356d189f Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Thu, 9 May 2024 10:54:11 +0100
Subject: [PATCH] Fix GPU patching of handles in local root signatures

---
 renderdoc/data/hlsl/raytracing.hlsl      | 16 ++++++++++------
 renderdoc/driver/d3d12/d3d12_manager.cpp |  8 ++++++++
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/renderdoc/data/hlsl/raytracing.hlsl b/renderdoc/data/hlsl/raytracing.hlsl
index 6ceed92fc..5086e39ff 100644
--- a/renderdoc/data/hlsl/raytracing.hlsl
+++ b/renderdoc/data/hlsl/raytracing.hlsl
@@ -136,9 +136,11 @@ void PatchTable(uint byteOffset)
   bufferToPatch.Store4(byteOffset, recordData.identifier[0]);
   bufferToPatch.Store4(byteOffset + 16, recordData.identifier[1]);
 
-  if(recordData.rootSigIndex & 0xffff != 0xffff)
+  uint rootSigIndex = (recordData.rootSigIndex & 0xffff);
+
+  if(rootSigIndex != 0xffff)
   {
-    RootSig sig = rootsigs[recordData.rootSigIndex];
+    RootSig sig = rootsigs[rootSigIndex];
 
     DescriptorHeapData heaps[2];
 
@@ -156,7 +158,7 @@ void PatchTable(uint byteOffset)
 
     for(uint i = 0; i < sig.numHandles; i++)
     {
-      GPUAddress wrappedHandlePtr = bufferToPatch.Load2(sig.handleOffsets[i]);
+      GPUAddress wrappedHandlePtr = bufferToPatch.Load2(byteOffset + sig.handleOffsets[i]);
 
       bool patched = false;
       for(int h = 0; h < 2; h++)
@@ -165,9 +167,11 @@ void PatchTable(uint byteOffset)
            lessThan(wrappedHandlePtr, heaps[h].wrapped_end))
         {
           // assume the byte offsets will all fit into the LSB 32-bits
-          uint index = sub(wrappedHandlePtr, wrapped_sampHeapBase).x / WRAPPED_DESCRIPTOR_STRIDE;
+          uint index = sub(wrappedHandlePtr, heaps[h].wrapped_base).x / WRAPPED_DESCRIPTOR_STRIDE;
+
           GPUAddress handleOffset = GPUAddress(index * heaps[h].unwrapped_stride, 0);
-          bufferToPatch.Store2(sig.handleOffsets[i], add(heaps[h].unwrapped_base, handleOffset));
+          bufferToPatch.Store2(byteOffset + sig.handleOffsets[i],
+                               add(heaps[h].unwrapped_base, handleOffset));
           patched = true;
         }
       }
@@ -175,7 +179,7 @@ void PatchTable(uint byteOffset)
       if(!patched)
       {
         // won't work but is our best effort
-        bufferToPatch.Store2(sig.handleOffsets[i], GPUAddress(0, 0));
+        bufferToPatch.Store2(byteOffset + sig.handleOffsets[i], GPUAddress(0, 0));
       }
     }
   }
diff --git a/renderdoc/driver/d3d12/d3d12_manager.cpp b/renderdoc/driver/d3d12/d3d12_manager.cpp
index db1a0b10d..bbac80b20 100644
--- a/renderdoc/driver/d3d12/d3d12_manager.cpp
+++ b/renderdoc/driver/d3d12/d3d12_manager.cpp
@@ -872,6 +872,14 @@ PatchedRayDispatch D3D12RaytracingResourceAndUtilHandler::PatchRayDispatch(
         exportIndex += (uint32_t)m_ExportDatabases[i]->ownExports.size();
       }
 
+      for(size_t i = 0; i < m_UniqueLocalRootSigs.size(); i++)
+      {
+        uint32_t *rootSigData = (uint32_t *)(lookupData.data() + RootSigOffset + RootSigStride * i);
+
+        rootSigData[0] = (uint32_t)m_UniqueLocalRootSigs[i].size();
+        memcpy(&rootSigData[1], m_UniqueLocalRootSigs[i].data(), m_UniqueLocalRootSigs[i].byteSize());
+      }
+
       D3D12GpuBufferAllocator::Inst()->Alloc(D3D12GpuBufferHeapType::UploadHeap,
                                              D3D12GpuBufferHeapMemoryFlag::Default,
                                              lookupData.size(), 256, &m_LookupBuffer);