FSRTestBuiltin/Assets/Scripts/Fsr2Pipeline.cs

using System;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;

namespace FidelityFX
{
    internal abstract class Fsr2Pipeline: IDisposable
    {
        private readonly Fsr2Callbacks _callbacks;
        protected readonly ComputeBuffer Constants;
        
        protected ComputeShader ComputeShader;
        protected int KernelIndex;
        
        // Shader resource views, i.e. read-only bindings
        protected static readonly int SrvInputColor = Shader.PropertyToID("r_input_color_jittered");
        protected static readonly int SrvOpaqueOnly = Shader.PropertyToID("r_input_opaque_only");
        protected static readonly int SrvInputMotionVectors = Shader.PropertyToID("r_input_motion_vectors");
        protected static readonly int SrvInputDepth = Shader.PropertyToID("r_input_depth");
        protected static readonly int SrvInputExposure = Shader.PropertyToID("r_input_exposure");
        protected static readonly int SrvRcasInput = Shader.PropertyToID("r_rcas_input");

        // Unordered access views, i.e. random read/write bindings
        protected static readonly int UavUpscaledOutput = Shader.PropertyToID("rw_upscaled_output");
        protected static readonly int UavExposureMipLumaChange = Shader.PropertyToID("rw_img_mip_shading_change");
        protected static readonly int UavExposureMip5 = Shader.PropertyToID("rw_img_mip_5");
        protected static readonly int UavAutoExposure = Shader.PropertyToID("rw_auto_exposure");
        protected static readonly int UavSpdAtomicCount = Shader.PropertyToID("rw_spd_global_atomic");

        // Constant buffer bindings
        protected static readonly int CbFsr2 = Shader.PropertyToID("cbFSR2");
        protected static readonly int CbSpd = Shader.PropertyToID("cbSPD");
        protected static readonly int CbRcas = Shader.PropertyToID("cbRCAS");
        protected static readonly int CbGenReactive = Shader.PropertyToID("cbGenerateReactive");

        protected Fsr2Pipeline(Fsr2Callbacks callbacks, ComputeBuffer constants)
        {
            _callbacks = callbacks;
            Constants = constants;
        }

        public virtual void Dispose()
        {
            UnloadComputeShader();
        }
        
        public abstract void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY);
        
        protected void LoadComputeShader(string name, Fsr2.InitializationFlags flags)
        {
            LoadComputeShader(name, flags, ref ComputeShader, out KernelIndex);
        }
        
        private void LoadComputeShader(string name, Fsr2.InitializationFlags flags, ref ComputeShader shaderRef, out int kernelIndex)
        {
            if (shaderRef == null)
                shaderRef = _callbacks.LoadComputeShader(name);

            kernelIndex = shaderRef.FindKernel("CS");

            bool useLut = (SystemInfo.computeSubGroupSize == 64);

            // This mirrors the permutation rules from the CreatePipeline* functions
            if ((flags & Fsr2.InitializationFlags.EnableHighDynamicRange) != 0) shaderRef.EnableKeyword("FFX_FSR2_OPTION_HDR_COLOR_INPUT");
            if ((flags & Fsr2.InitializationFlags.EnableDisplayResolutionMotionVectors) == 0) shaderRef.EnableKeyword("FFX_FSR2_OPTION_LOW_RESOLUTION_MOTION_VECTORS");
            if ((flags & Fsr2.InitializationFlags.EnableMotionVectorsJitterCancellation) != 0) shaderRef.EnableKeyword("FFX_FSR2_OPTION_JITTERED_MOTION_VECTORS");
            if ((flags & Fsr2.InitializationFlags.EnableDepthInverted) != 0) shaderRef.EnableKeyword("FFX_FSR2_OPTION_INVERTED_DEPTH");
            if (useLut) shaderRef.EnableKeyword("FFX_FSR2_OPTION_REPROJECT_USE_LANCZOS_TYPE");
            // TODO: enable FFX_HALF if FP16 is supported (except RCAS)
        }

        private void UnloadComputeShader()
        {
            UnloadComputeShader(ref ComputeShader);
        }
        
        private void UnloadComputeShader(ref ComputeShader shaderRef)
        {
            if (shaderRef == null)
                return;

            _callbacks.UnloadComputeShader(shaderRef);
            shaderRef = null;
        }
    }

    internal class Fsr2ComputeLuminancePyramidPipeline : Fsr2Pipeline
    {
        private readonly ComputeBuffer _spdConstants;
        private readonly RenderTexture _autoExposure;
        
        public Fsr2ComputeLuminancePyramidPipeline(Fsr2Callbacks callbacks, Fsr2.InitializationFlags flags, ComputeBuffer constants, ComputeBuffer spdConstants, RenderTexture autoExposure)
            : base(callbacks, constants)
        {
            _spdConstants = spdConstants;
            _autoExposure = autoExposure;
            
            LoadComputeShader("FSR2/ffx_fsr2_compute_luminance_pyramid_pass", flags);
        }

        public override void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY)
        {
            // Problems to solve:
            // - How do resources (render textures) relate to SRV/UAV bindings? How are those tied together?
            // - What about the SRV/UAVs that are not related to any resources? Where are those filled in?
            // - How do we clear the resources that need to be cleared at dispatch? (SetBufferData)
            // - Shouldn't we use a ComputeBuffer for resources that are one-dimensional and clearly not image data? e.g. SPD atomic counter & Lanczos LUT data

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, SrvInputColor, dispatchParams.Input);
            // Resource FSR2_SpdAtomicCounter: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R32_UINT, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(UavSpdAtomicCount, 1, 1, 0, FilterMode.Point, GraphicsFormat.R32_UInt, 1, true);   // FSR2_BIND_UAV_SPD_GLOBAL_ATOMIC
            // Resource FSR2_ExposureMips: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16_FLOAT, FFX_RESOURCE_FLAGS_ALIASABLE, mipCount = 0
            // See `scheduleDispatch` for the song and dance to bind UAV mip levels to each luminance mipmap... and this shader specifically wants mip levels 4 and 5
            // Looks like we can just bind two separate resources here, shouldn't be necessary to bother with mipmapping nonsense. Be sure to get the right dimensions though.
            commandBuffer.GetTemporaryRT(UavExposureMipLumaChange, dispatchParams.RenderSize.x >> 4, dispatchParams.RenderSize.y >> 4, 0, FilterMode.Point, GraphicsFormat.R16_SFloat, 1, true);    // FSR2_BIND_UAV_EXPOSURE_MIP_LUMA_CHANGE
            commandBuffer.GetTemporaryRT(UavExposureMip5, dispatchParams.RenderSize.x >> 5, dispatchParams.RenderSize.y >> 5, 0, FilterMode.Point, GraphicsFormat.R16_SFloat, 1, true);    // FSR2_BIND_UAV_EXPOSURE_MIP_5
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, UavAutoExposure, _autoExposure);   // FSR2_BIND_UAV_AUTO_EXPOSURE
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CbFsr2, Constants, 0, Marshal.SizeOf<Fsr2.Fsr2Constants>());
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CbSpd, _spdConstants, 0, Marshal.SizeOf<Fsr2.SpdConstants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, 1);
            
            // NOTE: since these temp RTs are not bound to a specific shader or kernel, we can set them globally one time and release them after dispatch.
            // That way we can share aliasable resources between shaders without any complicated management. 
            commandBuffer.ReleaseTemporaryRT(UavSpdAtomicCount);
            commandBuffer.ReleaseTemporaryRT(UavExposureMipLumaChange);
            commandBuffer.ReleaseTemporaryRT(UavExposureMip5);
        }
    }

    internal class Fsr2AccumulatePipeline : Fsr2Pipeline
    {
        public Fsr2AccumulatePipeline(Fsr2Callbacks callbacks, Fsr2.InitializationFlags flags, ComputeBuffer constants)
            : base(callbacks, constants)
        {
            LoadComputeShader("FSR2/ffx_fsr2_accumulate_pass", flags);
        }

        public override void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY)
        {
            //throw new NotImplementedException();
        }
    }

    internal class Fsr2AccumulateSharpenPipeline : Fsr2AccumulatePipeline
    {
        private readonly ComputeShader _shaderCopy;
        
        public Fsr2AccumulateSharpenPipeline(Fsr2Callbacks callbacks, Fsr2.InitializationFlags flags, ComputeBuffer constants)
            : base(callbacks, flags, constants)
        {
            // Simply loading the accumulate_pass compute shader will give us the same instance as the non-sharpen pipeline
            // So we have to clone the shader instance and set the extra keyword on the new copy
            _shaderCopy = UnityEngine.Object.Instantiate(ComputeShader);
            foreach (var keyword in ComputeShader.enabledKeywords)
            {
                _shaderCopy.EnableKeyword(keyword.name);
            }
            _shaderCopy.EnableKeyword("FFX_FSR2_OPTION_APPLY_SHARPENING");
        }

        public override void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY)
        {
            // Temporarily swap around the shaders so that the dispatch will bind and execute the correct one
            ComputeShader tmp = ComputeShader;
            ComputeShader = _shaderCopy;
            base.ScheduleDispatch(commandBuffer, dispatchParams, dispatchX, dispatchY);
            ComputeShader = tmp;
        }

        public override void Dispose()
        {
            // Since we instantiated this copy, we have to destroy it instead of unloading the shader resource
            UnityEngine.Object.Destroy(_shaderCopy);
            base.Dispose();
        }
    }
    
    internal class Fsr2RcasPipeline : Fsr2Pipeline
    {
        private readonly ComputeBuffer _rcasConstants;

        public Fsr2RcasPipeline(Fsr2Callbacks callbacks, Fsr2.InitializationFlags flags, ComputeBuffer constants, ComputeBuffer rcasConstants)
            : base(callbacks, constants)
        {
            _rcasConstants = rcasConstants;
            
            LoadComputeShader("FSR2/ffx_fsr2_rcas_pass", flags);
        }

        public override void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY)
        {
            // Run the RCAS sharpening filter on the upscaled image
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, SrvInputExposure, dispatchParams.Exposure);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, SrvRcasInput, dispatchParams.Input);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, UavUpscaledOutput, dispatchParams.Output);
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CbFsr2, Constants, 0, Marshal.SizeOf<Fsr2.Fsr2Constants>());
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CbRcas, _rcasConstants, 0, Marshal.SizeOf<Fsr2.RcasConstants>());

            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, 1);
        }
    }
}