Laurens-Packages/Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/RenderPass/Upscalers/PSSRUpscalerPlugin.cs

#if UNITY_PS5
using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine.Experimental.Rendering;
using WW1.PlayStation;

namespace UnityEngine.Rendering.HighDefinition.AMD.PSSR
{
    public class PSSRUpscalerPlugin: UpscalerPlugin
    {
        public override string name => "PSSR";
        public override bool isSupported => UnityEngine.PS5.Utility.isTrinityMode;
        public override bool isMLBased => true;

        private static readonly Queue<uint> ContextPool = new((int)PSSRPlugin.MaxNumContexts);
        private static bool _pluginInitialized;

        public override bool Load()
        {
            try
            {
                if (PSSRPlugin.Init() < 0)
                {
                    Debug.LogError("Failed to initialize PSSR plugin!");
                    _pluginInitialized = false;
                    return false;
                }
            }
            catch (DllNotFoundException)
            {
                Debug.LogError("PSSR plugin not found!");
                _pluginInitialized = false;
                return false;
            }

            for (uint index = 0; index < PSSRPlugin.MaxNumContexts; ++index)
            {
                ContextPool.Enqueue(index);
            }

            _pluginInitialized = true;
            return true;
        }

        public override bool IsLoaded()
        {
            return _pluginInitialized;
        }

        public override UpscalerContext CreateContext(in FSR2CommandInitializationData initSettings)
        {
            if (!_pluginInitialized)
            {
                Debug.LogWarning("PSSR plugin is not initialized!");
                return null;
            }

            if (!ContextPool.TryDequeue(out uint contextIndex))
            {
                Debug.LogError("Ran out of PSSR contexts to allocate!");
                return null;
            }

            PSSRUpscalerContext context = new();
            context.Init(initSettings, contextIndex);
            return context;
        }

        public override async void DestroyContext(UpscalerContext context)
        {
            var pssrContext = (PSSRUpscalerContext)context;
            await pssrContext.Destroy();
            ContextPool.Enqueue(pssrContext.ContextIndex);
        }

        public override bool GetRenderResolutionFromQualityMode(FSR2Quality qualityMode, uint displayWidth, uint displayHeight, out uint renderWidth, out uint renderHeight)
        {
            float ratio = GetUpscaleRatioFromQualityMode(qualityMode);
            renderWidth = (uint)Mathf.RoundToInt(displayWidth / ratio);
            renderHeight = (uint)Mathf.RoundToInt(displayHeight / ratio);
            return true;
        }

        public override float GetUpscaleRatioFromQualityMode(FSR2Quality qualityMode)
        {
            return qualityMode switch
            {
                FSR2Quality.Quality => 1.5f,
                FSR2Quality.Balanced => 1.7f,
                FSR2Quality.Performance => 2.0f,
                FSR2Quality.UltraPerformance => 3.0f,
                _ => 1.0f
            };
        }
    }

    public class PSSRUpscalerContext : UpscalerContext
    {
        private FSR2CommandInitializationData _initData;
        private uint _contextIndex;
        public uint ContextIndex => _contextIndex;

        private Texture2D _outputColor;

        private readonly PSSRPlugin.NativeData<PSSRPlugin.DispatchParams> _dispatchParams = new();
        private readonly PSSRPlugin.NativeData<PSSRPlugin.DestroyParams> _destroyParams = new();
        private readonly PSSRPlugin.NativeData<PSSRPlugin.CaptureParams> _captureParams = new();

        private bool _initialized;

        internal void Init(in FSR2CommandInitializationData initSettings, uint contextIndex)
        {
            _initData = initSettings;
            _contextIndex = contextIndex;

            PSSRPlugin.InitParams initParams;
            initParams.contextIndex = _contextIndex;
            initParams.displayWidth = initSettings.displaySizeWidth;
            initParams.displayHeight = initSettings.displaySizeHeight;
            initParams.maxRenderWidth = initSettings.maxRenderSizeWidth;
            initParams.maxRenderHeight = initSettings.maxRenderSizeHeight;
            initParams.autoKeepCopies = 1u;		// Allow native MFSR context to manage previous frame copies for us

            if (PSSRPlugin.CreateContext(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.
                _outputColor = Texture2D.CreateExternalTexture((int)initSettings.displaySizeWidth, (int)initSettings.displaySizeHeight, TextureFormat.RGBA32, false, true, outputColorTexturePtr);
                _dispatchParams.Initialize();
                _destroyParams.Initialize();
                _captureParams.Initialize();

                _initialized = true;
            }
        }

        internal async Awaitable Destroy()
        {
            // Wait until the GPU is done with the PSSR context before destroying it
            await Awaitable.EndOfFrameAsync();

            if (_initialized)
            {
                PSSRPlugin.DestroyContext(_contextIndex);
                _initialized = false;
            }

            if (_outputColor != null)
            {
                Object.Destroy(_outputColor);
                _outputColor = null;
            }

            _captureParams.Destroy();
            _destroyParams.Destroy();
            _dispatchParams.Destroy();
        }

        public override void Execute(CommandBuffer cmd, in FSR2CommandExecutionData executeData, in FSR2TextureTable textures)
        {
            if (!_initialized || executeData.renderSizeWidth < _initData.displaySizeWidth / 3 || executeData.renderSizeHeight < _initData.displaySizeHeight / 3)
            {
                cmd.Blit(textures.colorInput, textures.colorOutput);
                return;
            }

            cmd.BeginSample("PSSR");

            var flags = PSSRPlugin.OptionFlags.None;
            if ((_initData.ffxFsrFlags & FfxFsr2InitializationFlags.DepthInverted) != 0) flags |= PSSRPlugin.OptionFlags.ReverseDepth;
            if ((_initData.ffxFsrFlags & FfxFsr2InitializationFlags.EnableAutoExposure) != 0) flags |= PSSRPlugin.OptionFlags.AutoExposure;

            ref var dispatchParams = ref _dispatchParams.Value;
            dispatchParams.contextIndex = _contextIndex;

            dispatchParams.color = ToNativePtr(textures.colorInput);
            dispatchParams.depth = ToNativePtr(textures.depth);
            dispatchParams.prevDepth = IntPtr.Zero;    // Native MFSR context will manage previous frame copies internally
            dispatchParams.motionVectors = ToNativePtr(textures.motionVectors);
            dispatchParams.prevMotionVectors = IntPtr.Zero;    // Native MFSR context will manage previous frame copies internally
            dispatchParams.exposure = ToNativePtr(textures.exposureTexture);
            dispatchParams.reactiveMask = ToNativePtr(textures.biasColorMask);
            dispatchParams.outputColor = ToNativePtr(_outputColor);    // This is a specially prepared texture in the exact format that MFSR requires

            Matrix4x4 cameraViewMatrix = executeData.cameraViewMatrix;
            Vector4 cameraPosition = cameraViewMatrix.GetColumn(3);

            dispatchParams.renderWidth = executeData.renderSizeWidth;
            dispatchParams.renderHeight = executeData.renderSizeHeight;
            dispatchParams.jitter = new Vector2(executeData.jitterOffsetX, executeData.jitterOffsetY);
            dispatchParams.motionVectorScale = new Vector2(executeData.MVScaleX, executeData.MVScaleY);
            dispatchParams.camProjectionNoJitter = executeData.cameraProjectionMatrixNoJitter;
            dispatchParams.camForward = -cameraViewMatrix.GetColumn(2);
            dispatchParams.camUp = cameraViewMatrix.GetColumn(1);
            dispatchParams.camRight = cameraViewMatrix.GetColumn(0);
            dispatchParams.camPositionX = cameraPosition.x;
            dispatchParams.camPositionY = cameraPosition.y;
            dispatchParams.camPositionZ = cameraPosition.z;
            dispatchParams.camNear = executeData.cameraNear;
            dispatchParams.camFar = executeData.cameraFar;
            dispatchParams.preExposure = executeData.preExposure;
            dispatchParams.resetHistory = (uint)executeData.reset;
            dispatchParams.flags = flags;

            PSSRPlugin.IssuePluginEvent(cmd, PSSRPlugin.Event.Dispatch, _dispatchParams);

            {
                // Convert the PSSR output from R11G11B10 to the expected destination format
                Vector4 scaleBias = new Vector4((float)textures.colorOutput.width / _initData.displaySizeWidth, (float)textures.colorOutput.height / _initData.displaySizeHeight, 0, 0);
                cmd.SetRenderTarget(textures.colorOutput);
                Blitter.BlitColorAndDepth(cmd, _outputColor, null, scaleBias, 0, false);
            }

            cmd.EndSample("PSSR");
        }

        private static IntPtr ToNativePtr(Texture texture)
        {
            return texture != null ? texture.GetNativeTexturePtr() : IntPtr.Zero;
        }
    }
}
#endif  // UNITY_PS5