First working proof-of-concept implementation of PSSR

1 year ago · 8d973cf234
11 changed files with 404 additions and 0 deletions
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/PostProcessResources.asset
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/PostProcessResources.asset
@ -127,6 +127,7 @@ MonoBehaviour:
    multiScaleAORender: {fileID: 7200000, guid: 34a460e8a2e66c243a9c12024e5a798d, type: 3}
    multiScaleAOUpsample: {fileID: 7200000, guid: 600d6212b59bb40409d19d750b5fd1e9, type: 3}
    gaussianDownsample: {fileID: 7200000, guid: 6dba4103d23a7904fbc49099355aff3e, type: 3}
+    prepareInputs: {fileID: 7200000, guid: d27c88814f02b424088b06503e1bc9d5, type: 3}
    casSharpening: {fileID: 7200000, guid: 00e3ffafadd35564780d8a12adcbeff7, type: 3}
    fsr2Upscaler:
      computeLuminancePyramidPass: {fileID: 7200000, guid: 04c3480675e29a340808141e68d4cc8b, type: 3}
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling.cs
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling.cs
@ -180,6 +180,8 @@ namespace UnityEngine.Rendering.PostProcessing
 				{
 					UpscalerType.FSR2 when FSR2Upscaler.IsSupported => new FSR2Upscaler(),
 					UpscalerType.FSR3 when FSR3Upscaler.IsSupported => new FSR3Upscaler(),
+					//UpscalerType.SGSR2 when SGSR2Upscaler.IsSupported => new SGSR2Upscaler(),
+					UpscalerType.PSSR when PSSRUpscaler.IsSupported => new PSSRUpscaler(),
 					_ => new FSR2Upscaler(),	// Fallback for when the selected upscaler is not supported on the current hardware
 				};

--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5.meta
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5.meta
@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 32c5585f0ac1473aaa942ede8bc01712
+timeCreated: 1730718066
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5/PSSRPlugin.prx
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5/PSSRPlugin.prx
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5/PSSRPlugin.prx.meta
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PS5/PSSRPlugin.prx.meta
@ -0,0 +1,77 @@
+fileFormatVersion: 2
+guid: 82cb3c505a5252e4db0b661d9ccd67cb
+PluginImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  iconMap: {}
+  executionOrder: {}
+  defineConstraints: []
+  isPreloaded: 1
+  isOverridable: 1
+  isExplicitlyReferenced: 0
+  validateReferences: 1
+  platformData:
+  - first:
+      : Any
+    second:
+      enabled: 0
+      settings:
+        Exclude Editor: 1
+        Exclude GameCoreScarlett: 1
+        Exclude GameCoreXboxOne: 1
+        Exclude Linux64: 1
+        Exclude OSXUniversal: 1
+        Exclude PS4: 1
+        Exclude PS5: 0
+        Exclude Win: 1
+        Exclude Win64: 1
+  - first:
+      Any: 
+    second:
+      enabled: 0
+      settings: {}
+  - first:
+      Editor: Editor
+    second:
+      enabled: 0
+      settings:
+        CPU: AnyCPU
+        DefaultValueInitialized: true
+        OS: AnyOS
+  - first:
+      PS4: PS4
+    second:
+      enabled: 0
+      settings: {}
+  - first:
+      PS5: PS5
+    second:
+      enabled: 1
+      settings: {}
+  - first:
+      Standalone: Linux64
+    second:
+      enabled: 0
+      settings:
+        CPU: AnyCPU
+  - first:
+      Standalone: OSXUniversal
+    second:
+      enabled: 0
+      settings:
+        CPU: AnyCPU
+  - first:
+      Standalone: Win
+    second:
+      enabled: 0
+      settings:
+        CPU: AnyCPU
+  - first:
+      Standalone: Win64
+    second:
+      enabled: 0
+      settings:
+        CPU: AnyCPU
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PSSRUpscaler.cs
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PSSRUpscaler.cs
@ -0,0 +1,230 @@
+using System;
+using System.Runtime.InteropServices;
+using UnityEngine.Experimental.Rendering;
+
+namespace UnityEngine.Rendering.PostProcessing
+{
+	internal class PSSRUpscaler : Upscaler
+	{
+#if UNITY_PS5
+		public static bool IsSupported => UnityEngine.PS5.Utility.isTrinityMode;
+
+		private PSSRPlugin.DispatchParams _dispatchParams;
+		private IntPtr _dispatchParamsBuffer;
+
+		private RenderTexture _inputColor;
+		private RenderTexture _inputDepth;
+		private RenderTexture _outputColor;
+		private Texture2D _outputIntermediate;
+		
+		private bool _pluginInitialized;
+		private bool _contextInitialized;
+
+		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
+		{
+			if (!PSSRPlugin.InitPssr())
+			{
+				Debug.LogError("Failed to initialize PSSR plugin!");
+				_pluginInitialized = false;
+				return;
+			}
+
+			_pluginInitialized = true;
+
+			PSSRPlugin.InitParams initParams;
+			initParams.displayWidth = (uint)config.UpscaleSize.x;
+			initParams.displayHeight = (uint)config.UpscaleSize.y;
+			initParams.maxRenderWidth = (uint)config.MaxRenderSize.x;
+			initParams.maxRenderHeight = (uint)config.MaxRenderSize.y;
+
+			CreateRenderTexture(ref _inputColor, "PSSR Input Color", config.MaxRenderSize, context.sourceFormat, true);
+			CreateRenderTexture(ref _inputDepth, "PSSR Input Depth", config.MaxRenderSize, GraphicsFormat.R32_SFloat, true);
+			CreateRenderTexture(ref _outputColor, "PSSR Output Color", config.UpscaleSize, RenderTextureFormat.RGB111110Float);
+
+			if (PSSRPlugin.CreatePssrContext(ref initParams, out IntPtr outputColorTexturePtr) >= 0 && outputColorTexturePtr != IntPtr.Zero)
+			{
+				// PSSR requires an output color texture in a very particular format (k11_11_10Float with kStandard256B tile mode and a specific alignment) that Unity cannot create directly.
+				// So instead we let the plugin create that texture and then import it into Unity as a generic 32bpp texture from a native pointer.
+				_outputIntermediate = Texture2D.CreateExternalTexture(config.UpscaleSize.x, config.UpscaleSize.y, TextureFormat.RGBA32, false, true, outputColorTexturePtr);
+				_dispatchParamsBuffer = Marshal.AllocHGlobal(Marshal.SizeOf<PSSRPlugin.DispatchParams>());
+				_contextInitialized = true;
+			}
+		}
+
+		public override void DestroyContext()
+		{
+			base.DestroyContext();
+			
+			// TODO: this still race conditions sometimes
+			if (_contextInitialized)
+			{
+				var cmd = new CommandBuffer();
+				cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 2, IntPtr.Zero);
+				Graphics.ExecuteCommandBuffer(cmd);
+				cmd.Release();
+				
+				_contextInitialized = false;
+			}
+
+			DestroyRenderTexture(ref _outputColor);
+			DestroyRenderTexture(ref _inputDepth);
+			DestroyRenderTexture(ref _inputColor);
+
+			if (_dispatchParamsBuffer != IntPtr.Zero)
+			{
+				Marshal.FreeHGlobal(_dispatchParamsBuffer);
+				_dispatchParamsBuffer = IntPtr.Zero;
+			}
+
+			if (_pluginInitialized)
+			{
+				PSSRPlugin.ReleasePssr();
+				_pluginInitialized = false;
+			}
+		}
+
+		private static readonly int CameraMotionVectorsTexture = Shader.PropertyToID("_CameraMotionVectorsTexture");
+
+		public override void Render(PostProcessRenderContext context, Upscaling config)
+		{
+			var cmd = context.command;
+			if (!_pluginInitialized || !_contextInitialized)
+			{
+				cmd.BlitFullscreenTriangle(context.source, context.destination);
+				return;
+			}
+			
+			cmd.BeginSample("PSSR");
+			
+			// TODO: if PSSR needs a copy of the previous depth and motion vectors anyway, why not just double-buffer those resources ourselves and make good use of these otherwise dumb copies?
+			PrepareInputs(cmd, context, config, _inputColor, _inputDepth);
+			
+			_dispatchParams.color = ToNativePtr(_inputColor);
+			_dispatchParams.depth = ToNativePtr(_inputDepth);
+			_dispatchParams.motionVectors = ToNativePtr(Shader.GetGlobalTexture(CameraMotionVectorsTexture));
+			_dispatchParams.exposure = ToNativePtr(config.exposureSource switch
+			{
+				Upscaling.ExposureSource.Manual when config.exposure != null => config.exposure,
+				Upscaling.ExposureSource.Unity => context.autoExposureTexture,
+				_ => null,
+			});
+			_dispatchParams.reactiveMask = ToNativePtr(config.reactiveMask);
+			_dispatchParams.outputColor = ToNativePtr(_outputIntermediate);
+
+			var scaledRenderSize = config.GetScaledRenderSize(context.camera);
+			_dispatchParams.renderWidth = (uint)scaledRenderSize.x;
+			_dispatchParams.renderHeight = (uint)scaledRenderSize.y;
+			_dispatchParams.jitter = config.JitterOffset;		// TODO: may need to scale this (should be jitter in render pixels apparently)
+			_dispatchParams.motionVectorScale = new Vector2(-scaledRenderSize.x, -scaledRenderSize.y);
+			_dispatchParams.FromCamera(context.camera);
+			_dispatchParams.preExposure = config.preExposure;
+			_dispatchParams.resetHistory = config.Reset ? 1u : 0u;
+			
+			Marshal.StructureToPtr(_dispatchParams, _dispatchParamsBuffer, false);
+			cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 1, _dispatchParamsBuffer);
+			
+			// Convert the native texture into Unity-land using a blind texture data copy
+			cmd.CopyTexture(_outputIntermediate, _outputColor);
+			
+			// TODO: optional sharpening
+			cmd.BlitFullscreenTriangle(_outputColor, context.destination);
+			
+			cmd.EndSample("PSSR");
+		}
+		
+		private static IntPtr ToNativePtr(Texture texture)
+		{
+			return texture != null ? texture.GetNativeTexturePtr() : IntPtr.Zero;
+		}
+
+		private static class PSSRPlugin
+		{
+			private const string LibraryName = "PSSRPlugin.prx";
+
+			[DllImport(LibraryName)]
+			public static extern IntPtr GetRenderEventAndDataFunc();
+
+			[DllImport(LibraryName)]
+			public static extern bool InitPssr();
+
+			[DllImport(LibraryName)]
+			public static extern void ReleasePssr();
+
+			[DllImport(LibraryName)]
+			public static extern int CreatePssrContext(ref InitParams initParams, out IntPtr outputColorTexturePtr);
+
+			[Serializable, StructLayout(LayoutKind.Sequential)]
+			public struct InitParams
+			{
+				public uint displayWidth;
+				public uint displayHeight;
+				public uint maxRenderWidth;
+				public uint maxRenderHeight;
+			}
+
+			[Serializable, StructLayout(LayoutKind.Sequential)]
+			public struct DispatchParams
+			{
+				public IntPtr color;
+				public IntPtr depth;
+				public IntPtr motionVectors;
+				public IntPtr exposure;
+				public IntPtr reactiveMask;
+				public IntPtr outputColor;
+
+				public uint renderWidth;
+				public uint renderHeight;
+
+				public Vector2 jitter;
+				public Vector2 motionVectorScale;
+
+				public Matrix4x4 camProjection;
+				public Vector3 camForward;
+				public Vector3 camUp;
+				public Vector3 camRight;
+				public double camPositionX;
+				public double camPositionY;
+				public double camPositionZ;
+				public float pad;
+
+				public float camNear;
+				public float camFar;
+				public float preExposure;
+				public uint resetHistory;
+
+				public void FromCamera(Camera cam)
+				{
+					camProjection = cam.projectionMatrix;
+					camForward = cam.transform.forward;
+					camUp = cam.transform.up;
+					camRight = cam.transform.right;
+					camPositionX = cam.transform.position.x;
+					camPositionY = cam.transform.position.y;
+					camPositionZ = cam.transform.position.z;
+					pad = 0f;
+
+					camNear = cam.nearClipPlane;
+					camFar = cam.farClipPlane;
+				}
+			}
+		}
+#else
+		public static bool IsSupported => false;
+		
+		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
+		{
+			throw new NotImplementedException();
+		}
+		
+		public override void DestroyContext()
+		{
+			base.DestroyContext();
+		}
+		
+		public override void Render(PostProcessRenderContext context, Upscaling config)
+		{
+			throw new NotImplementedException();
+		}
+#endif
+	}
+}
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PSSRUpscaler.cs.meta
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/PSSRUpscaler.cs.meta
@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: dfe32e853f4042ab8ce14f5cdc3c3151
+timeCreated: 1730636826
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/Upscaler.cs
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/Effects/Upscaling/Upscaler.cs
@ -19,6 +19,55 @@ namespace UnityEngine.Rendering.PostProcessing
 		
 		public abstract void Render(PostProcessRenderContext context, Upscaling config);

+		private ConstantsBuffer<PrepareInputsConstants> _prepareInputsConstants;
+		
+		protected bool PrepareInputs(CommandBuffer cmd, PostProcessRenderContext context, Upscaling config, RenderTexture rwColor = null, RenderTexture rwDepth = null, RenderTexture rwMotionVectors = null)
+		{
+			ComputeShader shader = context.resources.computeShaders.prepareInputs;
+			if (shader == null)
+				return false;
+
+			_prepareInputsConstants ??= ConstantsBuffer<PrepareInputsConstants>.Create();
+			
+			Vector2Int scaledRenderSize = config.GetScaledRenderSize(context.camera);
+			
+			const int threadGroupWorkRegionDim = 8;
+			int dispatchSrcX = (scaledRenderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
+			int dispatchSrcY = (scaledRenderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
+			
+			cmd.BeginSample("Prepare Inputs");
+			
+			_prepareInputsConstants.Value.prepareColor = rwColor != null ? 1 : 0;
+			_prepareInputsConstants.Value.prepareDepth = rwDepth != null ? 1 : 0;
+			_prepareInputsConstants.Value.prepareMotionVectors = rwMotionVectors != null ? 1 : 0;
+			_prepareInputsConstants.UpdateBufferData(cmd);
+			
+			int kernelIndex = shader.FindKernel("CS");
+			cmd.SetComputeTextureParam(shader, kernelIndex, "InColor", context.source, 0, RenderTextureSubElement.Color);
+			cmd.SetComputeTextureParam(shader, kernelIndex, "InDepth", Upscaling.GetDepthTexture(context.camera), 0, RenderTextureSubElement.Depth);
+			cmd.SetComputeTextureParam(shader, kernelIndex, "InMotionVectors", BuiltinRenderTextureType.MotionVectors);
+			
+			cmd.SetComputeTextureParam(shader, kernelIndex, "OutColor", rwColor != null ? rwColor : BuiltinRenderTextureType.None);
+			cmd.SetComputeTextureParam(shader, kernelIndex, "OutDepth", rwDepth != null ? rwDepth : BuiltinRenderTextureType.None);
+			cmd.SetComputeTextureParam(shader, kernelIndex, "OutMotionVectors", rwMotionVectors != null ? rwMotionVectors : BuiltinRenderTextureType.None);
+			
+			cmd.SetComputeConstantBufferParam(shader, "Params", _prepareInputsConstants, 0, Marshal.SizeOf<PrepareInputsConstants>());
+			
+			cmd.DispatchCompute(shader, kernelIndex, dispatchSrcX, dispatchSrcY, 1);
+			
+			cmd.EndSample("Prepare Inputs");
+			return true;
+		}
+		
+		[Serializable, StructLayout(LayoutKind.Sequential)]
+		internal struct PrepareInputsConstants
+		{
+			public int prepareColor;
+			public int prepareDepth;
+			public int prepareMotionVectors;
+			public int pad;
+		}
+
 		private ConstantsBuffer<GenerateReactiveConstants> _reactiveMaskConstants;
 		private RenderTexture _reactiveMask;
 		
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/PostProcessResources.cs
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Runtime/PostProcessResources.cs
@ -216,6 +216,11 @@ namespace UnityEngine.Rendering.PostProcessing
            /// </summary>
            public ComputeShader gaussianDownsample;

+            /// <summary>
+            /// The compute shader used to prepare inputs for upscaling.
+            /// </summary>
+            public ComputeShader prepareInputs;
+
            /// <summary>
            /// The compute shader used by the CAS sharpening filter.
            /// </summary>
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Shaders/Builtins/PrepareInputs.compute
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Shaders/Builtins/PrepareInputs.compute
@ -0,0 +1,27 @@
+#pragma kernel CS
+
+Texture2D           InColor: register(t0);
+Texture2D<float>    InDepth: register(t1);
+//Texture2D<float2>   InMotionVectors: register(t2);
+
+RWTexture2D<float4> OutColor: register(u0);
+RWTexture2D<float>  OutDepth: register(u1);
+//RWTexture2D<float2> OutMotionVectors: register(u2);
+
+cbuffer Params: register(b0)
+{
+    int prepareColor;
+    int prepareDepth;
+    int prepareMotionVectors;
+    int pad;
+};
+
+[numthreads(8, 8, 1)]
+void CS(uint2 GroupId: SV_GroupID, uint2 GroupThreadId: SV_GroupThreadID)
+{
+    uint2 InputPos = GroupId * uint2(8, 8) + GroupThreadId;
+
+    if (prepareColor) OutColor[InputPos] = InColor[InputPos];
+    if (prepareDepth) OutDepth[InputPos] = InDepth[InputPos];
+    //if (prepareMotionVectors) OutMotionVectors[InputPos] = InMotionVectors[InputPos];
+}
--- a/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Shaders/Builtins/PrepareInputs.compute.meta
+++ b/Packages/com.unity.postprocessing@3.2.2/PostProcessing/Shaders/Builtins/PrepareInputs.compute.meta
@ -0,0 +1,7 @@
+fileFormatVersion: 2
+guid: d27c88814f02b424088b06503e1bc9d5
+ComputeShaderImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: