Add late command list hooks

Should allow to restore root signatures in annoying games

OptiScaler/hooks/D3D12_Hooks.cpp

@@ -63,6 +63,10 @@ using PFN_SetGraphicsRootSignature =
 static PFN_SetComputeRootSignature o_SetComputeRootSignature = nullptr;
 static PFN_SetGraphicsRootSignature o_SetGraphicsRootSignature = nullptr;
 
+static std::atomic_bool hookedLate = false;
+static PFN_SetComputeRootSignature o_SetComputeRootSignatureLate = nullptr;
+static PFN_SetGraphicsRootSignature o_SetGraphicsRootSignatureLate = nullptr;
+
 static ankerl::unordered_dense::map<ID3D12GraphicsCommandList*, ID3D12RootSignature*> computeSignatures;
 static ankerl::unordered_dense::map<ID3D12GraphicsCommandList*, ID3D12RootSignature*> graphicSignatures;
 static bool isUpscalerActive = false;
@@ -240,7 +244,7 @@ VALIDATE_HOOK(hkSetComputeRootSignature, PFN_SetComputeRootSignature)
 static void hkSetComputeRootSignature(ID3D12GraphicsCommandList* commandList, ID3D12RootSignature* pRootSignature)
 {
     if (Config::Instance()->RestoreComputeSignature.value_or_default() && !isUpscalerActive && commandList != nullptr &&
-        pRootSignature != nullptr)
+        pRootSignature != nullptr && !hookedLate)
     {
         std::unique_lock<std::shared_mutex> lock(computeSigatureMutex);
         computeSignatures.insert_or_assign(commandList, pRootSignature);
@@ -253,7 +257,7 @@ VALIDATE_HOOK(hkSetGraphicsRootSignature, PFN_SetGraphicsRootSignature)
 static void hkSetGraphicsRootSignature(ID3D12GraphicsCommandList* commandList, ID3D12RootSignature* pRootSignature)
 {
     if (Config::Instance()->RestoreGraphicSignature.value_or_default() && !isUpscalerActive && commandList != nullptr &&
-        pRootSignature != nullptr)
+        pRootSignature != nullptr && !hookedLate)
     {
         std::unique_lock<std::shared_mutex> lock(graphSigatureMutex);
         graphicSignatures.insert_or_assign(commandList, pRootSignature);
@@ -262,6 +266,61 @@ static void hkSetGraphicsRootSignature(ID3D12GraphicsCommandList* commandList, I
     o_SetGraphicsRootSignature(commandList, pRootSignature);
 }
 
+VALIDATE_HOOK(hkSetComputeRootSignatureLate, PFN_SetComputeRootSignature)
+static void hkSetComputeRootSignatureLate(ID3D12GraphicsCommandList* commandList, ID3D12RootSignature* pRootSignature)
+{
+    if (Config::Instance()->RestoreComputeSignature.value_or_default() && !isUpscalerActive && commandList != nullptr &&
+        pRootSignature != nullptr)
+    {
+        std::unique_lock<std::shared_mutex> lock(computeSigatureMutex);
+        computeSignatures.insert_or_assign(commandList, pRootSignature);
+    }
+
+    o_SetComputeRootSignatureLate(commandList, pRootSignature);
+}
+
+VALIDATE_HOOK(hkSetGraphicsRootSignatureLate, PFN_SetGraphicsRootSignature)
+static void hkSetGraphicsRootSignatureLate(ID3D12GraphicsCommandList* commandList, ID3D12RootSignature* pRootSignature)
+{
+    if (Config::Instance()->RestoreGraphicSignature.value_or_default() && !isUpscalerActive && commandList != nullptr &&
+        pRootSignature != nullptr)
+    {
+        std::unique_lock<std::shared_mutex> lock(graphSigatureMutex);
+        graphicSignatures.insert_or_assign(commandList, pRootSignature);
+    }
+
+    o_SetGraphicsRootSignatureLate(commandList, pRootSignature);
+}
+
+void D3D12Hooks::HookToCommandListLate(ID3D12GraphicsCommandList* commandList)
+{
+    if (o_SetComputeRootSignatureLate || o_SetGraphicsRootSignatureLate)
+        return;
+
+    // Get the vtable pointer
+    PVOID* pVTable = *(PVOID**) commandList;
+
+    o_SetComputeRootSignatureLate = (PFN_SetComputeRootSignature) pVTable[29];
+    o_SetGraphicsRootSignatureLate = (PFN_SetGraphicsRootSignature) pVTable[30];
+
+    if (o_SetComputeRootSignatureLate != nullptr || o_SetGraphicsRootSignatureLate != nullptr)
+    {
+        DetourTransactionBegin();
+        DetourUpdateThread(GetCurrentThread());
+
+        if (o_SetComputeRootSignatureLate != nullptr)
+            DetourAttach(&(PVOID&) o_SetComputeRootSignatureLate, hkSetComputeRootSignatureLate);
+
+        if (o_SetGraphicsRootSignatureLate != nullptr)
+            DetourAttach(&(PVOID&) o_SetGraphicsRootSignatureLate, hkSetGraphicsRootSignatureLate);
+
+        LOG_DEBUG("Hooked SetRootSignature functions Late");
+        hookedLate = true;
+
+        DetourTransactionCommit();
+    }
+}
+
 static void HookToCommandList(ID3D12Device* InDevice)
 {
     if (o_SetComputeRootSignature != nullptr || o_SetGraphicsRootSignature != nullptr)
@@ -1238,6 +1297,16 @@ void D3D12Hooks::Unhook()
 
 void D3D12Hooks::SetRootSignatureTracking(bool enable) { isUpscalerActive = !enable; }
 
+bool D3D12Hooks::CanRestoreComputeRootSignature(ID3D12GraphicsCommandList* cmdList)
+{
+    return computeSignatures.contains(cmdList);
+}
+
+bool D3D12Hooks::CanRestoreGraphicsRootSignature(ID3D12GraphicsCommandList* cmdList)
+{
+    return graphicSignatures.contains(cmdList);
+}
+
 void D3D12Hooks::RestoreComputeRootSignature(ID3D12GraphicsCommandList* cmdList)
 {
     if (Config::Instance()->RestoreComputeSignature.value_or_default() && computeSignatures.contains(cmdList))

OptiScaler/hooks/D3D12_Hooks.h

@@ -14,6 +14,9 @@ class D3D12Hooks
     static void HookDevice(ID3D12Device* device);
     static void Unhook();
     static void SetRootSignatureTracking(bool enable);
+    static bool CanRestoreComputeRootSignature(ID3D12GraphicsCommandList* cmdList);
+    static bool CanRestoreGraphicsRootSignature(ID3D12GraphicsCommandList* cmdList);
+    static void HookToCommandListLate(ID3D12GraphicsCommandList* commandList);
     static void RestoreComputeRootSignature(ID3D12GraphicsCommandList* cmdList);
     static void RestoreGraphicsRootSignature(ID3D12GraphicsCommandList* cmdList);
 };

OptiScaler/inputs/NVNGX_DLSS_Dx12.cpp

@@ -498,6 +498,7 @@ NVSDK_NGX_API NVSDK_NGX_Result NVSDK_NGX_D3D12_CreateFeature(ID3D12GraphicsComma
         Config::Instance()->RestoreGraphicSignature.value_or_default())
     {
         D3D12Hooks::SetRootSignatureTracking(false);
+        D3D12Hooks::HookToCommandListLate(InCmdList);
     }
 
     if (InFeatureID == NVSDK_NGX_Feature_SuperSampling)
@@ -768,6 +769,20 @@ NVSDK_NGX_API NVSDK_NGX_Result NVSDK_NGX_D3D12_EvaluateFeature(ID3D12GraphicsCom
     if (Config::Instance()->SkipFirstFrames.has_value() && evalCounter < Config::Instance()->SkipFirstFrames.value())
         return NVSDK_NGX_Result_Success;
 
+    if (Config::Instance()->RestoreComputeSignature.value_or_default() &&
+        !D3D12Hooks::CanRestoreComputeRootSignature(InCmdList))
+    {
+        LOG_DEBUG("Skipping upscaling because can't restore compute signature");
+        return NVSDK_NGX_Result_Success;
+    }
+
+    if (Config::Instance()->RestoreGraphicSignature.value_or_default() &&
+        !D3D12Hooks::CanRestoreGraphicsRootSignature(InCmdList))
+    {
+        LOG_DEBUG("Skipping upscaling because can't restore graphics signature");
+        return NVSDK_NGX_Result_Success;
+    }
+
     if (InCallback)
         LOG_INFO("callback exist");
 

OptiScaler/upscalers/fsr2/FSR2Feature_Dx12.cpp

@@ -325,7 +325,7 @@ bool FSR2FeatureDx12::Evaluate(ID3D12GraphicsCommandList* InCommandList, NVSDK_N
         LOG_WARN("Can't get motion vector scales!");
     }
 
-    if (!cfg.FsrUseFsrInputValues.value_or_default() ||
+    if (!Config::Instance()->FsrUseFsrInputValues.value_or_default() ||
         InParameters->Get("FSR.cameraNear", &params.cameraNear) != NVSDK_NGX_Result_Success)
     {
         if (DepthInverted())