From 3a23f8d87465f241eb61550d7b936e58cc897e77 Mon Sep 17 00:00:00 2001
From: Nico de Poel <ndepoel@gmail.com>
Date: Sat, 15 Mar 2025 17:20:40 +0100
Subject: [PATCH] Await end of frame on the CPU before destroying a context, if
 the GPU has a dispatch queued on the current frame. Fixes remaining crashes
 in D3D12 and Vulkan when enabling/disabling the plugin too rapidly.

---
 FSR2UnityPlugin.cpp | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/FSR2UnityPlugin.cpp b/FSR2UnityPlugin.cpp
index 5293552..3a8de3d 100644
--- a/FSR2UnityPlugin.cpp
+++ b/FSR2UnityPlugin.cpp
@@ -46,6 +46,8 @@ struct FSR2Feature
     uint32_t flags;
     bool isValid;
 
+    uint64_t dispatchFrameValue;
+
     FSR2TextureTable textureTable;
 };
 
@@ -53,6 +55,8 @@ static std::vector<FSR2Feature> s_Features;
 static std::queue<uint32_t> s_FeatureSlots;
 static std::mutex s_FeatureMutex;
 
+static HANDLE s_FrameFenceEventHandle = nullptr;
+
 // Unity plugin load event
 extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginLoad(IUnityInterfaces* unityInterfaces)
 {
@@ -190,6 +194,7 @@ extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
             return false;
 
         s_Device = ffxGetDeviceDX12(device);
+        s_FrameFenceEventHandle = CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);
 
         size_t scratchBufferSize = ffxFsr2GetScratchMemorySizeDX12();
         void* scratchBuffer = malloc(scratchBufferSize);
@@ -216,6 +221,29 @@ extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
 static void DestroyFeature(uint32_t featureSlot)
 {
     auto& feature = s_Features[featureSlot];
+
+    if (s_GraphicsD3D12 != nullptr)
+    {
+        // If there's still an upscale dispatch executing on the current frame, wait until rendering is finished before destroying its context.
+        // Thanks to https://alextardif.com/D3D11To12P1.html for explaining how fences work and how to make the CPU wait for a command queue to complete.
+        ID3D12Fence* frameFence = s_GraphicsD3D12->GetFrameFence();
+        if (feature.dispatchFrameValue > frameFence->GetCompletedValue())
+        {
+            frameFence->SetEventOnCompletion(feature.dispatchFrameValue, s_FrameFenceEventHandle);
+            WaitForSingleObjectEx(s_FrameFenceEventHandle, INFINITE, false);
+        }
+    }
+    else if (s_GraphicsVulkan != nullptr)
+    {
+        UnityVulkanRecordingState state;
+        s_GraphicsVulkan->CommandRecordingState(&state, kUnityVulkanGraphicsQueueAccess_DontCare);
+        if (feature.dispatchFrameValue > state.safeFrameNumber)
+        {
+            UnityVulkanInstance instance = s_GraphicsVulkan->Instance();
+            vkQueueWaitIdle(instance.graphicsQueue);
+        }
+    }
+
     if (feature.isValid)
     {
         ffxFsr2ContextDestroy(&feature.upscalingContext);
@@ -237,6 +265,12 @@ extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_ShutdownApi()
         free(s_Fsr2Interface.scratchBuffer);
         s_Fsr2Interface.scratchBuffer = nullptr;
     }
+
+    if (s_FrameFenceEventHandle != nullptr)
+    {
+        CloseHandle(s_FrameFenceEventHandle);
+        s_FrameFenceEventHandle = nullptr;
+    }
     
     s_Device = nullptr;
 }
@@ -378,6 +412,9 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             s_GraphicsD3D12->CommandRecordingState(&state);
             dispatchDescription.commandList = ffxGetCommandListDX12(state.commandList);
 
+            // Keep track of which frame this dispatch is happening on, so we can verify when it's finished
+            feature.dispatchFrameValue = s_GraphicsD3D12->GetNextFrameFenceValue();
+
             dispatchDescription.color = ffxGetResourceDX12(&feature.upscalingContext, (ID3D12Resource*)feature.textureTable.colorInput.image);
             dispatchDescription.depth = ffxGetResourceDX12(&feature.upscalingContext, (ID3D12Resource*)feature.textureTable.depth.image);
             dispatchDescription.motionVectors = ffxGetResourceDX12(&feature.upscalingContext, (ID3D12Resource*)feature.textureTable.motionVectors.image);
@@ -392,6 +429,9 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             s_GraphicsVulkan->CommandRecordingState(&state, kUnityVulkanGraphicsQueueAccess_DontCare);
             dispatchDescription.commandList = ffxGetCommandListVK(state.commandBuffer);
 
+            // Keep track of which frame this dispatch is happening on, so we can verify when it's finished
+            feature.dispatchFrameValue = state.currentFrameNumber;
+
             UnityVulkanInstance instance = s_GraphicsVulkan->Instance();
             dispatchDescription.color = GetVulkanTextureResource(instance, feature, feature.textureTable.colorInput);
             dispatchDescription.depth = GetVulkanTextureResource(instance, feature, feature.textureTable.depth);