FSRTestBuiltin/Assets/Scripts/Fsr2Context.cs

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

namespace FidelityFX
{
    public class Fsr2Context
    {
        private Fsr2.ContextDescription _contextDescription;

        private ComputeShader _prepareInputColorShader;
        private ComputeShader _depthClipShader;
        private ComputeShader _reconstructPreviousDepthShader;
        private ComputeShader _lockShader;
        private ComputeShader _accumulateShader;
        private ComputeShader _generateReactiveShader;
        private ComputeShader _rcasShader;
        private ComputeShader _computeLuminancePyramidShader;
        
        private ComputeBuffer _fsr2ConstantsBuffer;
        private readonly Fsr2Constants[] _fsr2ConstantsArray = { new Fsr2Constants() };

        private ComputeBuffer _spdConstantsBuffer;
        private readonly SpdConstants[] _spdConstantsArray = { new SpdConstants() };
    
        private ComputeBuffer _rcasConstantsBuffer;
        private readonly RcasConstants[] _rcasConstantsArray = new RcasConstants[1];

        public void Create(Fsr2.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;
            
            _fsr2ConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<Fsr2Constants>(), ComputeBufferType.Constant);
            _spdConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<SpdConstants>(), ComputeBufferType.Constant);
            _rcasConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<RcasConstants>(), ComputeBufferType.Constant);

            // Set defaults
            _fsr2ConstantsArray[0].displaySize = _contextDescription.DisplaySize;
            _fsr2ConstantsArray[0].displaySizeRcp = new Vector2(
                1.0f / _contextDescription.DisplaySize.x,
                1.0f / _contextDescription.DisplaySize.y);
            
            // 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);
            }
            
            InitShaders();
        }

        private void InitShaders()
        {
            LoadComputeShader("FSR2/ffx_fsr2_compute_luminance_pyramid_pass", ref _computeLuminancePyramidShader);
            LoadComputeShader("FSR2/ffx_fsr2_rcas_pass", ref _rcasShader);
            LoadComputeShader("FSR2/ffx_fsr2_prepare_input_color_pass", ref _prepareInputColorShader);
            LoadComputeShader("FSR2/ffx_fsr2_depth_clip_pass", ref _depthClipShader);
            LoadComputeShader("FSR2/ffx_fsr2_reconstruct_previous_depth_pass", ref _reconstructPreviousDepthShader);
            LoadComputeShader("FSR2/ffx_fsr2_lock_pass", ref _lockShader);
            LoadComputeShader("FSR2/ffx_fsr2_accumulate_pass", ref _accumulateShader);
            LoadComputeShader("FSR2/ffx_fsr2_autogen_reactive_pass", ref _generateReactiveShader);
        }

        private void LoadComputeShader(string name, ref ComputeShader shaderRef)
        {
            if (shaderRef == null)
                shaderRef = _contextDescription.Callbacks.LoadComputeShader(name);
        }
        
        public void Dispatch(Fsr2.DispatchDescription dispatchDescription)
        {
            _fsr2ConstantsArray[0].preExposure = dispatchDescription.PreExposure;
            _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);

            if (dispatchDescription.EnableSharpening)
            {
                int sharpnessIndex = Mathf.RoundToInt(Mathf.Clamp01(dispatchDescription.Sharpness) * (RcasConfigs.Count - 1));
                _rcasConstantsArray[0] = RcasConfigs[sharpnessIndex];
                _rcasConstantsBuffer.SetData(_rcasConstantsArray);
            
                // Run the RCAS sharpening filter on the upscaled image
                int rcasKernel = _rcasShader.FindKernel("CS");
                _rcasShader.SetTexture(rcasKernel, "r_exposure", dispatchDescription.Exposure);
                _rcasShader.SetTexture(rcasKernel, "r_rcas_input", dispatchDescription.Input);
                _rcasShader.SetTexture(rcasKernel, "rw_upscaled_output", dispatchDescription.Output);
                _rcasShader.SetConstantBuffer("cbFSR2", _fsr2ConstantsBuffer, 0, Marshal.SizeOf<Fsr2Constants>());
                _rcasShader.SetConstantBuffer("cbRCAS", _rcasConstantsBuffer, 0, Marshal.SizeOf<RcasConstants>());

                const int threadGroupWorkRegionDimRcas = 16;
                int threadGroupsX = (Screen.width + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                int threadGroupsY = (Screen.height + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
            
                _rcasShader.Dispatch(rcasKernel, threadGroupsX, threadGroupsY, 1);
            }
            else
            {
                Graphics.Blit(dispatchDescription.Input, dispatchDescription.Output);
            }
        }

        public void Destroy()
        {
            if (_rcasConstantsBuffer != null)
            {
                _rcasConstantsBuffer.Release();
                _rcasConstantsBuffer = null;
            }

            if (_spdConstantsBuffer != null)
            {
                _spdConstantsBuffer.Release();
                _spdConstantsBuffer = null;
            }

            if (_fsr2ConstantsBuffer != null)
            {
                _fsr2ConstantsBuffer.Release();
                _fsr2ConstantsBuffer = null;
            }
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct Fsr2Constants
        {
            public Vector2Int renderSize;
            public Vector2Int displaySize;
            public uint lumaMipDimensionsX, lumaMipDimensionsY;
            public uint lumaMipLevelToUse;
            public uint frameIndex;
            public Vector2 displaySizeRcp;
            public Vector2 jitterOffset;
            public Vector4 deviceToViewDepth;
            public Vector2 depthClipUVScale;
            public Vector2 postLockStatusUVScale;
            public Vector2 reactiveMaskDimRcp;
            public Vector2 motionVectorScale;
            public Vector2 downscaleFactor;
            public float preExposure;
            public float tanHalfFOV;
            public Vector2 motionVectorJitterCancellation;
            public float jitterPhaseCount;
            public float lockInitialLifetime;
            public float lockTickDelta;
            public float deltaTime;
            public float dynamicResChangeFactor;
            public float lumaMipRcp;
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct SpdConstants
        {
            public uint mips;
            public uint numWorkGroups;
            public uint workGroupOffsetX, workGroupOffsetY;
            public uint renderSizeX, renderSizeY;
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct RcasConstants
        {
            public RcasConstants(uint sharpness, uint halfSharp)
            {
                this.sharpness = sharpness;
                this.halfSharp = halfSharp;
                dummy0 = dummy1 = 0;
            }
        
            public readonly uint sharpness;
            public readonly uint halfSharp;
            public readonly uint dummy0;
            public readonly uint dummy1;
        }

        private static readonly List<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),
        };
    }
}