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 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 uint _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;
            
            // Generate the data for the LUT
            const uint lanczos2LutWidth = 128;
            short[] lanczos2Weights = new short[lanczos2LutWidth];
            for (uint currentLanczosWidthIndex = 0; currentLanczosWidthIndex < lanczos2LutWidth; ++currentLanczosWidthIndex)
            {
                float x = 2.0f * currentLanczosWidthIndex / (lanczos2LutWidth - 1);
                float y = Fsr2.Lanczos2(x);
                lanczos2Weights[currentLanczosWidthIndex] = (short)Mathf.Round(y * 32767.0f);
            }
            
            // TODO: create resources, i.e. render textures used for intermediate results.
            // Note that "aliasable" resources should be equivalent to GetTemporary render textures
            // UAVs *may* be an issue with the PS4 not handling simultaneous reading and writing to an RT properly
            // Unity does have Graphics.SetRandomWriteTarget for enabling UAV on ComputeBuffers or RTs
            // Unity doesn't do 1D textures so just default to Texture2D
            
            CreatePipelines();
        }

        // private void InitShaders()
        // {
        //     LoadComputeShader("FSR2/ffx_fsr2_compute_luminance_pyramid_pass", ref _computeLuminancePyramidShader, out _computeLuminancePyramidKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_rcas_pass", ref _rcasShader, out _rcasKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_prepare_input_color_pass", ref _prepareInputColorShader, out _prepareInputColorKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_depth_clip_pass", ref _depthClipShader, out _depthClipKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_reconstruct_previous_depth_pass", ref _reconstructPreviousDepthShader, out _reconstructPreviousDepthKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_lock_pass", ref _lockShader, out _lockKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_accumulate_pass", ref _accumulateShader, out _accumulateKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_autogen_reactive_pass", ref _generateReactiveShader, out _generateReactiveKernel);
        //     LoadComputeShader("FSR2/ffx_fsr2_tcr_autogen_pass", ref _tcrAutogenShader, out _tcrAutogenKernel);
        // }

        private void CreatePipelines()
        {
            _computeLuminancePyramidPipeline = new Fsr2ComputeLuminancePyramidPipeline(_contextDescription.Callbacks, _contextDescription.Flags, _fsr2ConstantsBuffer, _spdConstantsBuffer);
            _accumulatePipeline = new Fsr2AccumulatePipeline(_contextDescription.Callbacks, _contextDescription.Flags, _fsr2ConstantsBuffer);
            _accumulateSharpenPipeline = new Fsr2AccumulateSharpenPipeline(_contextDescription.Callbacks, _contextDescription.Flags, _fsr2ConstantsBuffer);
            _rcasPipeline = new Fsr2RcasPipeline(_contextDescription.Callbacks, _contextDescription.Flags, _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);
            
            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();
            
            // TODO: Should probably use a CommandBuffer here, to queue up all of the commands and dispatch them in one go
            // Use SetRandomWriteTarget to declare UAVs
            // Use ClearRandomWriteTargets to clear all UAVs at once
            
            if (_firstExecution)
            {
                // TODO: clear values
            }
            
            // TODO: setup resource indices for buffers that get swapped per frame
            
            bool resetAccumulation = dispatchParams.Reset || _firstExecution;
            _firstExecution = false;
            
            // TODO Register resources: ...
            
            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;

            if (resetAccumulation)
            {
                // TODO: clear reconstructed depth for max depth store
                
            }
            
            // Auto exposure
            SetupSpdConstants(dispatchParams, out var dispatchThreadGroupCount);
            
            // Initialize constant buffers data
            _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);
            _spdConstantsBuffer.SetData(_spdConstantsArray);
            
            // Compute luminance pyramid
            _computeLuminancePyramidPipeline.ScheduleDispatch(_commandBuffer, dispatchParams, dispatchThreadGroupCount.x, dispatchThreadGroupCount.y);

            // // Reconstruct previous depth
            // _commandBuffer.DispatchCompute(_reconstructPreviousDepthShader, _reconstructPreviousDepthKernel, dispatchSrcX, dispatchSrcY, 1);
            //
            // // Depth clip
            // _commandBuffer.DispatchCompute(_depthClipShader, _depthClipKernel, dispatchSrcX, dispatchSrcY, 1);
            //
            // // Lock
            // _commandBuffer.DispatchCompute(_lockShader, _lockKernel, dispatchSrcX, dispatchSrcY, 1);
            //

            bool sharpenEnabled = dispatchParams.EnableSharpening;
            
            // Accumulate
            var accumulatePipeline = sharpenEnabled ? _accumulateSharpenPipeline : _accumulatePipeline;
            accumulatePipeline.ScheduleDispatch(_commandBuffer, dispatchParams, 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, threadGroupsX, threadGroupsY);
            }
            else
            {
                _commandBuffer.Blit(dispatchParams.Input, dispatchParams.Output);
            }

            _resourceFrameIndex = (_resourceFrameIndex + 1) % MaxQueuedFrames;

            Graphics.ExecuteCommandBuffer(_commandBuffer);
            
            // TODO Unregister resources: release temp RT's
        }

        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.Input.width,
                dispatchParams.RenderSize.y > 0 ? dispatchParams.RenderSize.y : dispatchParams.Input.height);
            constants.maxRenderSize = _contextDescription.MaxRenderSize;
            constants.inputColorResourceDimensions =
                new Vector2Int(dispatchParams.Input.width, dispatchParams.Input.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 / 1000.0f);

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

            // Shading change usage of the SPD mip levels
            constants.lumaMipLevelToUse = 4;    // NOTE: this is derived from a bunch of auto-generated constant values in the FSR2 code

            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;
        }
    }
}