Add support for nvapi wrapped PSO create functions

renderdoc/driver/d3d11/d3d11_device.h

@@ -162,6 +162,20 @@ public:
   virtual BOOL STDMETHODCALLTYPE SetReal(IUnknown *);
   virtual IUnknown *STDMETHODCALLTYPE GetReal();
   virtual BOOL STDMETHODCALLTYPE SetShaderExtUAV(DWORD space, DWORD reg, BOOL global);
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc) {}
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc) {}
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedGraphicsPipelineState(const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                      uint32_t space, ID3D12PipelineState *realPSO)
+  {
+    return NULL;
+  }
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedComputePipelineState(const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                     uint32_t space, ID3D12PipelineState *realPSO)
+  {
+    return NULL;
+  }
 };
 
 struct WrappedAGS11 : public IAGSD3DDevice

renderdoc/driver/d3d12/d3d12_device.cpp

@@ -409,6 +409,38 @@ BOOL STDMETHODCALLTYPE WrappedNVAPI12::SetShaderExtUAV(DWORD space, DWORD reg, B
   return TRUE;
 }
 
+void WrappedNVAPI12::UnwrapDesc(D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc)
+{
+  pDesc->pRootSignature = Unwrap(pDesc->pRootSignature);
+}
+
+void WrappedNVAPI12::UnwrapDesc(D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc)
+{
+  pDesc->pRootSignature = Unwrap(pDesc->pRootSignature);
+}
+
+ID3D12PipelineState *WrappedNVAPI12::ProcessCreatedGraphicsPipelineState(
+    const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, uint32_t reg, uint32_t space,
+    ID3D12PipelineState *realPSO)
+{
+  ID3D12PipelineState *ret = NULL;
+  m_pDevice.SetShaderExt(GPUVendor::nVidia);
+  m_pDevice.ProcessCreatedGraphicsPSO(realPSO, reg, space, pDesc, __uuidof(ID3D12PipelineState),
+                                      (void **)&ret);
+  return ret;
+}
+
+ID3D12PipelineState *WrappedNVAPI12::ProcessCreatedComputePipelineState(
+    const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, uint32_t reg, uint32_t space,
+    ID3D12PipelineState *realPSO)
+{
+  ID3D12PipelineState *ret = NULL;
+  m_pDevice.SetShaderExt(GPUVendor::nVidia);
+  m_pDevice.ProcessCreatedComputePSO(realPSO, reg, space, pDesc, __uuidof(ID3D12PipelineState),
+                                     (void **)&ret);
+  return ret;
+}
+
 HRESULT STDMETHODCALLTYPE WrappedAGS12::QueryInterface(REFIID riid, void **ppvObject)
 {
   return E_NOINTERFACE;
@@ -3060,6 +3092,12 @@ void WrappedID3D12Device::SetShaderExtUAV(GPUVendor vendor, uint32_t reg, uint32
 INSTANTIATE_FUNCTION_SERIALISED(void, WrappedID3D12Device, SetShaderExtUAV, GPUVendor vendor,
                                 uint32_t reg, uint32_t space, bool global);
 
+void WrappedID3D12Device::SetShaderExt(GPUVendor vendor)
+{
+  // just overwrite, we don't expect to switch back and forth on a given device.
+  m_VendorEXT = vendor;
+}
+
 void WrappedID3D12Device::GetShaderExtUAV(uint32_t &reg, uint32_t &space)
 {
   if(m_ThreadLocalEXTUAVSlot != ~0ULL)

renderdoc/driver/d3d12/d3d12_device.h

@@ -500,6 +500,16 @@ public:
   virtual BOOL STDMETHODCALLTYPE SetReal(IUnknown *);
   virtual IUnknown *STDMETHODCALLTYPE GetReal();
   virtual BOOL STDMETHODCALLTYPE SetShaderExtUAV(DWORD space, DWORD reg, BOOL global);
+
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc);
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc);
+
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedGraphicsPipelineState(const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                      uint32_t space, ID3D12PipelineState *realPSO);
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedComputePipelineState(const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                     uint32_t space, ID3D12PipelineState *realPSO);
 };
 
 struct WrappedAGS12 : public IAGSD3DDevice
@@ -1088,6 +1098,7 @@ public:
   IMPLEMENT_FUNCTION_SERIALISED(void, SetShaderExtUAV, GPUVendor vendor, uint32_t reg,
                                 uint32_t space, bool global);
   void GetShaderExtUAV(uint32_t &reg, uint32_t &space);
+  void SetShaderExt(GPUVendor vendor);
 
   // Protected session
   ID3D12Fence *CreateProtectedSessionFence(ID3D12Fence *real);
@@ -1125,6 +1136,17 @@ public:
                                        D3D12_COMMAND_LIST_TYPE type, REFIID riid,
                                        void **ppCommandAllocator);
 
+  // these are separated from CreateGraphicsPipelineState and CreateComputePipelineState so that
+  // extension creation functions can pass in their custom-created pipeline state
+  void ProcessCreatedGraphicsPSO(ID3D12PipelineState *real, uint32_t vendorExtReg,
+                                 uint32_t vendorExtSpace,
+                                 const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, REFIID riid,
+                                 void **ppPipelineState);
+  void ProcessCreatedComputePSO(ID3D12PipelineState *real, uint32_t vendorExtReg,
+                                uint32_t vendorExtSpace,
+                                const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, REFIID riid,
+                                void **ppPipelineState);
+
   IMPLEMENT_FUNCTION_THREAD_SERIALISED(virtual HRESULT STDMETHODCALLTYPE, CreateGraphicsPipelineState,
                                        const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, REFIID riid,
                                        void **ppPipelineState);

renderdoc/driver/d3d12/d3d12_device_wrap.cpp

@@ -516,6 +516,131 @@ bool WrappedID3D12Device::Serialise_CreateGraphicsPipelineState(
   return true;
 }
 
+void WrappedID3D12Device::ProcessCreatedGraphicsPSO(ID3D12PipelineState *real,
+                                                    uint32_t vendorExtReg, uint32_t vendorExtSpace,
+                                                    const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc,
+                                                    REFIID riid, void **ppPipelineState)
+{
+  for(const D3D12_SHADER_BYTECODE &sh : {pDesc->VS, pDesc->HS, pDesc->DS, pDesc->GS, pDesc->PS})
+  {
+    if(sh.BytecodeLength > 0 && sh.pShaderBytecode &&
+       DXBC::DXBCContainer::CheckForDXIL(sh.pShaderBytecode, sh.BytecodeLength))
+      m_UsedDXIL = true;
+  }
+
+  WrappedID3D12PipelineState *wrapped = new WrappedID3D12PipelineState(real, this);
+
+  if(IsCaptureMode(m_State))
+  {
+    CACHE_THREAD_SERIALISER();
+
+    Chunk *vendorChunk = NULL;
+    if(m_VendorEXT != GPUVendor::Unknown)
+    {
+      if(UsesExtensionUAV(pDesc->VS, vendorExtReg, vendorExtSpace) ||
+         UsesExtensionUAV(pDesc->HS, vendorExtReg, vendorExtSpace) ||
+         UsesExtensionUAV(pDesc->DS, vendorExtReg, vendorExtSpace) ||
+         UsesExtensionUAV(pDesc->GS, vendorExtReg, vendorExtSpace) ||
+         UsesExtensionUAV(pDesc->PS, vendorExtReg, vendorExtSpace))
+      {
+        // don't set initparams until we've seen at least one shader actually created using the
+        // extensions.
+        m_InitParams.VendorExtensions = m_VendorEXT;
+
+        // if this shader uses the UAV slot registered for vendor extensions, serialise that out
+        // too
+        SCOPED_SERIALISE_CHUNK(D3D12Chunk::SetShaderExtUAV);
+        Serialise_SetShaderExtUAV(ser, m_VendorEXT, vendorExtReg, vendorExtSpace, true);
+        vendorChunk = scope.Get();
+      }
+    }
+
+    SCOPED_SERIALISE_CHUNK(D3D12Chunk::Device_CreateGraphicsPipeline);
+    Serialise_CreateGraphicsPipelineState(ser, pDesc, riid, (void **)&wrapped);
+
+    D3D12ResourceRecord *record = GetResourceManager()->AddResourceRecord(wrapped->GetResourceID());
+    record->type = Resource_PipelineState;
+    record->Length = 0;
+    wrapped->SetResourceRecord(record);
+
+    if(pDesc->pRootSignature)
+      record->AddParent(GetRecord(pDesc->pRootSignature));
+
+    if(vendorChunk)
+      record->AddChunk(vendorChunk);
+    record->AddChunk(scope.Get());
+  }
+  else
+  {
+    GetResourceManager()->AddLiveResource(wrapped->GetResourceID(), wrapped);
+
+    wrapped->graphics = new D3D12_EXPANDED_PIPELINE_STATE_STREAM_DESC(*pDesc);
+
+    D3D12_SHADER_BYTECODE *shaders[] = {
+        &wrapped->graphics->VS, &wrapped->graphics->HS, &wrapped->graphics->DS,
+        &wrapped->graphics->GS, &wrapped->graphics->PS,
+    };
+
+    for(size_t i = 0; i < ARRAY_COUNT(shaders); i++)
+    {
+      if(shaders[i]->BytecodeLength == 0 || shaders[i]->pShaderBytecode == NULL)
+      {
+        shaders[i]->pShaderBytecode = NULL;
+        shaders[i]->BytecodeLength = 0;
+      }
+      else
+      {
+        WrappedID3D12Shader *sh = WrappedID3D12Shader::AddShader(*shaders[i], this);
+        sh->AddRef();
+        if(m_GlobalEXTUAV != ~0U)
+          sh->SetShaderExtSlot(m_GlobalEXTUAV, m_GlobalEXTUAVSpace);
+        shaders[i]->pShaderBytecode = sh;
+      }
+    }
+
+    if(wrapped->graphics->InputLayout.NumElements)
+    {
+      wrapped->graphics->InputLayout.pInputElementDescs =
+          new D3D12_INPUT_ELEMENT_DESC[wrapped->graphics->InputLayout.NumElements];
+      memcpy((void *)wrapped->graphics->InputLayout.pInputElementDescs,
+             pDesc->InputLayout.pInputElementDescs,
+             sizeof(D3D12_INPUT_ELEMENT_DESC) * wrapped->graphics->InputLayout.NumElements);
+    }
+    else
+    {
+      wrapped->graphics->InputLayout.pInputElementDescs = NULL;
+    }
+
+    if(wrapped->graphics->StreamOutput.NumEntries)
+    {
+      wrapped->graphics->StreamOutput.pSODeclaration =
+          new D3D12_SO_DECLARATION_ENTRY[wrapped->graphics->StreamOutput.NumEntries];
+      memcpy((void *)wrapped->graphics->StreamOutput.pSODeclaration,
+             pDesc->StreamOutput.pSODeclaration,
+             sizeof(D3D12_SO_DECLARATION_ENTRY) * wrapped->graphics->StreamOutput.NumEntries);
+    }
+    else
+    {
+      wrapped->graphics->StreamOutput.pSODeclaration = NULL;
+    }
+
+    if(wrapped->graphics->StreamOutput.NumStrides)
+    {
+      wrapped->graphics->StreamOutput.pBufferStrides =
+          new UINT[wrapped->graphics->StreamOutput.NumStrides];
+      memcpy((void *)wrapped->graphics->StreamOutput.pBufferStrides,
+             pDesc->StreamOutput.pBufferStrides,
+             sizeof(UINT) * wrapped->graphics->StreamOutput.NumStrides);
+    }
+    else
+    {
+      wrapped->graphics->StreamOutput.pBufferStrides = NULL;
+    }
+  }
+
+  *ppPipelineState = (ID3D12PipelineState *)wrapped;
+}
+
 HRESULT WrappedID3D12Device::CreateGraphicsPipelineState(const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc,
                                                          REFIID riid, void **ppPipelineState)
 {
@@ -535,126 +660,12 @@ HRESULT WrappedID3D12Device::CreateGraphicsPipelineState(const D3D12_GRAPHICS_PI
 
   if(SUCCEEDED(ret))
   {
-    for(const D3D12_SHADER_BYTECODE &sh :
-        {unwrappedDesc.VS, unwrappedDesc.HS, unwrappedDesc.DS, unwrappedDesc.GS, unwrappedDesc.PS})
-    {
-      if(sh.BytecodeLength > 0 && sh.pShaderBytecode &&
-         DXBC::DXBCContainer::CheckForDXIL(sh.pShaderBytecode, sh.BytecodeLength))
-        m_UsedDXIL = true;
-    }
+    // use implicit register/space
+    uint32_t reg = ~0U, space = ~0U;
+    if(m_VendorEXT != GPUVendor::Unknown)
+      GetShaderExtUAV(reg, space);
 
-    WrappedID3D12PipelineState *wrapped = new WrappedID3D12PipelineState(real, this);
-
-    if(IsCaptureMode(m_State))
-    {
-      CACHE_THREAD_SERIALISER();
-
-      Chunk *vendorChunk = NULL;
-      if(m_VendorEXT != GPUVendor::Unknown)
-      {
-        uint32_t reg = ~0U, space = ~0U;
-        GetShaderExtUAV(reg, space);
-
-        if(UsesExtensionUAV(pDesc->VS, reg, space) || UsesExtensionUAV(pDesc->HS, reg, space) ||
-           UsesExtensionUAV(pDesc->DS, reg, space) || UsesExtensionUAV(pDesc->GS, reg, space) ||
-           UsesExtensionUAV(pDesc->PS, reg, space))
-        {
-          // don't set initparams until we've seen at least one shader actually created using the
-          // extensions.
-          m_InitParams.VendorExtensions = m_VendorEXT;
-
-          // if this shader uses the UAV slot registered for vendor extensions, serialise that out
-          // too
-          SCOPED_SERIALISE_CHUNK(D3D12Chunk::SetShaderExtUAV);
-          Serialise_SetShaderExtUAV(ser, m_VendorEXT, reg, space, true);
-          vendorChunk = scope.Get();
-        }
-      }
-
-      SCOPED_SERIALISE_CHUNK(D3D12Chunk::Device_CreateGraphicsPipeline);
-      Serialise_CreateGraphicsPipelineState(ser, pDesc, riid, (void **)&wrapped);
-
-      D3D12ResourceRecord *record = GetResourceManager()->AddResourceRecord(wrapped->GetResourceID());
-      record->type = Resource_PipelineState;
-      record->Length = 0;
-      wrapped->SetResourceRecord(record);
-
-      if(pDesc->pRootSignature)
-        record->AddParent(GetRecord(pDesc->pRootSignature));
-
-      if(vendorChunk)
-        record->AddChunk(vendorChunk);
-      record->AddChunk(scope.Get());
-    }
-    else
-    {
-      GetResourceManager()->AddLiveResource(wrapped->GetResourceID(), wrapped);
-
-      wrapped->graphics = new D3D12_EXPANDED_PIPELINE_STATE_STREAM_DESC(*pDesc);
-
-      D3D12_SHADER_BYTECODE *shaders[] = {
-          &wrapped->graphics->VS, &wrapped->graphics->HS, &wrapped->graphics->DS,
-          &wrapped->graphics->GS, &wrapped->graphics->PS,
-      };
-
-      for(size_t i = 0; i < ARRAY_COUNT(shaders); i++)
-      {
-        if(shaders[i]->BytecodeLength == 0 || shaders[i]->pShaderBytecode == NULL)
-        {
-          shaders[i]->pShaderBytecode = NULL;
-          shaders[i]->BytecodeLength = 0;
-        }
-        else
-        {
-          WrappedID3D12Shader *sh = WrappedID3D12Shader::AddShader(*shaders[i], this);
-          sh->AddRef();
-          if(m_GlobalEXTUAV != ~0U)
-            sh->SetShaderExtSlot(m_GlobalEXTUAV, m_GlobalEXTUAVSpace);
-          shaders[i]->pShaderBytecode = sh;
-        }
-      }
-
-      if(wrapped->graphics->InputLayout.NumElements)
-      {
-        wrapped->graphics->InputLayout.pInputElementDescs =
-            new D3D12_INPUT_ELEMENT_DESC[wrapped->graphics->InputLayout.NumElements];
-        memcpy((void *)wrapped->graphics->InputLayout.pInputElementDescs,
-               pDesc->InputLayout.pInputElementDescs,
-               sizeof(D3D12_INPUT_ELEMENT_DESC) * wrapped->graphics->InputLayout.NumElements);
-      }
-      else
-      {
-        wrapped->graphics->InputLayout.pInputElementDescs = NULL;
-      }
-
-      if(wrapped->graphics->StreamOutput.NumEntries)
-      {
-        wrapped->graphics->StreamOutput.pSODeclaration =
-            new D3D12_SO_DECLARATION_ENTRY[wrapped->graphics->StreamOutput.NumEntries];
-        memcpy((void *)wrapped->graphics->StreamOutput.pSODeclaration,
-               pDesc->StreamOutput.pSODeclaration,
-               sizeof(D3D12_SO_DECLARATION_ENTRY) * wrapped->graphics->StreamOutput.NumEntries);
-      }
-      else
-      {
-        wrapped->graphics->StreamOutput.pSODeclaration = NULL;
-      }
-
-      if(wrapped->graphics->StreamOutput.NumStrides)
-      {
-        wrapped->graphics->StreamOutput.pBufferStrides =
-            new UINT[wrapped->graphics->StreamOutput.NumStrides];
-        memcpy((void *)wrapped->graphics->StreamOutput.pBufferStrides,
-               pDesc->StreamOutput.pBufferStrides,
-               sizeof(UINT) * wrapped->graphics->StreamOutput.NumStrides);
-      }
-      else
-      {
-        wrapped->graphics->StreamOutput.pBufferStrides = NULL;
-      }
-    }
-
-    *ppPipelineState = (ID3D12PipelineState *)wrapped;
+    ProcessCreatedGraphicsPSO(real, reg, space, pDesc, riid, ppPipelineState);
   }
 
   return ret;
@@ -723,6 +734,66 @@ bool WrappedID3D12Device::Serialise_CreateComputePipelineState(
   return true;
 }
 
+void WrappedID3D12Device::ProcessCreatedComputePSO(ID3D12PipelineState *real, uint32_t vendorExtReg,
+                                                   uint32_t vendorExtSpace,
+                                                   const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc,
+                                                   REFIID riid, void **ppPipelineState)
+{
+  if(DXBC::DXBCContainer::CheckForDXIL(pDesc->CS.pShaderBytecode, pDesc->CS.BytecodeLength))
+    m_UsedDXIL = true;
+
+  WrappedID3D12PipelineState *wrapped = new WrappedID3D12PipelineState(real, this);
+
+  if(IsCaptureMode(m_State))
+  {
+    CACHE_THREAD_SERIALISER();
+
+    Chunk *vendorChunk = NULL;
+    if(m_VendorEXT != GPUVendor::Unknown)
+    {
+      if(UsesExtensionUAV(pDesc->CS, vendorExtReg, vendorExtSpace))
+      {
+        // don't set initparams until we've seen at least one shader actually created using the
+        // extensions.
+        m_InitParams.VendorExtensions = m_VendorEXT;
+
+        // if this shader uses the UAV slot registered for vendor extensions, serialise that out
+        // too
+        SCOPED_SERIALISE_CHUNK(D3D12Chunk::SetShaderExtUAV);
+        Serialise_SetShaderExtUAV(ser, m_VendorEXT, vendorExtReg, vendorExtSpace, true);
+        vendorChunk = scope.Get();
+      }
+    }
+
+    SCOPED_SERIALISE_CHUNK(D3D12Chunk::Device_CreateComputePipeline);
+    Serialise_CreateComputePipelineState(ser, pDesc, riid, (void **)&wrapped);
+
+    D3D12ResourceRecord *record = GetResourceManager()->AddResourceRecord(wrapped->GetResourceID());
+    record->type = Resource_PipelineState;
+    record->Length = 0;
+    wrapped->SetResourceRecord(record);
+
+    if(pDesc->pRootSignature)
+      record->AddParent(GetRecord(pDesc->pRootSignature));
+
+    if(vendorChunk)
+      record->AddChunk(vendorChunk);
+    record->AddChunk(scope.Get());
+  }
+  else
+  {
+    GetResourceManager()->AddLiveResource(wrapped->GetResourceID(), wrapped);
+
+    wrapped->compute = new D3D12_EXPANDED_PIPELINE_STATE_STREAM_DESC(*pDesc);
+
+    WrappedID3D12Shader *sh = WrappedID3D12Shader::AddShader(wrapped->compute->CS, this);
+    sh->AddRef();
+    wrapped->compute->CS.pShaderBytecode = sh;
+  }
+
+  *ppPipelineState = (ID3D12PipelineState *)wrapped;
+}
+
 HRESULT WrappedID3D12Device::CreateComputePipelineState(const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc,
                                                         REFIID riid, void **ppPipelineState)
 {
@@ -742,63 +813,12 @@ HRESULT WrappedID3D12Device::CreateComputePipelineState(const D3D12_COMPUTE_PIPE
 
   if(SUCCEEDED(ret))
   {
-    if(DXBC::DXBCContainer::CheckForDXIL(unwrappedDesc.CS.pShaderBytecode,
-                                         unwrappedDesc.CS.BytecodeLength))
-      m_UsedDXIL = true;
+    // use implicit register/space
+    uint32_t reg = ~0U, space = ~0U;
+    if(m_VendorEXT != GPUVendor::Unknown)
+      GetShaderExtUAV(reg, space);
 
-    WrappedID3D12PipelineState *wrapped = new WrappedID3D12PipelineState(real, this);
-
-    if(IsCaptureMode(m_State))
-    {
-      CACHE_THREAD_SERIALISER();
-
-      Chunk *vendorChunk = NULL;
-      if(m_VendorEXT != GPUVendor::Unknown)
-      {
-        uint32_t reg = ~0U, space = ~0U;
-        GetShaderExtUAV(reg, space);
-
-        if(UsesExtensionUAV(pDesc->CS, reg, space))
-        {
-          // don't set initparams until we've seen at least one shader actually created using the
-          // extensions.
-          m_InitParams.VendorExtensions = m_VendorEXT;
-
-          // if this shader uses the UAV slot registered for vendor extensions, serialise that out
-          // too
-          SCOPED_SERIALISE_CHUNK(D3D12Chunk::SetShaderExtUAV);
-          Serialise_SetShaderExtUAV(ser, m_VendorEXT, reg, space, true);
-          vendorChunk = scope.Get();
-        }
-      }
-
-      SCOPED_SERIALISE_CHUNK(D3D12Chunk::Device_CreateComputePipeline);
-      Serialise_CreateComputePipelineState(ser, pDesc, riid, (void **)&wrapped);
-
-      D3D12ResourceRecord *record = GetResourceManager()->AddResourceRecord(wrapped->GetResourceID());
-      record->type = Resource_PipelineState;
-      record->Length = 0;
-      wrapped->SetResourceRecord(record);
-
-      if(pDesc->pRootSignature)
-        record->AddParent(GetRecord(pDesc->pRootSignature));
-
-      if(vendorChunk)
-        record->AddChunk(vendorChunk);
-      record->AddChunk(scope.Get());
-    }
-    else
-    {
-      GetResourceManager()->AddLiveResource(wrapped->GetResourceID(), wrapped);
-
-      wrapped->compute = new D3D12_EXPANDED_PIPELINE_STATE_STREAM_DESC(*pDesc);
-
-      WrappedID3D12Shader *sh = WrappedID3D12Shader::AddShader(wrapped->compute->CS, this);
-      sh->AddRef();
-      wrapped->compute->CS.pShaderBytecode = sh;
-    }
-
-    *ppPipelineState = (ID3D12PipelineState *)wrapped;
+    ProcessCreatedComputePSO(real, reg, space, pDesc, riid, ppPipelineState);
   }
 
   return ret;

renderdoc/driver/ihv/nv/nvapi_hooks.cpp

@@ -152,6 +152,8 @@ private:
   HOOK_NVAPI(NvAPI_D3D12_IsNvShaderExtnOpCodeSupported, 0x3dfacec8);       \
   HOOK_NVAPI(NvAPI_D3D12_SetNvShaderExtnSlotSpace, 0xac2dfeb5);            \
   HOOK_NVAPI(NvAPI_D3D12_SetNvShaderExtnSlotSpaceLocalThread, 0x43d867c0); \
+  HOOK_NVAPI(NvAPI_D3D12_CreateGraphicsPipelineState, 0x2fc28856);         \
+  HOOK_NVAPI(NvAPI_D3D12_CreateComputePipelineState, 0x2762deac);          \
   WHITELIST_NVAPI(NvAPI_Unload, 0xd22bdd7e);                               \
   WHITELIST_NVAPI(NvAPI_GetErrorMessage, 0x6c2d048c);                      \
   WHITELIST_NVAPI(NvAPI_GetInterfaceVersionString, 0x01053fa5);
@@ -173,6 +175,7 @@ private:
     }
     else
     {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
       return NVAPI_INVALID_POINTER;
     }
   }
@@ -201,9 +204,17 @@ private:
 
         return ret;
       }
+      else
+      {
+        RDCERR("Couldn't retrieve ID3D12Device from RenderDoc-wrapped device");
+        return NVAPI_INVALID_POINTER;
+      }
+    }
+    else
+    {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
+      return NVAPI_INVALID_POINTER;
     }
-
-    return NVAPI_INVALID_POINTER;
   }
 
   static NvAPI_Status __cdecl NvAPI_D3D11_SetNvShaderExtnSlot_hook(__in IUnknown *pDev,
@@ -223,6 +234,7 @@ private:
     }
     else
     {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
       return NVAPI_INVALID_POINTER;
     }
   }
@@ -245,6 +257,7 @@ private:
     }
     else
     {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
       return NVAPI_INVALID_POINTER;
     }
   }
@@ -268,6 +281,7 @@ private:
     }
     else
     {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
       return NVAPI_INVALID_POINTER;
     }
   }
@@ -290,6 +304,175 @@ private:
     }
     else
     {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
+      return NVAPI_INVALID_POINTER;
+    }
+  }
+
+  static NvAPI_Status __cdecl NvAPI_D3D12_CreateGraphicsPipelineState_hook(
+      __in ID3D12Device *pDevice, __in const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pPSODesc,
+      NvU32 numExtensions, __in const NVAPI_D3D12_PSO_EXTENSION_DESC **ppExtensions,
+      __out ID3D12PipelineState **ppPSO)
+  {
+    // check that there's only supported extensions first, and extract the info we want.
+    uint32_t reg = 0, space = 0;
+    for(NvU32 i = 0; i < numExtensions; i++)
+    {
+      if(ppExtensions[i]->psoExtension != NV_PSO_SET_SHADER_EXTNENSION_SLOT_AND_SPACE)
+      {
+        RDCWARN("Unsupported D3D12 PSO extension: %d", ppExtensions[i]->psoExtension);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      // the versions don't look to be backwards compatible so we have to require an exact version
+      if(ppExtensions[i]->baseVersion != NV_PSO_EXTENSION_DESC_VER)
+      {
+        RDCERR("Unsupported PSO extension version %x, expected %x", ppExtensions[i]->baseVersion,
+               NV_PSO_EXTENSION_DESC_VER);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      const NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC_V1 *psoExt =
+          (const NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC_V1 *)ppExtensions[i];
+
+      if(psoExt->version != NV_SET_SHADER_EXTENSION_SLOT_DESC_VER)
+      {
+        RDCERR("Unsupported set-slot extension version %x, expected %x", psoExt->version,
+               NV_SET_SHADER_EXTENSION_SLOT_DESC_VER);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      reg = psoExt->uavSlot;
+      space = psoExt->registerSpace;
+    }
+
+    INVAPID3DDevice *nvapiDev = NULL;
+    HRESULT hr = pDevice->QueryInterface(__uuidof(INVAPID3DDevice), (void **)&nvapiDev);
+    // this will only succeed if it's our own wrapped device. It doesn't change the refcount, this
+    // is a COM-breaking backdoor
+    if(SUCCEEDED(hr))
+    {
+      IUnknown *real = nvapiDev->GetReal();
+
+      ID3D12Device *dev = NULL;
+      hr = real->QueryInterface(__uuidof(ID3D12Device), (void **)&dev);
+
+      if(SUCCEEDED(hr))
+      {
+        D3D12_GRAPHICS_PIPELINE_STATE_DESC desc = *pPSODesc;
+        nvapiDev->UnwrapDesc(&desc);
+
+        ID3D12PipelineState *realPSO = NULL;
+        NvAPI_Status ret = nvhooks.NvAPI_D3D12_CreateGraphicsPipelineState()(
+            dev, &desc, numExtensions, ppExtensions, &realPSO);
+
+        dev->Release();
+
+        if(ret == NVAPI_OK)
+        {
+          *ppPSO = nvapiDev->ProcessCreatedGraphicsPipelineState(pPSODesc, reg, space, realPSO);
+          return NVAPI_OK;
+        }
+        else
+        {
+          SAFE_RELEASE(realPSO);
+        }
+
+        return ret;
+      }
+      else
+      {
+        RDCERR("Couldn't retrieve ID3D12Device from RenderDoc-wrapped device");
+        return NVAPI_INVALID_POINTER;
+      }
+    }
+    else
+    {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
+      return NVAPI_INVALID_POINTER;
+    }
+  }
+
+  static NvAPI_Status __cdecl NvAPI_D3D12_CreateComputePipelineState_hook(
+      __in ID3D12Device *pDevice, __in const D3D12_COMPUTE_PIPELINE_STATE_DESC *pPSODesc,
+      NvU32 numExtensions, __in const NVAPI_D3D12_PSO_EXTENSION_DESC **ppExtensions,
+      __out ID3D12PipelineState **ppPSO)
+  {
+    // check that there's only supported extensions first, and extract the info we want.
+    uint32_t reg = 0, space = 0;
+    for(NvU32 i = 0; i < numExtensions; i++)
+    {
+      if(ppExtensions[i]->psoExtension != NV_PSO_SET_SHADER_EXTNENSION_SLOT_AND_SPACE)
+      {
+        RDCWARN("Unsupported D3D12 PSO extension: %d", ppExtensions[i]->psoExtension);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      // the versions don't look to be backwards compatible so we have to require an exact version
+      if(ppExtensions[i]->baseVersion != NV_PSO_EXTENSION_DESC_VER)
+      {
+        RDCERR("Unsupported PSO extension version %x, expected %x", ppExtensions[i]->baseVersion,
+               NV_PSO_EXTENSION_DESC_VER);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      const NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC_V1 *psoExt =
+          (const NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC_V1 *)ppExtensions[i];
+
+      if(psoExt->version != NV_SET_SHADER_EXTENSION_SLOT_DESC_VER)
+      {
+        RDCERR("Unsupported set-slot extension version %x, expected %x", psoExt->version,
+               NV_SET_SHADER_EXTENSION_SLOT_DESC_VER);
+        return NVAPI_NOT_SUPPORTED;
+      }
+
+      reg = psoExt->uavSlot;
+      space = psoExt->registerSpace;
+    }
+
+    INVAPID3DDevice *nvapiDev = NULL;
+    HRESULT hr = pDevice->QueryInterface(__uuidof(INVAPID3DDevice), (void **)&nvapiDev);
+    // this will only succeed if it's our own wrapped device. It doesn't change the refcount, this
+    // is a COM-breaking backdoor
+    if(SUCCEEDED(hr))
+    {
+      IUnknown *real = nvapiDev->GetReal();
+
+      ID3D12Device *dev = NULL;
+      hr = real->QueryInterface(__uuidof(ID3D12Device), (void **)&dev);
+
+      if(SUCCEEDED(hr))
+      {
+        D3D12_COMPUTE_PIPELINE_STATE_DESC desc = *pPSODesc;
+        nvapiDev->UnwrapDesc(&desc);
+
+        ID3D12PipelineState *realPSO = NULL;
+        NvAPI_Status ret = nvhooks.NvAPI_D3D12_CreateComputePipelineState()(
+            dev, &desc, numExtensions, ppExtensions, &realPSO);
+
+        dev->Release();
+
+        if(ret == NVAPI_OK)
+        {
+          *ppPSO = nvapiDev->ProcessCreatedComputePipelineState(pPSODesc, reg, space, realPSO);
+          return NVAPI_OK;
+        }
+        else
+        {
+          SAFE_RELEASE(realPSO);
+        }
+
+        return ret;
+      }
+      else
+      {
+        RDCERR("Couldn't retrieve ID3D12Device from RenderDoc-wrapped device");
+        return NVAPI_INVALID_POINTER;
+      }
+    }
+    else
+    {
+      RDCERR("Didn't pass RenderDoc-wrapped device to nvapi function");
       return NVAPI_INVALID_POINTER;
     }
   }

renderdoc/driver/ihv/nv/nvapi_wrapper.cpp

@@ -116,6 +116,22 @@ public:
     return FALSE;
   }
 
+  // only used on capture
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc) {}
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc) {}
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedGraphicsPipelineState(const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                      uint32_t space, ID3D12PipelineState *realPSO)
+  {
+    return NULL;
+  }
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE
+  ProcessCreatedComputePipelineState(const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, uint32_t reg,
+                                     uint32_t space, ID3D12PipelineState *realPSO)
+  {
+    return NULL;
+  }
+
 private:
   friend INVAPID3DDevice *InitialiseNVAPIReplay();
 

renderdoc/driver/ihv/nv/nvapi_wrapper.h

@@ -72,12 +72,26 @@ enum class NvShaderAtomic
   CompareAndSwap = 9,
 };
 
+struct D3D12_GRAPHICS_PIPELINE_STATE_DESC;
+struct D3D12_COMPUTE_PIPELINE_STATE_DESC;
+interface ID3D12PipelineState;
+
 MIDL_INTERFACE("DA122FC2-0F60-4904-AEA4-5ED1D2E1D19F")
 INVAPID3DDevice : public IUnknown
 {
   virtual BOOL STDMETHODCALLTYPE SetReal(IUnknown * device) = 0;
   virtual IUnknown *STDMETHODCALLTYPE GetReal() = 0;
   virtual BOOL STDMETHODCALLTYPE SetShaderExtUAV(DWORD space, DWORD reg, BOOL global) = 0;
+
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_GRAPHICS_PIPELINE_STATE_DESC * pDesc) = 0;
+  virtual void STDMETHODCALLTYPE UnwrapDesc(D3D12_COMPUTE_PIPELINE_STATE_DESC * pDesc) = 0;
+
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE ProcessCreatedGraphicsPipelineState(
+      const D3D12_GRAPHICS_PIPELINE_STATE_DESC *pDesc, uint32_t reg, uint32_t space,
+      ID3D12PipelineState *realPSO) = 0;
+  virtual ID3D12PipelineState *STDMETHODCALLTYPE ProcessCreatedComputePipelineState(
+      const D3D12_COMPUTE_PIPELINE_STATE_DESC *pDesc, uint32_t reg, uint32_t space,
+      ID3D12PipelineState *realPSO) = 0;
 };
 
 INVAPID3DDevice *InitialiseNVAPIReplay();