Add indirect patching of shader records for indirect ray dispatch

* We use a worst-case estimate of how much scratch buffer space will be needed
  as there's no way to know beforehand on the CPU.

renderdoc/data/hlsl/hlsl_cbuffers.h

@@ -34,6 +34,7 @@
 #define float2 Vec2f
 #define float3 Vec3f
 #define uint4 Vec4u
+#define uint3 Vec3u
 #define uint2 Vec2u
 #define int4 Vec4i
 #define float4 Vec4f
@@ -334,6 +335,32 @@ struct LocalRootSigData
 
 #define WRAPPED_DESCRIPTOR_STRIDE 64
 
+cbuffer RayIndirectDispatchCB REG(b0)
+{
+  GPUAddress scratchBuffer;
+
+  uint commandSigDispatchOffset;
+  uint commandSigStride;
+  uint commandSigSize;
+  uint maxCommandCount;    // MaxCommandCount to clamp to. We also set the top bit if there is no count buffer
+};
+
+struct PatchingExecute
+{
+  // D3D12PatchRayDispatchParam::RecordCB
+  uint shaderrecord_stride;
+  uint shaderrecord_count;
+  // D3D12PatchRayDispatchParam::SourceBuffer
+  GPUAddress sourceData;
+  // D3D12PatchRayDispatchParam::DestBuffer
+  GPUAddress destData;
+  // Dispatch itself
+  uint3 dispatchDim;
+  uint padding1;
+
+  uint2 padding2;
+};
+
 cbuffer DebugSampleOperation REG(b0)
 {
   float4 debugSampleUV;

renderdoc/data/hlsl/raytracing.hlsl

@@ -82,7 +82,8 @@ struct DescriptorHeapData
   uint unwrapped_stride;
 };
 
-uint CopyData(uint byteOffset, uint dataOffset, uint dataEnd)
+uint CopyData(ByteAddressBuffer src, RWByteAddressBuffer dst, uint byteOffset, uint dataOffset,
+              uint dataEnd)
 {
   uint4 data = 0;
   uint readBytes = 0;
@@ -90,26 +91,26 @@ uint CopyData(uint byteOffset, uint dataOffset, uint dataEnd)
   // copy larger chunks if possible
   if(dataOffset + 16 <= dataEnd)
   {
-    data = patchSource.Load4(byteOffset + dataOffset);
-    patchDest.Store4(byteOffset + dataOffset, data);
+    data = src.Load4(byteOffset + dataOffset);
+    dst.Store4(byteOffset + dataOffset, data);
     readBytes = 16;
   }
   else if(dataOffset + 12 <= dataEnd)
   {
-    data = patchSource.Load3(byteOffset + dataOffset).xyzx;
-    patchDest.Store3(byteOffset + dataOffset, data.xyz);
+    data = src.Load3(byteOffset + dataOffset).xyzx;
+    dst.Store3(byteOffset + dataOffset, data.xyz);
     readBytes = 12;
   }
   else if(dataOffset + 8 <= dataEnd)
   {
-    data = patchSource.Load2(byteOffset + dataOffset).xyxx;
-    patchDest.Store2(byteOffset + dataOffset, data.xy);
+    data = src.Load2(byteOffset + dataOffset).xyxx;
+    dst.Store2(byteOffset + dataOffset, data.xy);
     readBytes = 8;
   }
   else if(dataOffset + 4 <= dataEnd)
   {
-    data = patchSource.Load(byteOffset + dataOffset).xxxx;
-    patchDest.Store(byteOffset + dataOffset, data.x);
+    data = src.Load(byteOffset + dataOffset).xxxx;
+    dst.Store(byteOffset + dataOffset, data.x);
     readBytes = 4;
   }
   else
@@ -193,7 +194,7 @@ void PatchTable(uint byteOffset)
       // copy any gap from where we last processed
       for(uint b = firstUncopiedByte; b < paramOffset;)
       {
-        b = CopyData(byteOffset, b, paramOffset);
+        b = CopyData(patchSource, patchDest, byteOffset, b, paramOffset);
       }
       firstUncopiedByte = paramOffset + 8;
 
@@ -253,7 +254,7 @@ void PatchTable(uint byteOffset)
     // copy any remaining trailing bytes
     for(uint b = firstUncopiedByte; b < shaderrecord_stride;)
     {
-      b = CopyData(byteOffset, b, shaderrecord_stride);
+      b = CopyData(patchSource, patchDest, byteOffset, b, shaderrecord_stride);
     }
   }
   else
@@ -261,7 +262,7 @@ void PatchTable(uint byteOffset)
     // no root sig data to patch, just copy the whole stride
     for(uint b = firstUncopiedByte; b < shaderrecord_stride;)
     {
-      b = CopyData(byteOffset, b, shaderrecord_stride);
+      b = CopyData(patchSource, patchDest, byteOffset, b, shaderrecord_stride);
     }
   }
 }
@@ -272,3 +273,179 @@ void PatchTable(uint byteOffset)
   if(dispatchThread.x < shaderrecord_count)
     PatchTable(shaderrecord_stride * dispatchThread.x);
 }
+
+// define these structs in hlsl for simplicity
+
+struct D3D12_GPU_VIRTUAL_ADDRESS_RANGE
+{
+  GPUAddress StartAddress;
+  uint2 SizeInBytes;    // we ignore the high order bits in .y
+};
+
+struct D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE
+{
+  GPUAddress StartAddress;
+  uint2 SizeInBytes;      // we ignore the high order bits in .y
+  uint2 StrideInBytes;    // we ignore the high order bits in .y
+};
+
+struct D3D12_DISPATCH_RAYS_DESC
+{
+  D3D12_GPU_VIRTUAL_ADDRESS_RANGE RayGenerationShaderRecord;
+  D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE MissShaderTable;
+  D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE HitGroupTable;
+  D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE CallableShaderTable;
+  uint Width;
+  uint Height;
+  uint Depth;
+};
+
+struct D3D12_DISPATCH_RAYS_DESC_RAW
+{
+  uint4 a[6];
+  uint b;
+};
+
+ByteAddressBuffer applicationExecuteArguments : register(t0);
+ByteAddressBuffer applicationCountBuffer : register(t1);
+RWByteAddressBuffer patchedExecuteArguments : register(u0);
+RWStructuredBuffer<PatchingExecute> internalExecuteArguments : register(u1);
+RWByteAddressBuffer internalExecuteCount : register(u2);
+
+GPUAddress AlignRecordAddress(GPUAddress x)
+{
+  GPUAddress ret = add(x, GPUAddress(0xff, 0));
+  ret.x &= ~0xff;
+  return ret;
+}
+
+[numthreads(1, 1, 1)] void RENDERDOC_PrepareRayIndirectExecuteCS() {
+  uint numCommands = maxCommandCount;
+  if(numCommands & 0x80000000U)
+  {
+    numCommands &= 0x7fffffffU;
+  }
+  else
+  {
+    numCommands = min(numCommands, applicationCountBuffer.Load(0));
+  }
+
+  GPUAddress outputBufferLocation = scratchBuffer;
+  uint dispatchIndex = 0;
+
+  PatchingExecute execute = (PatchingExecute)0xccddeeff;
+
+  // loop over all commands the application ran (probably just one)
+  for(uint i = 0; i < numCommands; i++)
+  {
+    uint commandOffset = i * commandSigStride;
+
+    // load the indirect arguments in batch as possible
+    D3D12_DISPATCH_RAYS_DESC_RAW raw = {
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 0 * 16),
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 1 * 16),
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 2 * 16),
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 3 * 16),
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 4 * 16),
+        applicationExecuteArguments.Load4(commandSigDispatchOffset + commandOffset + 5 * 16),
+        applicationExecuteArguments.Load(commandSigDispatchOffset + commandOffset + 6 * 16),
+    };
+
+    D3D12_DISPATCH_RAYS_DESC appCommand = (D3D12_DISPATCH_RAYS_DESC)raw;
+
+    // align the address before each record we're going to patch into
+    outputBufferLocation = AlignRecordAddress(outputBufferLocation);
+
+    // raygen - set up an execute to patch the record
+    {
+      execute.shaderrecord_stride = appCommand.RayGenerationShaderRecord.SizeInBytes.x;
+      execute.shaderrecord_count = 1;
+      execute.sourceData = appCommand.RayGenerationShaderRecord.StartAddress;
+      execute.destData = outputBufferLocation;
+      execute.dispatchDim = uint3(1, 1, 1);
+      internalExecuteArguments[dispatchIndex++] = execute;
+
+      // update the dispatch desc to point to the new location
+      appCommand.RayGenerationShaderRecord.StartAddress = outputBufferLocation;
+
+      outputBufferLocation =
+          add(outputBufferLocation, appCommand.RayGenerationShaderRecord.SizeInBytes);
+    }
+
+    outputBufferLocation = AlignRecordAddress(outputBufferLocation);
+
+    if(appCommand.MissShaderTable.SizeInBytes.x > 0)
+    {
+      execute.shaderrecord_stride = appCommand.MissShaderTable.StrideInBytes.x;
+      execute.shaderrecord_count = appCommand.MissShaderTable.SizeInBytes.x /
+                                   max(1, appCommand.MissShaderTable.StrideInBytes.x);
+      execute.sourceData = appCommand.MissShaderTable.StartAddress;
+      execute.destData = outputBufferLocation;
+      execute.dispatchDim = uint3(
+          (execute.shaderrecord_count + (RECORD_PATCH_THREADS - 1)) / RECORD_PATCH_THREADS, 1, 1);
+      internalExecuteArguments[dispatchIndex++] = execute;
+
+      appCommand.MissShaderTable.StartAddress = outputBufferLocation;
+
+      outputBufferLocation = add(outputBufferLocation, appCommand.MissShaderTable.SizeInBytes);
+    }
+
+    outputBufferLocation = AlignRecordAddress(outputBufferLocation);
+
+    if(appCommand.HitGroupTable.SizeInBytes.x > 0)
+    {
+      execute.shaderrecord_stride = appCommand.HitGroupTable.StrideInBytes.x;
+      execute.shaderrecord_count =
+          appCommand.HitGroupTable.SizeInBytes.x / max(1, appCommand.HitGroupTable.StrideInBytes.x);
+      execute.sourceData = appCommand.HitGroupTable.StartAddress;
+      execute.destData = outputBufferLocation;
+      execute.dispatchDim = uint3(
+          (execute.shaderrecord_count + (RECORD_PATCH_THREADS - 1)) / RECORD_PATCH_THREADS, 1, 1);
+      internalExecuteArguments[dispatchIndex++] = execute;
+
+      appCommand.HitGroupTable.StartAddress = outputBufferLocation;
+
+      outputBufferLocation = add(outputBufferLocation, appCommand.HitGroupTable.SizeInBytes);
+    }
+
+    outputBufferLocation = AlignRecordAddress(outputBufferLocation);
+
+    if(appCommand.CallableShaderTable.SizeInBytes.x > 0)
+    {
+      execute.shaderrecord_stride = appCommand.CallableShaderTable.StrideInBytes.x;
+      execute.shaderrecord_count = appCommand.CallableShaderTable.SizeInBytes.x /
+                                   max(1, appCommand.CallableShaderTable.StrideInBytes.x);
+      execute.sourceData = appCommand.CallableShaderTable.StartAddress;
+      execute.destData = outputBufferLocation;
+      execute.dispatchDim = uint3(
+          (execute.shaderrecord_count + (RECORD_PATCH_THREADS - 1)) / RECORD_PATCH_THREADS, 1, 1);
+      internalExecuteArguments[dispatchIndex++] = execute;
+
+      appCommand.CallableShaderTable.StartAddress = outputBufferLocation;
+
+      outputBufferLocation = add(outputBufferLocation, appCommand.CallableShaderTable.SizeInBytes);
+    }
+
+    // copy all data apart from the indirect draw (if anything) to the patched argument buffer
+    // no root sig data to patch, just copy the whole stride
+    for(uint b = 0; b < commandSigDispatchOffset;)
+    {
+      b = CopyData(applicationExecuteArguments, patchedExecuteArguments, commandOffset, b,
+                   commandSigDispatchOffset);
+    }
+
+    // store the indirect command (with updated addresses) to the patched buffer
+    raw = (D3D12_DISPATCH_RAYS_DESC_RAW)appCommand;
+
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 0 * 16, raw.a[0]);
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 1 * 16, raw.a[1]);
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 2 * 16, raw.a[2]);
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 3 * 16, raw.a[3]);
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 4 * 16, raw.a[4]);
+    patchedExecuteArguments.Store4(commandSigDispatchOffset + commandOffset + 5 * 16, raw.a[5]);
+    patchedExecuteArguments.Store(commandSigDispatchOffset + commandOffset + 6 * 16, raw.b);
+  }
+
+  // store the number of patching indirect dispatches we'll do, up to 4 per the application's number
+  internalExecuteCount.Store(0, dispatchIndex);
+}

renderdoc/driver/d3d12/d3d12_command_list_wrap.cpp

@@ -391,8 +391,9 @@ HRESULT WrappedID3D12GraphicsCommandList::ResetInternal(ID3D12CommandAllocator *
     // holding references until their fences are appropriately signalled.
     for(PatchedRayDispatch::Resources &r : m_RayDispatches)
     {
-      r.lookupBuffer->Release();
-      r.patchScratchBuffer->Release();
+      SAFE_RELEASE(r.lookupBuffer);
+      SAFE_RELEASE(r.patchScratchBuffer);
+      SAFE_RELEASE(r.argumentBuffer);
     }
     m_RayDispatches.clear();
 
@@ -4075,6 +4076,8 @@ bool WrappedID3D12GraphicsCommandList::Serialise_ExecuteIndirect(
 
     BakedCmdListInfo &cmdInfo = m_Cmd->m_BakedCmdListInfo[m_Cmd->m_LastCmdListID];
 
+    const D3D12RenderState &state = cmdInfo.state;
+
     if(IsActiveReplaying(m_State))
     {
       WrappedID3D12CommandSignature *comSig = (WrappedID3D12CommandSignature *)pCommandSignature;
@@ -4200,6 +4203,32 @@ bool WrappedID3D12GraphicsCommandList::Serialise_ExecuteIndirect(
                   (pCountBuffer ? pCountBuffer->GetGPUVirtualAddress() : 0) + CountBufferOffset,
                   MaxCommandCount);
 
+          ID3D12Resource *argBuffer = Unwrap(patched.first);
+          UINT64 argOffset = patched.second;
+
+          PatchedRayDispatch patchedDispatch = {};
+          if(comSig->sig.raytraced)
+          {
+            patchedDispatch =
+                GetResourceManager()->GetRaytracingResourceAndUtilHandler()->PatchIndirectRayDispatch(
+                    Unwrap(pCommandList), state.heaps, comSig, MaxCommandCount, patched.first,
+                    patched.second, pCountBuffer, CountBufferOffset);
+
+            argBuffer = patchedDispatch.resources.argumentBuffer->Resource();
+            argOffset = patchedDispatch.resources.argumentBuffer->Offset();
+
+            // restore state that would have been mutated by the patching process
+            Unwrap(pCommandList)
+                ->SetComputeRootSignature(Unwrap(
+                    GetResourceManager()->GetCurrentAs<ID3D12RootSignature>(state.compute.rootsig)));
+            Unwrap4((ID3D12GraphicsCommandList4 *)pCommandList)
+                ->SetPipelineState1(
+                    Unwrap(GetResourceManager()->GetCurrentAs<ID3D12StateObject>(state.stateobj)));
+            state.ApplyComputeRootElementsUnwrapped(Unwrap(pCommandList));
+          }
+
+          m_Cmd->m_RayDispatches.push_back(patchedDispatch.resources);
+
           if(m_Cmd->m_FirstEventID <= 1)
           {
             // if we're replaying part-way into a multidraw we just clamp the count
@@ -4238,8 +4267,8 @@ bool WrappedID3D12GraphicsCommandList::Serialise_ExecuteIndirect(
           }
 
           if(count > 0)
-            Unwrap(list)->ExecuteIndirect(Unwrap(pCommandSignature), count, Unwrap(patched.first),
-                                          patched.second, NULL, 0);
+            Unwrap(list)->ExecuteIndirect(Unwrap(pCommandSignature), count, argBuffer, argOffset,
+                                          NULL, 0);
         }
       }
 
@@ -4276,8 +4305,34 @@ bool WrappedID3D12GraphicsCommandList::Serialise_ExecuteIndirect(
           (pCountBuffer ? pCountBuffer->GetGPUVirtualAddress() : 0) + CountBufferOffset,
           MaxCommandCount);
 
-      Unwrap(list)->ExecuteIndirect(comSig->GetReal(), MaxCommandCount, Unwrap(patched.first),
-                                    patched.second, Unwrap(pCountBuffer), CountBufferOffset);
+      ID3D12Resource *argBuffer = Unwrap(patched.first);
+      UINT64 argOffset = patched.second;
+
+      PatchedRayDispatch patchedDispatch = {};
+      if(comSig->sig.raytraced)
+      {
+        patchedDispatch =
+            GetResourceManager()->GetRaytracingResourceAndUtilHandler()->PatchIndirectRayDispatch(
+                Unwrap(list), state.heaps, comSig, MaxCommandCount, patched.first, patched.second,
+                pCountBuffer, CountBufferOffset);
+
+        argBuffer = patchedDispatch.resources.argumentBuffer->Resource();
+        argOffset = patchedDispatch.resources.argumentBuffer->Offset();
+
+        // restore state that would have been mutated by the patching process
+        Unwrap(pCommandList)
+            ->SetComputeRootSignature(Unwrap(
+                GetResourceManager()->GetCurrentAs<ID3D12RootSignature>(state.compute.rootsig)));
+        Unwrap4((ID3D12GraphicsCommandList4 *)pCommandList)
+            ->SetPipelineState1(
+                Unwrap(GetResourceManager()->GetCurrentAs<ID3D12StateObject>(state.stateobj)));
+        state.ApplyComputeRootElementsUnwrapped(Unwrap(pCommandList));
+      }
+
+      m_Cmd->m_RayDispatches.push_back(patchedDispatch.resources);
+
+      Unwrap(list)->ExecuteIndirect(comSig->GetReal(), MaxCommandCount, argBuffer, argOffset,
+                                    Unwrap(pCountBuffer), CountBufferOffset);
 
       const uint32_t sigSize = (uint32_t)comSig->sig.arguments.size();
 
@@ -4362,9 +4417,30 @@ void WrappedID3D12GraphicsCommandList::ExecuteIndirect(ID3D12CommandSignature *p
                                                        ID3D12Resource *pCountBuffer,
                                                        UINT64 CountBufferOffset)
 {
-  SERIALISE_TIME_CALL(m_pList->ExecuteIndirect(Unwrap(pCommandSignature), MaxCommandCount,
-                                               Unwrap(pArgumentBuffer), ArgumentBufferOffset,
-                                               Unwrap(pCountBuffer), CountBufferOffset));
+  ID3D12Resource *argBuffer = Unwrap(pArgumentBuffer);
+  UINT64 argOffset = ArgumentBufferOffset;
+
+  PatchedRayDispatch patchedDispatch = {};
+  if(((WrappedID3D12CommandSignature *)pCommandSignature)->sig.raytraced)
+  {
+    patchedDispatch =
+        GetResourceManager()->GetRaytracingResourceAndUtilHandler()->PatchIndirectRayDispatch(
+            m_pList, m_CaptureComputeState.heaps, pCommandSignature, MaxCommandCount,
+            pArgumentBuffer, ArgumentBufferOffset, pCountBuffer, CountBufferOffset);
+
+    argBuffer = patchedDispatch.resources.argumentBuffer->Resource();
+    argOffset = patchedDispatch.resources.argumentBuffer->Offset();
+
+    // restore state that would have been mutated by the patching process
+    m_pList->SetComputeRootSignature(Unwrap(GetResourceManager()->GetCurrentAs<ID3D12RootSignature>(
+        m_CaptureComputeState.compute.rootsig)));
+    m_pList4->SetPipelineState1(Unwrap(
+        GetResourceManager()->GetCurrentAs<ID3D12StateObject>(m_CaptureComputeState.stateobj)));
+    m_CaptureComputeState.ApplyComputeRootElementsUnwrapped(m_pList);
+  }
+
+  SERIALISE_TIME_CALL(m_pList->ExecuteIndirect(Unwrap(pCommandSignature), MaxCommandCount, argBuffer,
+                                               argOffset, Unwrap(pCountBuffer), CountBufferOffset));
   if(IsCaptureMode(m_State))
   {
     CACHE_THREAD_SERIALISER();
@@ -4380,6 +4456,13 @@ void WrappedID3D12GraphicsCommandList::ExecuteIndirect(ID3D12CommandSignature *p
     m_ListRecord->MarkResourceFrameReferenced(GetResID(pCommandSignature), eFrameRef_Read);
     m_ListRecord->MarkResourceFrameReferenced(GetResID(pArgumentBuffer), eFrameRef_Read);
     m_ListRecord->MarkResourceFrameReferenced(GetResID(pCountBuffer), eFrameRef_Read);
+
+    // during capture track the ray dispatches so the memory can be freed dynamically. On replay we
+    if(patchedDispatch.resources.lookupBuffer)
+    {
+      // free all the memory at the end of each replay
+      m_RayDispatches.push_back(patchedDispatch.resources);
+    }
   }
 }
 

renderdoc/driver/d3d12/d3d12_commands.cpp

@@ -643,6 +643,7 @@ void WrappedID3D12CommandQueue::CheckAndFreeRayDispatches()
     {
       SAFE_RELEASE(ray.patchScratchBuffer);
       SAFE_RELEASE(ray.lookupBuffer);
+      SAFE_RELEASE(ray.argumentBuffer);
     }
   }
 
@@ -1434,8 +1435,9 @@ void WrappedID3D12GraphicsCommandList::AddRayDispatches(rdcarray<PatchedRayDispa
   for(const PatchedRayDispatch::Resources &r : m_RayDispatches)
   {
     dispatches.push_back(r);
-    r.lookupBuffer->AddRef();
-    r.patchScratchBuffer->AddRef();
+    SAFE_ADDREF(r.lookupBuffer);
+    SAFE_ADDREF(r.patchScratchBuffer);
+    SAFE_ADDREF(r.argumentBuffer);
   }
 }
 

renderdoc/driver/d3d12/d3d12_device.cpp

@@ -4958,8 +4958,9 @@ void WrappedID3D12Device::ReplayLog(uint32_t startEventID, uint32_t endEventID,
 
     for(PatchedRayDispatch::Resources &r : cmd.m_RayDispatches)
     {
-      r.lookupBuffer->Release();
-      r.patchScratchBuffer->Release();
+      SAFE_RELEASE(r.lookupBuffer);
+      SAFE_RELEASE(r.patchScratchBuffer);
+      SAFE_RELEASE(r.argumentBuffer);
     }
     cmd.m_RayDispatches.clear();
 

renderdoc/driver/d3d12/d3d12_device_wrap.cpp

@@ -1813,76 +1813,8 @@ bool WrappedID3D12Device::Serialise_CreateCommandSignature(SerialiserType &ser,
     }
     else
     {
-      WrappedID3D12CommandSignature *wrapped = new WrappedID3D12CommandSignature(ret, this);
-
-      wrapped->sig.ByteStride = Descriptor.ByteStride;
-      wrapped->sig.arguments.assign(Descriptor.pArgumentDescs, Descriptor.NumArgumentDescs);
-
-      wrapped->sig.graphics = true;
-      wrapped->sig.PackedByteSize = 0;
-
-      // From MSDN, command signatures are either graphics or compute so just search for dispatches:
-      // "A given command signature is either an action or a compute command signature. If a command
-      // signature contains a drawing operation, then it is a graphics command signature. Otherwise,
-      // the command signature must contain a dispatch operation, and it is a compute command
-      // signature."
-      for(uint32_t i = 0; i < Descriptor.NumArgumentDescs; i++)
-      {
-        switch(Descriptor.pArgumentDescs[i].Type)
-        {
-          case D3D12_INDIRECT_ARGUMENT_TYPE_DRAW:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_DRAW_ARGUMENTS);
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_DRAW_INDEXED:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_DRAW_INDEXED_ARGUMENTS);
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_DISPATCH_ARGUMENTS);
-            wrapped->sig.graphics = false;
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH_MESH:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_DISPATCH_MESH_ARGUMENTS);
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH_RAYS:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_DISPATCH_RAYS_DESC);
-            wrapped->sig.raytraced = true;
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_CONSTANT:
-          {
-            wrapped->sig.PackedByteSize +=
-                sizeof(uint32_t) * Descriptor.pArgumentDescs[i].Constant.Num32BitValuesToSet;
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_VERTEX_BUFFER_VIEW:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_VERTEX_BUFFER_VIEW);
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_INDEX_BUFFER_VIEW:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_INDEX_BUFFER_VIEW);
-            break;
-          }
-          case D3D12_INDIRECT_ARGUMENT_TYPE_CONSTANT_BUFFER_VIEW:
-          case D3D12_INDIRECT_ARGUMENT_TYPE_SHADER_RESOURCE_VIEW:
-          case D3D12_INDIRECT_ARGUMENT_TYPE_UNORDERED_ACCESS_VIEW:
-          {
-            wrapped->sig.PackedByteSize += sizeof(D3D12_GPU_VIRTUAL_ADDRESS);
-            break;
-          }
-          default: RDCERR("Unexpected argument type! %d", Descriptor.pArgumentDescs[i].Type); break;
-        }
-      }
+      WrappedID3D12CommandSignature *wrapped =
+          new WrappedID3D12CommandSignature(ret, this, Descriptor);
 
       ret = wrapped;
 
@@ -1929,7 +1861,7 @@ HRESULT WrappedID3D12Device::CreateCommandSignature(const D3D12_COMMAND_SIGNATUR
         return ret;
       }
 
-      wrapped = new WrappedID3D12CommandSignature(real, this);
+      wrapped = new WrappedID3D12CommandSignature(real, this, *pDesc);
     }
 
     if(IsCaptureMode(m_State))

renderdoc/driver/d3d12/d3d12_manager.cpp

@@ -24,6 +24,7 @@
 
 #include "d3d12_manager.h"
 #include <algorithm>
+#include "core/settings.h"
 #include "driver/dx/official/d3dcompiler.h"
 #include "driver/dxgi/dxgi_common.h"
 #include "d3d12_command_list.h"
@@ -34,6 +35,9 @@
 
 #include "data/hlsl/hlsl_cbuffers.h"
 
+RDOC_CONFIG(uint32_t, D3D12_Debug_RTIndirectEstimateOverride, 0,
+            "Override how many bytes are reserved for shader tables in each indirect ray dispatch");
+
 void D3D12Descriptor::Init(const D3D12_SAMPLER_DESC2 *pDesc)
 {
   if(pDesc)
@@ -891,8 +895,8 @@ PatchedRayDispatch D3D12RaytracingResourceAndUtilHandler::PatchRayDispatch(
 
   // set up general patching data - lookup buffers and so on
 
-  unwrappedCmd->SetPipelineState(m_RayPatchingData.pipe);
-  unwrappedCmd->SetComputeRootSignature(m_RayPatchingData.rootSig);
+  unwrappedCmd->SetPipelineState(m_RayPatchingData.descPatchPipe);
+  unwrappedCmd->SetComputeRootSignature(m_RayPatchingData.descPatchRootSig);
   unwrappedCmd->SetComputeRoot32BitConstants((UINT)D3D12PatchRayDispatchParam::GeneralCB,
                                              sizeof(cbufferData) / sizeof(uint32_t), &cbufferData, 0);
   unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::StateObjectData,
@@ -994,6 +998,165 @@ PatchedRayDispatch D3D12RaytracingResourceAndUtilHandler::PatchRayDispatch(
   // ref we had when we created it.
   ret.resources.patchScratchBuffer = scratchBuffer;
 
+  ret.resources.argumentBuffer = NULL;
+
+  return ret;
+}
+
+PatchedRayDispatch D3D12RaytracingResourceAndUtilHandler::PatchIndirectRayDispatch(
+    ID3D12GraphicsCommandList *unwrappedCmd, rdcarray<ResourceId> heaps,
+    ID3D12CommandSignature *pCommandSignature, UINT MaxCommandCount, ID3D12Resource *pArgumentBuffer,
+    UINT64 ArgumentBufferOffset, ID3D12Resource *pCountBuffer, UINT64 CountBufferOffset)
+{
+  PatchedRayDispatch ret = {};
+
+  D3D12MarkerRegion region(unwrappedCmd, "PatchIndirectRayDispatch");
+
+  PrepareRayDispatchBuffer(NULL);
+
+  D3D12GpuBuffer *scratchBuffer = NULL;
+
+  // :( some games have fixed sizes, so any other estimate could fail.
+  // games without fixed sizes would be reasonably served by 3 or 4 multiplied by the number of BLAS
+  // in the largest TLAS, multiplied by the largest local root signature size.
+  uint32_t patchDataSize = 20 * 1024 * 1024 * MaxCommandCount;
+
+  if(D3D12_Debug_RTIndirectEstimateOverride() > 0)
+    patchDataSize = RDCMAX(patchDataSize, D3D12_Debug_RTIndirectEstimateOverride());
+
+  m_GPUBufferAllocator.Alloc(D3D12GpuBufferHeapType::DefaultHeapWithUav,
+                             D3D12GpuBufferHeapMemoryFlag::Default, patchDataSize,
+                             D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT, &scratchBuffer);
+
+  WrappedID3D12CommandSignature *comSig = (WrappedID3D12CommandSignature *)pCommandSignature;
+
+  D3D12GpuBuffer *argsBuffer = NULL;
+
+  // the args buffer contains both patched arguments for the application as well as patched
+  // arguments & count for our indirect patching
+
+  uint64_t applicationArgsSize = AlignUp16(MaxCommandCount * comSig->sig.ByteStride);
+  uint64_t patchingArgsSize = AlignUp16(MaxCommandCount * 4 * AlignUp16(sizeof(PatchingExecute)));
+
+  m_GPUBufferAllocator.Alloc(D3D12GpuBufferHeapType::DefaultHeapWithUav,
+                             D3D12GpuBufferHeapMemoryFlag::Default,
+                             applicationArgsSize + patchingArgsSize + 4,
+                             D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT, &argsBuffer);
+
+  RDCCOMPILE_ASSERT(WRAPPED_DESCRIPTOR_STRIDE == sizeof(D3D12Descriptor),
+                    "Shader descriptor stride is wrong");
+
+  RayDispatchPatchCB cbufferData = {};
+
+  for(ResourceId heapId : heaps)
+  {
+    WrappedID3D12DescriptorHeap *heap =
+        (WrappedID3D12DescriptorHeap *)m_wrappedDevice->GetResourceManager()
+            ->GetCurrentAs<ID3D12DescriptorHeap>(heapId);
+
+    if(heap->GetDescriptors()->GetType() == D3D12DescriptorType::Sampler)
+    {
+      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->GetOriginalGPUBase();
+      cbufferData.unwrapped_srvHeapBase = heap->GetGPU(0).ptr;
+      cbufferData.wrapped_srvHeapSize = heap->GetNumDescriptors() * sizeof(D3D12Descriptor);
+      cbufferData.unwrapped_heapStrides |= uint32_t(heap->GetUnwrappedIncrement()) << 16;
+    }
+  }
+
+  cbufferData.numPatchingAddrs = m_NumPatchingAddrs;
+
+  RayIndirectDispatchCB prepInfo = {};
+
+  prepInfo.commandSigDispatchOffset = comSig->sig.PackedByteSize - sizeof(D3D12_DISPATCH_RAYS_DESC);
+  prepInfo.commandSigSize = comSig->sig.PackedByteSize;
+  prepInfo.commandSigStride = comSig->sig.ByteStride;
+  prepInfo.maxCommandCount = MaxCommandCount;
+  prepInfo.scratchBuffer = scratchBuffer->Address();
+
+  // set up general patching data - lookup buffers and so on
+
+  unwrappedCmd->SetPipelineState(m_RayPatchingData.indirectPrepPipe);
+  unwrappedCmd->SetComputeRootSignature(m_RayPatchingData.indirectPrepRootSig);
+  unwrappedCmd->SetComputeRoot32BitConstants((UINT)D3D12IndirectPrepParam::GeneralCB,
+                                             sizeof(prepInfo) / sizeof(uint32_t), &prepInfo, 0);
+  unwrappedCmd->SetComputeRootShaderResourceView(
+      (UINT)D3D12IndirectPrepParam::AppExecuteArgs,
+      pArgumentBuffer->GetGPUVirtualAddress() + ArgumentBufferOffset);
+  unwrappedCmd->SetComputeRootShaderResourceView(
+      (UINT)D3D12IndirectPrepParam::AppCount,
+      pCountBuffer ? pCountBuffer->GetGPUVirtualAddress() + CountBufferOffset
+                   : pArgumentBuffer->GetGPUVirtualAddress());
+  unwrappedCmd->SetComputeRootUnorderedAccessView((UINT)D3D12IndirectPrepParam::PatchedExecuteArgs,
+                                                  argsBuffer->Address());
+  unwrappedCmd->SetComputeRootUnorderedAccessView((UINT)D3D12IndirectPrepParam::InternalExecuteArgs,
+                                                  argsBuffer->Address() + applicationArgsSize);
+  unwrappedCmd->SetComputeRootUnorderedAccessView(
+      (UINT)D3D12IndirectPrepParam::InternalExecuteCount,
+      argsBuffer->Address() + applicationArgsSize + patchingArgsSize);
+
+  // prepare our actual indirect patching by setting up destination space locations as well as
+  // patching the actual arguments buffer we'll return
+  unwrappedCmd->Dispatch(1, 1, 1);
+
+  // this is ready for the application to use (once we patch the things it refers to), and for us to use below
+  D3D12_RESOURCE_BARRIER barrier = {};
+  barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
+  barrier.Transition.pResource = argsBuffer->Resource();
+  barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
+  barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_INDIRECT_ARGUMENT;
+  unwrappedCmd->ResourceBarrier(1, &barrier);
+
+  unwrappedCmd->SetPipelineState(m_RayPatchingData.descPatchPipe);
+  unwrappedCmd->SetComputeRootSignature(m_RayPatchingData.descPatchRootSig);
+  unwrappedCmd->SetComputeRoot32BitConstants((UINT)D3D12PatchRayDispatchParam::GeneralCB,
+                                             sizeof(cbufferData) / sizeof(uint32_t), &cbufferData, 0);
+  unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::StateObjectData,
+                                                 m_LookupAddrs[0]);
+  unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::RecordData,
+                                                 m_LookupAddrs[1]);
+  unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::RootSigData,
+                                                 m_LookupAddrs[2]);
+  unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::AddrPatchData,
+                                                 m_LookupAddrs[3]);
+
+  RayDispatchShaderRecordCB recordInfo = {};
+  // these will be overwritten by the execute indirect, but set them to something to be safe
+  unwrappedCmd->SetComputeRoot32BitConstants((UINT)D3D12PatchRayDispatchParam::RecordCB,
+                                             sizeof(recordInfo) / sizeof(uint32_t), &recordInfo, 0);
+  unwrappedCmd->SetComputeRootShaderResourceView((UINT)D3D12PatchRayDispatchParam::SourceBuffer,
+                                                 m_LookupAddrs[0]);
+  unwrappedCmd->SetComputeRootUnorderedAccessView((UINT)D3D12PatchRayDispatchParam::DestBuffer,
+                                                  scratchBuffer->Address());
+  // execute the indirect arguments buffer that we prepared in the dispatch above to do the actual
+  // shader record patching/unwrapping
+  unwrappedCmd->ExecuteIndirect(m_RayPatchingData.indirectComSig, MaxCommandCount * 4,
+                                argsBuffer->Resource(), argsBuffer->Offset() + applicationArgsSize,
+                                argsBuffer->Resource(),
+                                argsBuffer->Offset() + applicationArgsSize + patchingArgsSize);
+
+  // scratch buffer has now been patched and is ready to use as well
+  barrier.Transition.pResource = scratchBuffer->Resource();
+  barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
+  barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE;
+  unwrappedCmd->ResourceBarrier(1, &barrier);
+
+  // we have our own ref, the patch data has its ref too that will be held while the list is
+  // submittable. Each submission will also get a ref to keep this referenced lookup buffer alive until then
+  m_LookupBuffer->AddRef();
+  ret.resources.lookupBuffer = m_LookupBuffer;
+
+  // the patch and arguments buffers are not owned by us, so the refcounting is the same as above
+  // but it takes the ref we had when we created it.
+  ret.resources.patchScratchBuffer = scratchBuffer;
+  ret.resources.argumentBuffer = argsBuffer;
+
   return ret;
 }
 
@@ -1103,6 +1266,14 @@ void D3D12RaytracingResourceAndUtilHandler::PrepareRayDispatchBuffer(
 
 void D3D12RaytracingResourceAndUtilHandler::InitRayDispatchPatchingResources()
 {
+  D3D12ShaderCache *shaderCache = m_wrappedDevice->GetShaderCache();
+
+  if(shaderCache == NULL)
+  {
+    RDCERR("Shadercache not available");
+    return;
+  }
+
   // need 5x 2-DWORD root buffers, the rest we can have for constants.
   // this could be made another buffer to track but it fits in push constants so we'll use them
   RDCCOMPILE_ASSERT(
@@ -1197,53 +1368,195 @@ void D3D12RaytracingResourceAndUtilHandler::InitRayDispatchPatchingResources()
 
   RDCASSERT(rootParameters.size() == uint32_t(D3D12PatchRayDispatchParam::Count));
 
-  D3D12ShaderCache *shaderCache = m_wrappedDevice->GetShaderCache();
+  ID3DBlob *rootSig = shaderCache->MakeRootSig(rootParameters, D3D12_ROOT_SIGNATURE_FLAG_NONE);
 
-  if(shaderCache != NULL)
+  if(rootSig)
   {
-    ID3DBlob *rootSig = shaderCache->MakeRootSig(rootParameters, D3D12_ROOT_SIGNATURE_FLAG_NONE);
+    HRESULT result = m_wrappedDevice->GetReal()->CreateRootSignature(
+        0, rootSig->GetBufferPointer(), rootSig->GetBufferSize(), __uuidof(ID3D12RootSignature),
+        (void **)&m_RayPatchingData.descPatchRootSig);
 
-    if(rootSig)
+    if(!SUCCEEDED(result))
+      RDCERR("Unable to create root signature for dispatch patching");
+
+    // PipelineState
+    ID3DBlob *shader = NULL;
+    rdcstr hlsl = GetEmbeddedResource(raytracing_hlsl);
+    shaderCache->GetShaderBlob(hlsl.c_str(), "RENDERDOC_PatchRayDispatchCS",
+                               D3DCOMPILE_WARNINGS_ARE_ERRORS, {}, "cs_5_0", &shader);
+
+    if(shader)
     {
-      HRESULT result = m_wrappedDevice->GetReal()->CreateRootSignature(
-          0, rootSig->GetBufferPointer(), rootSig->GetBufferSize(), __uuidof(ID3D12RootSignature),
-          (void **)&m_RayPatchingData.rootSig);
+      D3D12_COMPUTE_PIPELINE_STATE_DESC pipeline;
+      pipeline.Flags = D3D12_PIPELINE_STATE_FLAG_NONE;
+      pipeline.NodeMask = 0;
+      pipeline.CS = {(void *)shader->GetBufferPointer(), shader->GetBufferSize()};
+      pipeline.CachedPSO = {NULL, 0};
+      pipeline.pRootSignature = m_RayPatchingData.descPatchRootSig;
+
+      result = m_wrappedDevice->GetReal()->CreateComputePipelineState(
+          &pipeline, __uuidof(ID3D12PipelineState), (void **)&m_RayPatchingData.descPatchPipe);
 
       if(!SUCCEEDED(result))
-        RDCERR("Unable to create root signature for dispatch patching");
-
-      // PipelineState
-      ID3DBlob *shader = NULL;
-      rdcstr hlsl = GetEmbeddedResource(raytracing_hlsl);
-      shaderCache->GetShaderBlob(hlsl.c_str(), "RENDERDOC_PatchRayDispatchCS",
-                                 D3DCOMPILE_WARNINGS_ARE_ERRORS, {}, "cs_5_0", &shader);
-
-      if(shader)
-      {
-        D3D12_COMPUTE_PIPELINE_STATE_DESC pipeline;
-        pipeline.Flags = D3D12_PIPELINE_STATE_FLAG_NONE;
-        pipeline.NodeMask = 0;
-        pipeline.CS = {(void *)shader->GetBufferPointer(), shader->GetBufferSize()};
-        pipeline.CachedPSO = {NULL, 0};
-        pipeline.pRootSignature = m_RayPatchingData.rootSig;
-
-        result = m_wrappedDevice->GetReal()->CreateComputePipelineState(
-            &pipeline, __uuidof(ID3D12PipelineState), (void **)&m_RayPatchingData.pipe);
-
-        if(!SUCCEEDED(result))
-          RDCERR("Unable to create pipeline for dispatch patching");
-      }
-      else
-      {
-        RDCERR("Failed to get shader for dispatch patching");
-      }
-
-      SAFE_RELEASE(rootSig);
+        RDCERR("Unable to create pipeline for dispatch patching");
     }
+    else
+    {
+      RDCERR("Failed to get shader for dispatch patching");
+    }
+
+    SAFE_RELEASE(rootSig);
   }
-  else
+
+  // need 5x 2-DWORD root buffers, the rest we can have for constants.
+  // this could be made another buffer to track but it fits in push constants so we'll use them
+  RDCCOMPILE_ASSERT((sizeof(RayIndirectDispatchCB) / sizeof(uint32_t)) +
+                            (uint32_t(D3D12IndirectPrepParam::Count) - 2) * 2 <
+                        64,
+                    "Root signature constants are too large");
+
+  rootParameters.clear();
+  rootParameters.reserve((uint16_t)D3D12IndirectPrepParam::Count);
+
   {
-    RDCERR("Shadercache not available");
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Constants.ShaderRegister = 0;
+    rootParam.Constants.RegisterSpace = 0;
+    rootParam.Constants.Num32BitValues = sizeof(RayIndirectDispatchCB) / sizeof(uint32_t);
+    rootParameters.push_back(rootParam);
+  }
+
+  {
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_SRV;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Descriptor.ShaderRegister = 0;
+    rootParam.Descriptor.RegisterSpace = 0;
+    rootParam.Descriptor.Flags = D3D12_ROOT_DESCRIPTOR_FLAG_NONE;
+    rootParameters.push_back(rootParam);
+  }
+
+  {
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_SRV;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Descriptor.ShaderRegister = 1;
+    rootParam.Descriptor.RegisterSpace = 0;
+    rootParam.Descriptor.Flags = D3D12_ROOT_DESCRIPTOR_FLAG_NONE;
+    rootParameters.push_back(rootParam);
+  }
+
+  {
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_UAV;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Descriptor.ShaderRegister = 0;
+    rootParam.Descriptor.RegisterSpace = 0;
+    rootParam.Descriptor.Flags = D3D12_ROOT_DESCRIPTOR_FLAG_NONE;
+    rootParameters.push_back(rootParam);
+  }
+
+  {
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_UAV;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Descriptor.ShaderRegister = 1;
+    rootParam.Descriptor.RegisterSpace = 0;
+    rootParam.Descriptor.Flags = D3D12_ROOT_DESCRIPTOR_FLAG_NONE;
+    rootParameters.push_back(rootParam);
+  }
+
+  {
+    D3D12_ROOT_PARAMETER1 rootParam;
+    rootParam.ParameterType = D3D12_ROOT_PARAMETER_TYPE_UAV;
+    rootParam.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+    rootParam.Descriptor.ShaderRegister = 2;
+    rootParam.Descriptor.RegisterSpace = 0;
+    rootParam.Descriptor.Flags = D3D12_ROOT_DESCRIPTOR_FLAG_NONE;
+    rootParameters.push_back(rootParam);
+  }
+
+  RDCASSERT(rootParameters.size() == uint32_t(D3D12IndirectPrepParam::Count));
+
+  rootSig = shaderCache->MakeRootSig(rootParameters, D3D12_ROOT_SIGNATURE_FLAG_NONE);
+
+  if(rootSig)
+  {
+    HRESULT result = m_wrappedDevice->GetReal()->CreateRootSignature(
+        0, rootSig->GetBufferPointer(), rootSig->GetBufferSize(), __uuidof(ID3D12RootSignature),
+        (void **)&m_RayPatchingData.indirectPrepRootSig);
+
+    if(!SUCCEEDED(result))
+      RDCERR("Unable to create root signature for indirect execute patching");
+
+    // PipelineState
+    ID3DBlob *shader = NULL;
+    rdcstr hlsl = GetEmbeddedResource(raytracing_hlsl);
+    shaderCache->GetShaderBlob(hlsl.c_str(), "RENDERDOC_PrepareRayIndirectExecuteCS",
+                               D3DCOMPILE_WARNINGS_ARE_ERRORS, {}, "cs_5_0", &shader);
+
+    if(shader)
+    {
+      D3D12_COMPUTE_PIPELINE_STATE_DESC pipeline;
+      pipeline.Flags = D3D12_PIPELINE_STATE_FLAG_NONE;
+      pipeline.NodeMask = 0;
+      pipeline.CS = {(void *)shader->GetBufferPointer(), shader->GetBufferSize()};
+      pipeline.CachedPSO = {NULL, 0};
+      pipeline.pRootSignature = m_RayPatchingData.indirectPrepRootSig;
+
+      result = m_wrappedDevice->GetReal()->CreateComputePipelineState(
+          &pipeline, __uuidof(ID3D12PipelineState), (void **)&m_RayPatchingData.indirectPrepPipe);
+
+      if(!SUCCEEDED(result))
+        RDCERR("Unable to create pipeline for indirect execute patching");
+    }
+    else
+    {
+      RDCERR("Failed to get shader for indirect execute patching");
+    }
+
+    SAFE_RELEASE(rootSig);
+  }
+
+  {
+    D3D12_INDIRECT_ARGUMENT_DESC args[] = {
+        {D3D12_INDIRECT_ARGUMENT_TYPE_CONSTANT},
+        {D3D12_INDIRECT_ARGUMENT_TYPE_SHADER_RESOURCE_VIEW},
+        {D3D12_INDIRECT_ARGUMENT_TYPE_UNORDERED_ACCESS_VIEW},
+        {D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH},
+    };
+
+    args[0].Constant.DestOffsetIn32BitValues = 0;
+    args[0].Constant.Num32BitValuesToSet = sizeof(RayDispatchShaderRecordCB) / sizeof(uint32_t);
+    args[0].Constant.RootParameterIndex = (uint32_t)D3D12PatchRayDispatchParam::RecordCB;
+
+    RDCCOMPILE_ASSERT(sizeof(RayDispatchShaderRecordCB) == offsetof(PatchingExecute, sourceData),
+                      "Start of PatchingExecute is not one RayDispatchShaderRecordCB");
+
+    args[1].ShaderResourceView.RootParameterIndex =
+        (uint32_t)D3D12PatchRayDispatchParam::SourceBuffer;
+    args[2].UnorderedAccessView.RootParameterIndex = (uint32_t)D3D12PatchRayDispatchParam::DestBuffer;
+
+    RDCCOMPILE_ASSERT(sizeof(PatchingExecute) == sizeof(GPUAddress) * 2 +
+                                                     sizeof(RayDispatchShaderRecordCB) +
+                                                     sizeof(Vec4u) + sizeof(uint32_t) * 2,
+                      "PatchingExecute has changed size");
+    RDCCOMPILE_ASSERT(sizeof(PatchingExecute) % 16 == 0,
+                      "PatchingExecute is not 16-byte aligned, add explicit padding");
+
+    D3D12_COMMAND_SIGNATURE_DESC desc = {};
+    desc.ByteStride = (UINT)AlignUp16(sizeof(PatchingExecute));
+    desc.NumArgumentDescs = ARRAY_COUNT(args);
+    desc.pArgumentDescs = args;
+
+    HRESULT hr = m_wrappedDevice->GetReal()->CreateCommandSignature(
+        &desc, m_RayPatchingData.descPatchRootSig, __uuidof(ID3D12CommandSignature),
+        (void **)&m_RayPatchingData.indirectComSig);
+
+    if(!SUCCEEDED(hr))
+      RDCERR("Unable to create command signature for indirect execute patching");
   }
 }
 

renderdoc/driver/d3d12/d3d12_manager.h

@@ -997,6 +997,17 @@ enum class D3D12PatchTLASBuildParam
   Count
 };
 
+enum class D3D12IndirectPrepParam
+{
+  GeneralCB,
+  AppExecuteArgs,
+  AppCount,
+  PatchedExecuteArgs,
+  InternalExecuteArgs,
+  InternalExecuteCount,
+  Count,
+};
+
 enum class D3D12PatchRayDispatchParam
 {
   GeneralCB,
@@ -1025,6 +1036,8 @@ struct PatchedRayDispatch
     D3D12GpuBuffer *lookupBuffer;
     // the scratch buffer used for patching's fence.
     D3D12GpuBuffer *patchScratchBuffer;
+    // the argument buffer used for indirect executes.
+    D3D12GpuBuffer *argumentBuffer;
 
     // for convenience, when these resources are referenced in a queue they get a fence value to
     // indicate when they're safe to release. This values are unset when returned from patching or
@@ -1063,8 +1076,11 @@ public:
     SAFE_RELEASE(m_gpuFence);
     SAFE_RELEASE(m_accStructPatchInfo.m_rootSignature);
     SAFE_RELEASE(m_accStructPatchInfo.m_pipeline);
-    SAFE_RELEASE(m_RayPatchingData.rootSig);
-    SAFE_RELEASE(m_RayPatchingData.pipe);
+    SAFE_RELEASE(m_RayPatchingData.descPatchRootSig);
+    SAFE_RELEASE(m_RayPatchingData.descPatchPipe);
+    SAFE_RELEASE(m_RayPatchingData.indirectComSig);
+    SAFE_RELEASE(m_RayPatchingData.indirectPrepPipe);
+    SAFE_RELEASE(m_RayPatchingData.indirectPrepRootSig);
   }
 
   void InitInternalResources();
@@ -1079,6 +1095,12 @@ public:
   PatchedRayDispatch PatchRayDispatch(ID3D12GraphicsCommandList4 *unwrappedCmd,
                                       rdcarray<ResourceId> heaps,
                                       const D3D12_DISPATCH_RAYS_DESC &desc);
+  PatchedRayDispatch PatchIndirectRayDispatch(ID3D12GraphicsCommandList *unwrappedCmd,
+                                              rdcarray<ResourceId> heaps,
+                                              ID3D12CommandSignature *pCommandSignature,
+                                              UINT MaxCommandCount, ID3D12Resource *pArgumentBuffer,
+                                              UINT64 ArgumentBufferOffset,
+                                              ID3D12Resource *pCountBuffer, UINT64 CountBufferOffset);
 
   void AddPendingASBuilds(ID3D12Fence *fence, UINT64 waitValue,
                           const rdcarray<std::function<bool()>> &callbacks);
@@ -1127,8 +1149,11 @@ private:
   // pipeline data for patching ray dispatches
   struct
   {
-    ID3D12RootSignature *rootSig = NULL;
-    ID3D12PipelineState *pipe = NULL;
+    ID3D12RootSignature *descPatchRootSig = NULL;
+    ID3D12PipelineState *descPatchPipe = NULL;
+    ID3D12RootSignature *indirectPrepRootSig = NULL;
+    ID3D12PipelineState *indirectPrepPipe = NULL;
+    ID3D12CommandSignature *indirectComSig = NULL;
   } m_RayPatchingData;
 
   struct PendingASBuild

renderdoc/driver/d3d12/d3d12_resources.h

@@ -380,9 +380,78 @@ public:
     TypeEnum = Resource_CommandSignature,
   };
 
-  WrappedID3D12CommandSignature(ID3D12CommandSignature *real, WrappedID3D12Device *device)
+  WrappedID3D12CommandSignature(ID3D12CommandSignature *real, WrappedID3D12Device *device,
+                                const D3D12_COMMAND_SIGNATURE_DESC &Descriptor)
       : WrappedDeviceChild12(real, device)
   {
+    sig.ByteStride = Descriptor.ByteStride;
+    sig.arguments.assign(Descriptor.pArgumentDescs, Descriptor.NumArgumentDescs);
+
+    sig.graphics = true;
+    sig.PackedByteSize = 0;
+
+    // From MSDN, command signatures are either graphics or compute so just search for dispatches:
+    // "A given command signature is either an action or a compute command signature. If a command
+    // signature contains a drawing operation, then it is a graphics command signature. Otherwise,
+    // the command signature must contain a dispatch operation, and it is a compute command
+    // signature."
+    for(uint32_t i = 0; i < Descriptor.NumArgumentDescs; i++)
+    {
+      switch(Descriptor.pArgumentDescs[i].Type)
+      {
+        case D3D12_INDIRECT_ARGUMENT_TYPE_DRAW:
+        {
+          sig.PackedByteSize += sizeof(D3D12_DRAW_ARGUMENTS);
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_DRAW_INDEXED:
+        {
+          sig.PackedByteSize += sizeof(D3D12_DRAW_INDEXED_ARGUMENTS);
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH:
+        {
+          sig.PackedByteSize += sizeof(D3D12_DISPATCH_ARGUMENTS);
+          sig.graphics = false;
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH_MESH:
+        {
+          sig.PackedByteSize += sizeof(D3D12_DISPATCH_MESH_ARGUMENTS);
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_DISPATCH_RAYS:
+        {
+          sig.PackedByteSize += sizeof(D3D12_DISPATCH_RAYS_DESC);
+          sig.raytraced = true;
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_CONSTANT:
+        {
+          sig.PackedByteSize +=
+              sizeof(uint32_t) * Descriptor.pArgumentDescs[i].Constant.Num32BitValuesToSet;
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_VERTEX_BUFFER_VIEW:
+        {
+          sig.PackedByteSize += sizeof(D3D12_VERTEX_BUFFER_VIEW);
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_INDEX_BUFFER_VIEW:
+        {
+          sig.PackedByteSize += sizeof(D3D12_INDEX_BUFFER_VIEW);
+          break;
+        }
+        case D3D12_INDIRECT_ARGUMENT_TYPE_CONSTANT_BUFFER_VIEW:
+        case D3D12_INDIRECT_ARGUMENT_TYPE_SHADER_RESOURCE_VIEW:
+        case D3D12_INDIRECT_ARGUMENT_TYPE_UNORDERED_ACCESS_VIEW:
+        {
+          sig.PackedByteSize += sizeof(D3D12_GPU_VIRTUAL_ADDRESS);
+          break;
+        }
+        default: RDCERR("Unexpected argument type! %d", Descriptor.pArgumentDescs[i].Type); break;
+      }
+    }
   }
   virtual ~WrappedID3D12CommandSignature() { Shutdown(); }
 };

renderdoc/maths/vec.h

@@ -203,6 +203,24 @@ struct Vec4u
   };
 };
 
+struct Vec3u
+{
+  Vec3u(uint32_t X = 0, uint32_t Y = 0, uint32_t Z = 0)
+  {
+    x = X;
+    y = Y;
+    z = Z;
+  }
+  union
+  {
+    struct
+    {
+      uint32_t x, y, z;
+    };
+    uint32_t uv[3];
+  };
+};
+
 struct Vec4i
 {
   Vec4i(int32_t X = 0, int32_t Y = 0, int32_t Z = 0, int32_t W = 0)