FSRTestBuiltin/Assets/Scripts/Fsr2Context.cs

using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Rendering;

namespace FidelityFX
{
    public class Fsr2Context
    {
        private const int MaxQueuedFrames = 16;
        
        private Fsr2.ContextDescription _contextDescription;
        private CommandBuffer _commandBuffer;
        
        private Fsr2Pipeline _depthClipPipeline;
        private Fsr2Pipeline _reconstructPreviousDepthPipeline;
        private Fsr2Pipeline _lockPipeline;
        private Fsr2Pipeline _accumulatePipeline;
        private Fsr2Pipeline _accumulateSharpenPipeline;
        private Fsr2Pipeline _rcasPipeline;
        private Fsr2Pipeline _computeLuminancePyramidPipeline;
        private Fsr2Pipeline _generateReactivePipeline;
        private Fsr2Pipeline _tcrAutogeneratePipeline;

        private readonly Fsr2Resources _resources = new Fsr2Resources();

        private ComputeBuffer _fsr2ConstantsBuffer;
        private readonly Fsr2.Fsr2Constants[] _fsr2ConstantsArray = { new Fsr2.Fsr2Constants() };
        private ref Fsr2.Fsr2Constants Constants => ref _fsr2ConstantsArray[0];

        private ComputeBuffer _spdConstantsBuffer;
        private readonly Fsr2.SpdConstants[] _spdConstantsArray = { new Fsr2.SpdConstants() };
        private ref Fsr2.SpdConstants SpdConsts => ref _spdConstantsArray[0];
    
        private ComputeBuffer _rcasConstantsBuffer;
        private readonly Fsr2.RcasConstants[] _rcasConstantsArray = new Fsr2.RcasConstants[1];
        private ref Fsr2.RcasConstants RcasConsts => ref _rcasConstantsArray[0];

        private ComputeBuffer _generateReactiveConstantsBuffer;
        private readonly Fsr2.GenerateReactiveConstants[] _generateReactiveConstantsArray = { new Fsr2.GenerateReactiveConstants() };
        private ref Fsr2.GenerateReactiveConstants GenReactiveConsts => ref _generateReactiveConstantsArray[0];

        private bool _firstExecution;
        private Vector2 _previousJitterOffset;
        private int _resourceFrameIndex;

        public void Create(Fsr2.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;
            _commandBuffer = new CommandBuffer { name = "FSR2" };
            
            _fsr2ConstantsBuffer = CreateConstantBuffer<Fsr2.Fsr2Constants>();
            _spdConstantsBuffer = CreateConstantBuffer<Fsr2.SpdConstants>();
            _rcasConstantsBuffer = CreateConstantBuffer<Fsr2.RcasConstants>();

            // Set defaults
            _firstExecution = true;
            _resourceFrameIndex = 0;
            
            Constants.displaySize = _contextDescription.DisplaySize;
            
            _resources.Create(_contextDescription);
            CreatePipelines();
        }

        // private void InitShaders()
        // {
        //     LoadComputeShader("FSR2/ffx_fsr2_autogen_reactive_pass", ref _generateReactiveShader, out _generateReactiveKernel);
        // }

        private void CreatePipelines()
        {
            _computeLuminancePyramidPipeline = new Fsr2ComputeLuminancePyramidPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer, _spdConstantsBuffer);
            _reconstructPreviousDepthPipeline = new Fsr2ReconstructPreviousDepthPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _depthClipPipeline = new Fsr2DepthClipPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _lockPipeline = new Fsr2LockPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _accumulatePipeline = new Fsr2AccumulatePipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _accumulateSharpenPipeline = new Fsr2AccumulateSharpenPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _rcasPipeline = new Fsr2RcasPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer, _rcasConstantsBuffer);
        }
        
        public void Destroy()
        {
            DestroyPipeline(ref _tcrAutogeneratePipeline);
            DestroyPipeline(ref _generateReactivePipeline);
            DestroyPipeline(ref _computeLuminancePyramidPipeline);
            DestroyPipeline(ref _rcasPipeline);
            DestroyPipeline(ref _accumulateSharpenPipeline);
            DestroyPipeline(ref _accumulatePipeline);
            DestroyPipeline(ref _lockPipeline);
            DestroyPipeline(ref _reconstructPreviousDepthPipeline);
            DestroyPipeline(ref _depthClipPipeline);
            
            _resources.Destroy();
            
            DestroyConstantBuffer(ref _rcasConstantsBuffer);
            DestroyConstantBuffer(ref _spdConstantsBuffer);
            DestroyConstantBuffer(ref _fsr2ConstantsBuffer);

            _commandBuffer.Dispose();
            _commandBuffer = null;
        }

        public void Dispatch(Fsr2.DispatchDescription dispatchParams)
        {
            // TODO: validation & debug checking
            
            _commandBuffer.Clear();
            
            if (_firstExecution)
            {
                _commandBuffer.SetRenderTarget(_resources.LockStatus[0]);
                _commandBuffer.ClearRenderTarget(false, true, Color.clear);
                _commandBuffer.SetRenderTarget(_resources.LockStatus[1]);
                _commandBuffer.ClearRenderTarget(false, true, Color.clear);
            }
            
            int frameIndex = _resourceFrameIndex % 2;
            bool resetAccumulation = dispatchParams.Reset || _firstExecution;
            _firstExecution = false;

            // TODO: auto-exposure flag
            if (dispatchParams.Exposure == null) dispatchParams.Exposure = _resources.DefaultExposure;
            if (dispatchParams.Reactive == null) dispatchParams.Reactive = _resources.DefaultReactive;
            Fsr2Pipeline.RegisterResources(_commandBuffer, _contextDescription, dispatchParams);
            
            SetupConstants(dispatchParams, resetAccumulation);
            
            // Reactive mask bias
            const int threadGroupWorkRegionDim = 8;
            int dispatchSrcX = (Constants.renderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchSrcY = (Constants.renderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchDstX = (_contextDescription.DisplaySize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchDstY = (_contextDescription.DisplaySize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;

            // Clear reconstructed depth for max depth store.
            if (resetAccumulation)
            {
                _commandBuffer.SetRenderTarget(_resources.LockStatus[frameIndex ^ 1]);
                _commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                _commandBuffer.SetRenderTarget(_resources.InternalUpscaled[frameIndex ^ 1]);
                _commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                _commandBuffer.SetRenderTarget(_resources.SceneLuminance);
                _commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                // Auto exposure always used to track luma changes in locking logic
                _commandBuffer.SetRenderTarget(_resources.AutoExposure);
                _commandBuffer.ClearRenderTarget(false, true, new Color(-1f, 1e8f, 0f, 0f));
            }
            
            // Auto exposure
            SetupSpdConstants(dispatchParams, out var dispatchThreadGroupCount);
            
            // Initialize constant buffers data
            _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);
            _spdConstantsBuffer.SetData(_spdConstantsArray);
            
            // Compute luminance pyramid
            _computeLuminancePyramidPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, dispatchThreadGroupCount.x, dispatchThreadGroupCount.y);

            // Reconstruct previous depth
            _reconstructPreviousDepthPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            // Depth clip
            _depthClipPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            // Create locks
            _lockPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            bool sharpenEnabled = dispatchParams.EnableSharpening;
            
            // Accumulate
            var accumulatePipeline = sharpenEnabled ? _accumulateSharpenPipeline : _accumulatePipeline;
            accumulatePipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, dispatchDstX, dispatchDstY);

            if (sharpenEnabled)
            {
                // Compute the constants
                SetupRcasConstants(dispatchParams);
                _rcasConstantsBuffer.SetData(_rcasConstantsArray);
                
                // Dispatch RCAS
                const int threadGroupWorkRegionDimRcas = 16;
                int threadGroupsX = (Screen.width + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                int threadGroupsY = (Screen.height + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                _rcasPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, frameIndex, threadGroupsX, threadGroupsY);
            }

            _resourceFrameIndex = (_resourceFrameIndex + 1) % MaxQueuedFrames;

            Graphics.ExecuteCommandBuffer(_commandBuffer);
            
            Fsr2Pipeline.UnregisterResources(_commandBuffer);
        }

        private void SetupConstants(Fsr2.DispatchDescription dispatchParams, bool resetAccumulation)
        {
            ref Fsr2.Fsr2Constants constants = ref Constants;
            
            constants.jitterOffset = dispatchParams.JitterOffset;
            constants.renderSize = new Vector2Int(
                dispatchParams.RenderSize.x > 0 ? dispatchParams.RenderSize.x : dispatchParams.ColorDepth.width,
                dispatchParams.RenderSize.y > 0 ? dispatchParams.RenderSize.y : dispatchParams.ColorDepth.height);
            constants.maxRenderSize = _contextDescription.MaxRenderSize;
            constants.inputColorResourceDimensions = new Vector2Int(dispatchParams.ColorDepth.width, dispatchParams.ColorDepth.height);

            // Compute the horizontal FOV for the shader from the vertical one
            float aspectRatio = (float)dispatchParams.RenderSize.x / dispatchParams.RenderSize.y;
            float cameraAngleHorizontal = Mathf.Atan(Mathf.Tan(dispatchParams.CameraFovAngleVertical / 2.0f) * aspectRatio) * 2.0f;
            constants.tanHalfFOV = Mathf.Tan(cameraAngleHorizontal * 0.5f);
            constants.viewSpaceToMetersFactor = (dispatchParams.ViewSpaceToMetersFactor > 0.0f) ? dispatchParams.ViewSpaceToMetersFactor : 1.0f;

            // Compute params to enable device depth to view space depth computation in shader
            constants.deviceToViewDepth = SetupDeviceDepthToViewSpaceDepthParams(dispatchParams);
            
            // To be updated if resource is larger than the actual image size
            constants.downscaleFactor = new Vector2((float)constants.renderSize.x / _contextDescription.DisplaySize.x, (float)constants.renderSize.y / _contextDescription.DisplaySize.y);
            constants.previousFramePreExposure = constants.preExposure;
            constants.preExposure = (dispatchParams.PreExposure != 0) ? dispatchParams.PreExposure : 1.0f;
            
            // Motion vector data
            Vector2Int motionVectorsTargetSize = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDisplayResolutionMotionVectors) != 0 ? constants.displaySize : constants.renderSize;
            constants.motionVectorScale = dispatchParams.MotionVectorScale / motionVectorsTargetSize;
            
            // Compute jitter cancellation
            if ((_contextDescription.Flags & Fsr2.InitializationFlags.EnableMotionVectorsJitterCancellation) != 0)
            {
                constants.motionVectorJitterCancellation = (_previousJitterOffset - constants.jitterOffset) / motionVectorsTargetSize;
                _previousJitterOffset = constants.jitterOffset;
            }

            int jitterPhaseCount = Fsr2.GetJitterPhaseCount(dispatchParams.RenderSize.x, _contextDescription.DisplaySize.x);
            if (resetAccumulation || constants.jitterPhaseCount == 0)
            {
                constants.jitterPhaseCount = jitterPhaseCount;
            }
            else
            {
                int jitterPhaseCountDelta = (int)(jitterPhaseCount - constants.jitterPhaseCount);
                if (jitterPhaseCountDelta > 0)
                    constants.jitterPhaseCount++;
                else if (jitterPhaseCountDelta < 0)
                    constants.jitterPhaseCount--;
            }
            
            // Convert delta time to seconds and clamp to [0, 1]
            constants.deltaTime = Mathf.Clamp01(dispatchParams.FrameTimeDelta);

            if (resetAccumulation)
                constants.frameIndex = 0;
            else
                constants.frameIndex++;

            // Shading change usage of the SPD mip levels
            constants.lumaMipLevelToUse = Fsr2Pipeline.ShadingChangeMipLevel;

            float mipDiv = 2 << constants.lumaMipLevelToUse;
            constants.lumaMipDimensions.x = (int)(constants.maxRenderSize.x / mipDiv);
            constants.lumaMipDimensions.y = (int)(constants.maxRenderSize.y / mipDiv);
        }
        
        private Vector4 SetupDeviceDepthToViewSpaceDepthParams(Fsr2.DispatchDescription dispatchParams)
        {
            bool inverted = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDepthInverted) != 0;
            bool infinite = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDepthInfinite) != 0;

            // make sure it has no impact if near and far plane values are swapped in dispatch params
            // the flags "inverted" and "infinite" will decide what transform to use
            float min = Mathf.Min(dispatchParams.CameraNear, dispatchParams.CameraFar);
            float max = Mathf.Max(dispatchParams.CameraNear, dispatchParams.CameraFar);

            if (inverted)
            {
                (min, max) = (max, min);
            }

            float q = max / (min - max);
            float d = -1.0f;

            Vector4 matrixElemC = new Vector4(q, -1.0f - Mathf.Epsilon, q, 0.0f + Mathf.Epsilon);
            Vector4 matrixElemE = new Vector4(q * min, -min - Mathf.Epsilon, q * min, max);
            
            // Revert x and y coords
            float aspect = (float)dispatchParams.RenderSize.x / dispatchParams.RenderSize.y;
            float cotHalfFovY = Mathf.Cos(0.5f * dispatchParams.CameraFovAngleVertical) / Mathf.Sin(0.5f * dispatchParams.CameraFovAngleVertical);

            int matrixIndex = (inverted ? 2 : 0) + (infinite ? 1 : 0);
            return new Vector4(
                d * matrixElemC[matrixIndex],
                matrixElemE[matrixIndex],
                aspect / cotHalfFovY,
                1.0f / cotHalfFovY);
        }

        private void SetupRcasConstants(Fsr2.DispatchDescription dispatchParams)
        {
            int sharpnessIndex = Mathf.RoundToInt(Mathf.Clamp01(dispatchParams.Sharpness) * (RcasConfigs.Count - 1));
            RcasConsts = RcasConfigs[sharpnessIndex];
        }

        private void SetupSpdConstants(Fsr2.DispatchDescription dispatchParams, out Vector2Int dispatchThreadGroupCount)
        {
            RectInt rectInfo = new RectInt(0, 0, dispatchParams.RenderSize.x, dispatchParams.RenderSize.y);
            SpdSetup(rectInfo, out dispatchThreadGroupCount, out var workGroupOffset, out var numWorkGroupsAndMips);
            
            // Downsample
            ref Fsr2.SpdConstants spdConstants = ref SpdConsts;
            spdConstants.numWorkGroups = (uint)numWorkGroupsAndMips.x;
            spdConstants.mips = (uint)numWorkGroupsAndMips.y;
            spdConstants.workGroupOffsetX = (uint)workGroupOffset.x;
            spdConstants.workGroupOffsetY = (uint)workGroupOffset.y;
            spdConstants.renderSizeX = (uint)dispatchParams.RenderSize.x;
            spdConstants.renderSizeY = (uint)dispatchParams.RenderSize.y;
        }

        private static void SpdSetup(RectInt rectInfo, 
            out Vector2Int dispatchThreadGroupCount, out Vector2Int workGroupOffset, out Vector2Int numWorkGroupsAndMips, int mips = -1)
        {
            workGroupOffset = new Vector2Int(rectInfo.x / 64, rectInfo.y / 64);

            int endIndexX = (rectInfo.x + rectInfo.width - 1) / 64;
            int endIndexY = (rectInfo.y + rectInfo.height - 1) / 64;

            dispatchThreadGroupCount = new Vector2Int(endIndexX + 1 - workGroupOffset.x, endIndexY + 1 - workGroupOffset.y);

            numWorkGroupsAndMips = new Vector2Int(dispatchThreadGroupCount.x * dispatchThreadGroupCount.y, mips);
            if (mips < 0)
            {
                float resolution = Math.Max(rectInfo.width, rectInfo.height);
                numWorkGroupsAndMips.y = Math.Min(Mathf.FloorToInt(Mathf.Log(resolution, 2.0f)), 12);
            }
        }

        /// <summary>
        /// The FSR2 C++ codebase uses floats bitwise converted to ints to pass sharpness parameters to the RCAS shader.
        /// This is not possible in C# without enabling unsafe code compilation, so to avoid that we instead use a table of precomputed values.
        /// </summary>
        private static readonly List<Fsr2.RcasConstants> RcasConfigs = new()
        {
            new(1048576000u, 872428544u),
            new(1049178080u, 877212745u),
            new(1049823372u, 882390168u),
            new(1050514979u, 887895276u),
            new(1051256227u, 893859143u),
            new(1052050675u, 900216232u),
            new(1052902144u, 907032080u),
            new(1053814727u, 914306687u),
            new(1054792807u, 922105590u),
            new(1055841087u, 930494326u),
            new(1056964608u, 939538432u),
            new(1057566688u, 944322633u),
            new(1058211980u, 949500056u),
            new(1058903587u, 955005164u),
            new(1059644835u, 960969031u),
            new(1060439283u, 967326120u),
            new(1061290752u, 974141968u),
            new(1062203335u, 981416575u),
            new(1063181415u, 989215478u),
            new(1064229695u, 997604214u),
            new(1065353216u, 1006648320),
        };
        
        private static ComputeBuffer CreateConstantBuffer<TConstants>() where TConstants: struct
        {
            return new ComputeBuffer(1, Marshal.SizeOf<TConstants>(), ComputeBufferType.Constant);
        }
        
        private static void DestroyConstantBuffer(ref ComputeBuffer bufferRef)
        {
            if (bufferRef == null)
                return;
            
            bufferRef.Release();
            bufferRef = null;
        }

        private static void DestroyPipeline(ref Fsr2Pipeline pipeline)
        {
            if (pipeline == null)
                return;
            
            pipeline.Dispose();
            pipeline = null;
        }
    }
}