Reworked plugin structure so that up to four MFSR context can be created and used, with Unity being able to select which context to use.

Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PSSRUpscaler.cs

@@ -43,6 +43,7 @@ namespace UnityEngine.Rendering.PostProcessing
 			}
 
 			PSSRPlugin.InitParams initParams;
+			initParams.contextIndex = 0;
 			initParams.displayWidth = (uint)config.UpscaleSize.x;
 			initParams.displayHeight = (uint)config.UpscaleSize.y;
 			initParams.maxRenderWidth = (uint)config.MaxRenderSize.x;
@@ -68,11 +69,14 @@ namespace UnityEngine.Rendering.PostProcessing
 		{
 			base.DestroyContext();
 			
-			// TODO: this still race conditions sometimes
 			if (_contextInitialized)
 			{
+				_dispatchParams.contextIndex = 0;
+				Marshal.StructureToPtr(_dispatchParams, _dispatchParamsBuffer, false);
+				
+				// Destroying the context from the render thread reduces the risk of race conditions causing crashes
 				var cmd = new CommandBuffer();
-				cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 2, IntPtr.Zero);
+				cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 2, _dispatchParamsBuffer);
 				Graphics.ExecuteCommandBuffer(cmd);
 				cmd.Release();
 				
@@ -134,7 +138,8 @@ namespace UnityEngine.Rendering.PostProcessing
 			var flags = PSSRPlugin.OptionFlags.None;
 			if (SystemInfo.usesReversedZBuffer) flags |= PSSRPlugin.OptionFlags.ReverseDepth;
 			if (config.exposureSource == Upscaling.ExposureSource.Auto) flags |= PSSRPlugin.OptionFlags.AutoExposure;
-			
+
+			_dispatchParams.contextIndex = 0;
 			_dispatchParams.color = ToNativePtr(_inputColor);
 			_dispatchParams.depth = ToNativePtr(_inputDepth[frameIndex]);
 			_dispatchParams.prevDepth = ToNativePtr(_inputDepth[frameIndex ^ 1]);
@@ -167,7 +172,7 @@ namespace UnityEngine.Rendering.PostProcessing
 			
 			if (config.performSharpenPass)
 			{
-				// PSSR output is already pretty sharp and we don't want to over-sharpen the image, so cut the sharpness range by half
+				// PSSR output is already pretty sharp, and we don't want to over-sharpen the image, so cut the sharpness range by half
 				ApplySharpening(cmd, context, config.UpscaleSize, config.sharpness * 0.5f, _outputColor, context.destination);
 			}
 			else
@@ -206,6 +211,8 @@ namespace UnityEngine.Rendering.PostProcessing
 			[Serializable, StructLayout(LayoutKind.Sequential)]
 			public struct InitParams
 			{
+				public uint contextIndex;
+				
 				public uint displayWidth;
 				public uint displayHeight;
 				public uint maxRenderWidth;
@@ -217,6 +224,8 @@ namespace UnityEngine.Rendering.PostProcessing
 			[Serializable, StructLayout(LayoutKind.Sequential)]
 			public struct DispatchParams
 			{
+				public uint contextIndex;
+				
 				public IntPtr color;
 				public IntPtr depth;
 				public IntPtr prevDepth;

Tools/PSSRPlugin/pssrplugin.cpp

@@ -23,16 +23,21 @@ static IUnityGraphicsPS5* s_GraphicsPS5 = nullptr;          // Old Gnmp-based gr
 static IUnityGraphicsAgcPS5* s_GraphicsAgcPS5 = nullptr;    // New Agc-based graphics API (NGGC)
 static UnityGfxRenderer s_RendererType = kUnityGfxRendererNull;
 
-MfsrPhysicalMemoryBlock* s_mfsrSharedResourcesMemoryBlocks = nullptr;
-MfsrSharedResources s_mfsrSharedResources = nullptr;
+static bool s_mfsrInitialized = false;
+static MfsrPhysicalMemoryBlock* s_mfsrSharedResourcesMemoryBlocks = nullptr;
+static MfsrSharedResources s_mfsrSharedResources = nullptr;
 
-// Just use a single context for now to keep things simple
 // "Multiple contexts (up to 4) can be generated to perform multiple MFSR processes within an application."
-static bool s_mfsrInitialized = false;
-static Core::BasicContext s_AgcContext;
-static MfsrPhysicalMemoryBlock* s_mfsrMemoryBlocks = nullptr;
-static MfsrContext s_mfsrContext = nullptr;
-static Core::Texture s_outputColorTexture;
+static const int MaxNumContexts = 4;
+
+struct PssrContext
+{
+    Core::BasicContext agcContext;
+    MfsrPhysicalMemoryBlock* mfsrMemoryBlocks = nullptr;
+    MfsrContext mfsrContext = nullptr;
+    Core::Texture outputColorTexture;
+};
+static PssrContext s_contexts[MaxNumContexts];
 
 static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType);
 static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data);
@@ -200,6 +205,8 @@ extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API ReleasePssr()
 
 typedef struct pssr_init_params_s
 {
+    uint32_t contextIndex;
+
     uint32_t displayWidth;
     uint32_t displayHeight;
     uint32_t maxRenderWidth;
@@ -210,6 +217,8 @@ typedef struct pssr_init_params_s
 
 typedef struct pssr_dispatch_params_s
 {
+    uint32_t contextIndex;
+
     sce::Agc::Core::Texture* color;
     sce::Agc::Core::Texture* depth;
     sce::Agc::Core::Texture* prevDepth;
@@ -240,6 +249,16 @@ typedef struct pssr_dispatch_params_s
 
 extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(const pssr_init_params_t* params, Core::Texture** outputColorTexture)
 {
+    if (params->contextIndex < 0 || params->contextIndex >= MaxNumContexts)
+    {
+        std::stringstream msg;
+        msg << "Invalid PSSR context index: " << params->contextIndex;
+        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
+        return -1;
+    }
+
+    PssrContext& context = s_contexts[params->contextIndex];
+
     MfsrContextInitParameters initParams = {};
     initParams.init();
     initParams.m_outputColorWidth = params->displayWidth;
@@ -268,7 +287,7 @@ extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(
         MfsrContextInitRequirement mfsrInitRequirement = {};
         getMfsrContextInitRequirement(&mfsrInitRequirement, &initParams);
 
-        s_mfsrMemoryBlocks = new MfsrPhysicalMemoryBlock[mfsrInitRequirement.m_physicalMemoryBlockCount];
+        context.mfsrMemoryBlocks = new MfsrPhysicalMemoryBlock[mfsrInitRequirement.m_physicalMemoryBlockCount];
 
         size_t size = mfsrInitRequirement.m_physicalMemoryBlockSize;
         size_t align = mfsrInitRequirement.m_physicalMemoryBlockAlignment;
@@ -279,12 +298,12 @@ extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(
 
         for (uint32_t i = 0; i < mfsrInitRequirement.m_physicalMemoryBlockCount; i++)
         {
-            s_mfsrMemoryBlocks[i].m_size = size;
-            s_mfsrMemoryBlocks[i].m_offset = m_mfsrMemoryPhysicalAddress + mfsrInitRequirement.m_physicalMemoryBlockSize * i;
+            context.mfsrMemoryBlocks[i].m_size = size;
+            context.mfsrMemoryBlocks[i].m_offset = m_mfsrMemoryPhysicalAddress + mfsrInitRequirement.m_physicalMemoryBlockSize * i;
         }
 
         initParams.m_sharedResources = s_mfsrSharedResources;
-        initParams.m_physicalMemoryBlocks = s_mfsrMemoryBlocks;
+        initParams.m_physicalMemoryBlocks = context.mfsrMemoryBlocks;
         initParams.m_physicalMemoryBlockCount = mfsrInitRequirement.m_physicalMemoryBlockCount;
     }
 
@@ -309,7 +328,7 @@ extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(
             texSpec.m_dataAddress = s_GraphicsAgcPS5->AllocateGPUMemory(sizeAlign.m_size, MfsrAlignment::kOutputTexture);
         }
 
-        int res = Core::initialize(&s_outputColorTexture, &texSpec);
+        int res = Core::initialize(&context.outputColorTexture, &texSpec);
         if (res != SCE_OK)
         {
             std::stringstream msg;
@@ -318,20 +337,20 @@ extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(
             return res;
         }
 
-        *outputColorTexture = &s_outputColorTexture;
+        *outputColorTexture = &context.outputColorTexture;
     }
 
     // Set up command buffer for NGGC
     if (s_RendererType == kUnityGfxRendererPS5NGGC && s_GraphicsAgcPS5 != nullptr)
     {
         const size_t agcContextSize = 128 * 1024;
-        s_AgcContext.m_dcb.init(s_GraphicsAgcPS5->AllocateGPUMemory(agcContextSize, Alignment::kCommandBuffer), agcContextSize);
-        s_AgcContext.m_bdr.init(&s_AgcContext.m_dcb, &s_AgcContext.m_dcb);
-        s_AgcContext.m_sb.init(256, &s_AgcContext.m_dcb, &s_AgcContext.m_dcb);
+        context.agcContext.m_dcb.init(s_GraphicsAgcPS5->AllocateGPUMemory(agcContextSize, Alignment::kCommandBuffer), agcContextSize);
+        context.agcContext.m_bdr.init(&context.agcContext.m_dcb, &context.agcContext.m_dcb);
+        context.agcContext.m_sb.init(256, &context.agcContext.m_dcb, &context.agcContext.m_dcb);
     }
 
     // Finally, create the actual MFSR context
-    int res = createMfsrContext(&s_mfsrContext, &initParams);
+    int res = createMfsrContext(&context.mfsrContext, &initParams);
     if (res != SCE_OK)
     {
         std::stringstream msg;
@@ -344,8 +363,18 @@ extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(
     return SCE_OK;
 }
 
-extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API DestroyPssrContext()
+extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API DestroyPssrContext(uint32_t contextIndex)
 {
+    if (contextIndex < 0 || contextIndex >= MaxNumContexts)
+    {
+        std::stringstream msg;
+        msg << "Invalid PSSR context index: " << contextIndex;
+        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
+        return;
+    }
+
+    PssrContext& context = s_contexts[contextIndex];
+
     // TODO: this still crashes often due to race conditions between dispatch above and release here
     // The main cause will be that the MFSR context is released before the already dispatched commands are executed on the GPU
     // Just adding mutexes everywhere is not going to fix that
@@ -353,36 +382,36 @@ extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API DestroyPssrContext()
     // 
     // Another issue is that we reuse the same static s_mfsrContext pointer when recreating the MFSR context, which can lead to situations where dispatch uses the wrong context object
     // Multi-buffering those pointers will help with that
-    if (s_mfsrContext != nullptr)
+    if (context.mfsrContext != nullptr)
     {
-        releaseMfsrContext(s_mfsrContext);
-        s_mfsrContext = nullptr;
+        releaseMfsrContext(context.mfsrContext);
+        context.mfsrContext = nullptr;
     }
 
-    if (s_GraphicsAgcPS5 != nullptr && s_AgcContext.m_dcb.m_bottom != nullptr)
+    if (s_GraphicsAgcPS5 != nullptr && context.agcContext.m_dcb.m_bottom != nullptr)
     {
-        s_GraphicsAgcPS5->ReleaseGPUMemory(s_AgcContext.m_dcb.m_bottom);
-        s_AgcContext.m_dcb.clear();
+        s_GraphicsAgcPS5->ReleaseGPUMemory(context.agcContext.m_dcb.m_bottom);
+        context.agcContext.m_dcb.clear();
     }
 
-    if (s_outputColorTexture.getDataAddress())
+    if (context.outputColorTexture.getDataAddress())
     {
         if (s_GraphicsPS5 != nullptr)
         {
-            s_GraphicsPS5->ReleaseGPUMemory(s_outputColorTexture.getDataAddress());
+            s_GraphicsPS5->ReleaseGPUMemory(context.outputColorTexture.getDataAddress());
         }
         else if (s_GraphicsAgcPS5 != nullptr)
         {
-            s_GraphicsAgcPS5->ReleaseGPUMemory(s_outputColorTexture.getDataAddress());
+            s_GraphicsAgcPS5->ReleaseGPUMemory(context.outputColorTexture.getDataAddress());
         }
 
-        s_outputColorTexture.init();
+        context.outputColorTexture.init();
     }
 
-    if (s_mfsrMemoryBlocks != nullptr)
+    if (context.mfsrMemoryBlocks != nullptr)
     {
-        delete[] s_mfsrMemoryBlocks;
-        s_mfsrMemoryBlocks = nullptr;
+        delete[] context.mfsrMemoryBlocks;
+        context.mfsrMemoryBlocks = nullptr;
     }
 
     UNITY_LOG(s_Log, "Destroyed PSSR context");
@@ -411,10 +440,18 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
         }
         case 1: // Execute PSSR
         {
-            if (s_mfsrContext == nullptr)
+            auto* params = (pssr_dispatch_params_t*)data;
+            if (params->contextIndex < 0 || params->contextIndex >= MaxNumContexts)
+            {
+                std::stringstream msg;
+                msg << "Invalid PSSR context index: " << params->contextIndex;
+                UNITY_LOG_ERROR(s_Log, msg.str().c_str());
                 break;
+            }
 
-            auto* params = (pssr_dispatch_params_t*)data;
+            PssrContext& context = s_contexts[params->contextIndex];
+            if (context.mfsrContext == nullptr)
+                break;
 
             // How we obtain a command buffer to dispatch to depends on which graphics API we're using in Unity
             sce::Agc::DrawCommandBuffer* cmd = nullptr;
@@ -425,7 +462,7 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             }
             else if (s_RendererType == kUnityGfxRendererPS5NGGC)
             {
-                Core::BasicContext& ctx = s_AgcContext.reset();
+                Core::BasicContext& ctx = context.agcContext.reset();
                 cmd = &ctx.m_dcb;
             }
 
@@ -457,7 +494,7 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             dispatchParams.m_reset = params->resetHistory != 0;
             dispatchParams.m_flags = params->flags;
 
-            dispatchMfsr(s_mfsrContext, cmd, &dispatchParams);
+            dispatchMfsr(context.mfsrContext, cmd, &dispatchParams);
 
             cmd->popMarker();
 
@@ -469,7 +506,8 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
         }
         case 2: // Destroy PSSR context
         {
-            DestroyPssrContext();
+            uint32_t* contextIndex = (uint32_t*)data;
+            DestroyPssrContext(*contextIndex);
             break;
         }
     }