From 0a1844ef16c97c858627bf8da37409af022268c2 Mon Sep 17 00:00:00 2001
From: Dustin Graves <quic_dgraves@quicinc.com>
Date: Tue, 22 Oct 2024 16:48:29 -0600
Subject: [PATCH] Clear CachedPSO data for CreatePipelineState replay

When the replay --use-cached-psos option is not specified, clear the
cached PSO data from the ID3D12Device2::CreatePipelineState parameters,
similar to what is currently done for the ID3D12Device
CreateComputePipelineState and CreateGraphicsPipelineState methods.
---
 .../decode/dx12_replay_consumer_base.cpp      | 33 +++++++++++++++++++
 framework/decode/dx12_replay_consumer_base.h  |  6 ++++
 .../dx12_generators/replay_overrides.json     |  3 ++
 .../generated_dx12_replay_consumer.cpp        | 14 ++++----
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/framework/decode/dx12_replay_consumer_base.cpp b/framework/decode/dx12_replay_consumer_base.cpp
index e2ef484f5b..025bfbbd49 100644
--- a/framework/decode/dx12_replay_consumer_base.cpp
+++ b/framework/decode/dx12_replay_consumer_base.cpp
@@ -3580,6 +3580,39 @@ HRESULT Dx12ReplayConsumerBase::OverrideCreateComputePipelineState(
     return replay_result;
 }
 
+HRESULT Dx12ReplayConsumerBase::OverrideCreatePipelineState(
+    DxObjectInfo*                                                   device_object_info,
+    HRESULT                                                         original_result,
+    StructPointerDecoder<Decoded_D3D12_PIPELINE_STATE_STREAM_DESC>* pDesc,
+    Decoded_GUID                                                    riid,
+    HandlePointerDecoder<void*>*                                    ppPipelineState)
+{
+    GFXRECON_UNREFERENCED_PARAMETER(original_result);
+
+    GFXRECON_ASSERT(device_object_info != nullptr);
+    GFXRECON_ASSERT(device_object_info->object != nullptr);
+
+    auto device = static_cast<ID3D12Device2*>(device_object_info->object);
+
+    auto pDesc2 = pDesc->GetPointer();
+    if (!options_.use_cached_psos)
+    {
+        auto desc = pDesc->GetMetaStructPointer();
+        GFXRECON_ASSERT(desc != nullptr);
+
+        if (desc->cached_pso.decoded_value != nullptr)
+        {
+            desc->cached_pso.decoded_value->pCachedBlob           = nullptr;
+            desc->cached_pso.decoded_value->CachedBlobSizeInBytes = 0;
+        }
+    }
+
+    HRESULT replay_result =
+        device->CreatePipelineState(pDesc2, *riid.decoded_value, ppPipelineState->GetHandlePointer());
+
+    return replay_result;
+}
+
 HRESULT
 Dx12ReplayConsumerBase::OverrideSetFullscreenState(DxObjectInfo* swapchain_info,
                                                    HRESULT       original_result,
diff --git a/framework/decode/dx12_replay_consumer_base.h b/framework/decode/dx12_replay_consumer_base.h
index 17c25a17cf..6029cf702e 100644
--- a/framework/decode/dx12_replay_consumer_base.h
+++ b/framework/decode/dx12_replay_consumer_base.h
@@ -741,6 +741,12 @@ class Dx12ReplayConsumerBase : public Dx12Consumer
                                                Decoded_GUID                                                     riid,
                                                HandlePointerDecoder<void*>* pipelineState);
 
+    HRESULT OverrideCreatePipelineState(DxObjectInfo* device_object_info,
+                                        HRESULT       original_result,
+                                        StructPointerDecoder<Decoded_D3D12_PIPELINE_STATE_STREAM_DESC>* pDesc,
+                                        Decoded_GUID                                                    riid,
+                                        HandlePointerDecoder<void*>* ppPipelineState);
+
     HRESULT OverrideSetFullscreenState(DxObjectInfo* swapchain_info,
                                        HRESULT       original_result,
                                        BOOL          Fullscreen,
diff --git a/framework/generated/dx12_generators/replay_overrides.json b/framework/generated/dx12_generators/replay_overrides.json
index 23c9b2c799..24659c4eab 100644
--- a/framework/generated/dx12_generators/replay_overrides.json
+++ b/framework/generated/dx12_generators/replay_overrides.json
@@ -39,6 +39,9 @@
     "ID3D12Device1": {
       "CreatePipelineLibrary": "OverrideCreatePipelineLibrary"
     },
+    "ID3D12Device2": {
+      "CreatePipelineState": "OverrideCreatePipelineState"
+    },
     "ID3D12Device3": {
       "EnqueueMakeResident": "OverrideEnqueueMakeResident",
       "OpenExistingHeapFromAddress": "OverrideOpenExistingHeapFromAddress"
diff --git a/framework/generated/generated_dx12_replay_consumer.cpp b/framework/generated/generated_dx12_replay_consumer.cpp
index 60d1f889d2..c4500a7671 100644
--- a/framework/generated/generated_dx12_replay_consumer.cpp
+++ b/framework/generated/generated_dx12_replay_consumer.cpp
@@ -10214,14 +10214,16 @@ void Dx12ReplayConsumer::Process_ID3D12Device2_CreatePipelineState(
             ppPipelineState);
         MapStructObjects(pDesc->GetMetaStructPointer(), GetObjectInfoTable(), GetGpuVaTable());
         if(!ppPipelineState->IsNull()) ppPipelineState->SetHandleLength(1);
-        auto out_p_ppPipelineState    = ppPipelineState->GetPointer();
-        auto out_hp_ppPipelineState   = ppPipelineState->GetHandlePointer();
-        auto replay_result = reinterpret_cast<ID3D12Device2*>(replay_object->object)->CreatePipelineState(pDesc->GetPointer(),
-                                                                                                          *riid.decoded_value,
-                                                                                                          out_hp_ppPipelineState);
+        DxObjectInfo object_info_ppPipelineState{};
+        ppPipelineState->SetConsumerData(0, &object_info_ppPipelineState);
+        auto replay_result = OverrideCreatePipelineState(replay_object,
+                                                         return_value,
+                                                         pDesc,
+                                                         riid,
+                                                         ppPipelineState);
         if (SUCCEEDED(replay_result))
         {
-            AddObject(out_p_ppPipelineState, out_hp_ppPipelineState, format::ApiCall_ID3D12Device2_CreatePipelineState);
+            AddObject(ppPipelineState->GetPointer(), ppPipelineState->GetHandlePointer(), std::move(object_info_ppPipelineState), format::ApiCall_ID3D12Device2_CreatePipelineState);
         }
         CheckReplayResult("ID3D12Device2_CreatePipelineState", return_value, replay_result);
         CustomReplayPostCall<format::ApiCallId::ApiCall_ID3D12Device2_CreatePipelineState>::Dispatch(