FSRPackage/Runtime/FrameInterpolation/FrameInterpolationContext.cs

using System;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Assertions;
using UnityEngine.Profiling;
using UnityEngine.Rendering;

namespace FidelityFX.FrameGen
{
    public class FrameInterpolationContext
    {
        private FrameInterpolation.ContextDescription _contextDescription;

        private FrameInterpolationPass _reconstructAndDilatePass;
        private FrameInterpolationPass _setupPass;
        private FrameInterpolationPass _reconstructPreviousDepthPass;
        private FrameInterpolationPass _gameMotionVectorFieldPass;
        private FrameInterpolationPass _opticalFlowVectorFieldPass;
        private FrameInterpolationPass _disocclusionMaskPass;
        private FrameInterpolationPass _interpolationPass;
        private FrameInterpolationPass _inpaintingPyramidPass;
        private FrameInterpolationPass _inpaintingPass;
        private FrameInterpolationPass _gameVectorFieldInpaintingPyramidPass;
        private FrameInterpolationPass _debugViewPass;
        
        private readonly FrameInterpolationResources _resources = new FrameInterpolationResources();
        
        private ComputeBuffer _frameInterpolationConstantsBuffer;
        private readonly FrameInterpolation.Constants[] _frameInterpolationConstantsArray = { new FrameInterpolation.Constants() };
        private ref FrameInterpolation.Constants Constants => ref _frameInterpolationConstantsArray[0];
        
        private ComputeBuffer _spdConstantsBuffer;
        private readonly FrameInterpolation.InpaintingPyramidConstants[] _spdConstantsArray = { new FrameInterpolation.InpaintingPyramidConstants() };
        private ref FrameInterpolation.InpaintingPyramidConstants SpdConstants => ref _spdConstantsArray[0];
        
        private readonly CustomSampler _sampler = CustomSampler.Create("Frame Interpolation");
        private readonly CustomSampler _prepareSampler = CustomSampler.Create("Frame Interpolation - Prepare");

        private bool _firstExecution;
        private bool _asyncSupported;
        private ulong _previousFrameID;
        private ulong _dispatchCount;
        
        public void Create(in FrameInterpolation.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;

            _frameInterpolationConstantsBuffer = CreateConstantBuffer<FrameInterpolation.Constants>();
            _spdConstantsBuffer = CreateConstantBuffer<FrameInterpolation.InpaintingPyramidConstants>();

            _firstExecution = true;
            _dispatchCount = 0;
            _previousFrameID = 0;
            _asyncSupported = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableAsyncSupport) == FrameInterpolation.InitializationFlags.EnableAsyncSupport;

            Constants.maxRenderSize = _contextDescription.maxRenderSize;
            Constants.displaySize = _contextDescription.displaySize;
            Constants.displaySizeRcp.x = 1.0f / _contextDescription.displaySize.x;
            Constants.displaySizeRcp.y = 1.0f / _contextDescription.displaySize.y;
            Constants.interpolationRectBase = Vector2Int.zero;
            Constants.interpolationRectSize = _contextDescription.displaySize;
            
            _resources.Create(_contextDescription);
            CreatePasses();
        }

        private void CreatePasses()
        {
            _reconstructAndDilatePass = new FrameInterpolationReconstructAndDilatePass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _setupPass = new FrameInterpolationSetupPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _reconstructPreviousDepthPass = new FrameInterpolationReconstructPreviousDepthPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _gameMotionVectorFieldPass = new FrameInterpolationGameMotionVectorFieldPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _opticalFlowVectorFieldPass = new FrameInterpolationOpticalFlowVectorFieldPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _disocclusionMaskPass = new FrameInterpolationDisocclusionMaskPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _interpolationPass = new FrameInterpolationInterpolationPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _inpaintingPyramidPass = new FrameInterpolationInpaintingPyramidPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer, _spdConstantsBuffer);
            _inpaintingPass = new FrameInterpolationInpaintingPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _gameVectorFieldInpaintingPyramidPass = new FrameInterpolationGameVectorFieldInpaintingPyramidPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer, _spdConstantsBuffer);
            _debugViewPass = new FrameInterpolationDebugViewPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
        }

        public void Destroy()
        {
            DestroyPass(ref _debugViewPass);
            DestroyPass(ref _gameVectorFieldInpaintingPyramidPass);
            DestroyPass(ref _inpaintingPass);
            DestroyPass(ref _inpaintingPyramidPass);
            DestroyPass(ref _interpolationPass);
            DestroyPass(ref _disocclusionMaskPass);
            DestroyPass(ref _opticalFlowVectorFieldPass);
            DestroyPass(ref _gameMotionVectorFieldPass);
            DestroyPass(ref _reconstructPreviousDepthPass);
            DestroyPass(ref _setupPass);
            DestroyPass(ref _reconstructAndDilatePass);
            
            _resources.Destroy();
            
            DestroyConstantBuffer(ref _spdConstantsBuffer);
            DestroyConstantBuffer(ref _frameInterpolationConstantsBuffer);
        }

        public void Prepare(CommandBuffer commandBuffer, FrameInterpolation.PrepareDescription prepareDescription)
        {
            commandBuffer.BeginSample(_prepareSampler);

            int doubleBufferId = _asyncSupported ? (int)(prepareDescription.frameID & 1) : 0;

            Constants.renderSize = prepareDescription.renderSize;
            Constants.jitter = prepareDescription.jitterOffset;
            Vector2Int motionVectorsTargetSize = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableDisplayResolutionMotionVectors) != 0 ? Constants.displaySize : Constants.renderSize;
            Constants.motionVectorScale.x = prepareDescription.motionVectorScale.x / motionVectorsTargetSize.x;
            Constants.motionVectorScale.y = prepareDescription.motionVectorScale.y / motionVectorsTargetSize.y;
            
            commandBuffer.SetBufferData(_frameInterpolationConstantsBuffer, _frameInterpolationConstantsArray);
            
            Assert.IsTrue(prepareDescription.depth.IsValid);
            Assert.IsTrue(prepareDescription.motionVectors.IsValid);
            
            // clear estimated depth resources
            {
                bool inverted = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableDepthInverted) == FrameInterpolation.InitializationFlags.EnableDepthInverted;
                commandBuffer.SetRenderTarget(_resources.ReconstructedDepth[doubleBufferId]);
                commandBuffer.ClearRenderTarget(false, true, inverted ? Color.clear : Color.white);
            }

            int renderDispatchSizeX = (prepareDescription.renderSize.x + 7) / 8;
            int renderDispatchSizeY = (prepareDescription.renderSize.y + 7) / 8;
            ((FrameInterpolationReconstructAndDilatePass)_reconstructAndDilatePass).ScheduleDispatch(commandBuffer, prepareDescription, doubleBufferId, renderDispatchSizeX, renderDispatchSizeY);
            
            commandBuffer.EndSample(_prepareSampler);
        }
        
        public void Dispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchDescription)
        {
            commandBuffer.BeginSample(_sampler);

            bool reset = _dispatchCount == 0 || dispatchDescription.reset;
            Assert.IsTrue(!_asyncSupported || reset || dispatchDescription.frameID > _previousFrameID,
                "When async support is enabled, and the reset flag is not set, frame ID must increment in each dispatch.");

            bool frameIdDecreased = dispatchDescription.frameID < _previousFrameID;
            bool frameIdSkipped = (dispatchDescription.frameID - _previousFrameID) > 1;
            bool disjointFrameId = frameIdDecreased || frameIdSkipped;
            _previousFrameID = dispatchDescription.frameID;
            _dispatchCount++;   // TODO: this is pointless, it does the same as _firstExecution, no need to do any counting

            Constants.renderSize = dispatchDescription.renderSize;
            Constants.displaySize = dispatchDescription.displaySize;
            Constants.displaySizeRcp.x = 1.0f / dispatchDescription.displaySize.x;
            Constants.displaySizeRcp.y = 1.0f / dispatchDescription.displaySize.y;
            Constants.upscalerTargetSize = dispatchDescription.interpolationRect.size;
            Constants.mode = 0;
            Constants.reset = (reset || disjointFrameId) ? 1 : 0;
            Constants.deltaTime = dispatchDescription.frameTimeDelta;
            Constants.HUDLessAttachedFactor = dispatchDescription.currentBackBuffer_HUDLess.IsValid ? 1 : 0;

            Constants.opticalFlowScale = dispatchDescription.opticalFlowScale;
            Constants.opticalFlowBlockSize = dispatchDescription.opticalFlowBlockSize;
            Constants.dispatchFlags = (uint)dispatchDescription.flags;

            Constants.cameraNear = dispatchDescription.cameraNear;
            Constants.cameraFar = dispatchDescription.cameraFar;

            Constants.interpolationRectBase = dispatchDescription.interpolationRect.position;
            Constants.interpolationRectSize = dispatchDescription.interpolationRect.size;

            // Debug bar
            Constants.debugBarColor.x = DebugBarColorSequence[_debugIndex * 3 + 0];
            Constants.debugBarColor.y = DebugBarColorSequence[_debugIndex * 3 + 1];
            Constants.debugBarColor.z = DebugBarColorSequence[_debugIndex * 3 + 2];
            _debugIndex = (_debugIndex + 1) % (DebugBarColorSequence.Length / 3);

            Constants.backBufferTransferFunction = (uint)dispatchDescription.backbufferTransferFunction;
            Constants.minMaxLuminance = dispatchDescription.minMaxLuminance;

            float aspectRatio = dispatchDescription.renderSize.x / (float)dispatchDescription.renderSize.y;
            float cameraAngleHorizontal = Mathf.Atan(Mathf.Tan(dispatchDescription.cameraFovAngleVertical / 2) * aspectRatio) * 2;
            Constants.tanHalfFOV = Mathf.Tan(cameraAngleHorizontal * 0.5f);
            Constants.deviceToViewDepth = SetupDeviceDepthToViewSpaceDepthParams(dispatchDescription);
            
            commandBuffer.SetBufferData(_frameInterpolationConstantsBuffer, _frameInterpolationConstantsArray);
            
            int doubleBufferId = _asyncSupported ? (int)(dispatchDescription.frameID & 1) : 0;
            
            int displayDispatchSizeX = (dispatchDescription.displaySize.x + 7) / 8;
            int displayDispatchSizeY = (dispatchDescription.displaySize.y + 7) / 8;

            int renderDispatchSizeX = (dispatchDescription.renderSize.x + 7) / 8;
            int renderDispatchSizeY = (dispatchDescription.renderSize.y + 7) / 8;

            int opticalFlowDispatchSizeX = (int)(dispatchDescription.displaySize.x / (float)dispatchDescription.opticalFlowBlockSize + 7) / 8;
            int opticalFlowDispatchSizeY = (int)(dispatchDescription.displaySize.y / (float)dispatchDescription.opticalFlowBlockSize + 7) / 8;

            bool executePreparationPasses = (Constants.reset == 0);

            // TODO: do we have to clear some buffers somewhere on reset? Probably uninitialized RTs are kept to NULL in FFX source... we might want to clear them. Specifically the inputs for InterpolationPass.
            
            // Schedule work for the interpolation command list
            _setupPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, renderDispatchSizeX, renderDispatchSizeY);

            // TODO delegates are reference types, i.e. this allocates some memory every frame
            Action dispatchGameVectorFieldInpaintingPyramid = () =>
            {
                SetupSpdConstants(dispatchDescription.renderSize, out var dispatchThreadGroupCount);
                commandBuffer.SetBufferData(_spdConstantsBuffer, _spdConstantsArray);
                _gameVectorFieldInpaintingPyramidPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, dispatchThreadGroupCount.x, dispatchThreadGroupCount.y);
            };
            
            // Only execute FG data preparation passes when reset wasn't triggered
            if (executePreparationPasses)
            {
                // Clear estimated depth resources
                bool inverted = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableDepthInverted) == FrameInterpolation.InitializationFlags.EnableDepthInverted;
                commandBuffer.SetRenderTarget(_resources.ReconstructedDepthInterpolatedFrame);
                commandBuffer.ClearRenderTarget(false, true, inverted ? Color.clear : Color.white);
                
                _reconstructPreviousDepthPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, renderDispatchSizeX, renderDispatchSizeY);
                _gameMotionVectorFieldPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, renderDispatchSizeX, renderDispatchSizeY);
                dispatchGameVectorFieldInpaintingPyramid();
                _opticalFlowVectorFieldPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, opticalFlowDispatchSizeX, opticalFlowDispatchSizeY);
                _disocclusionMaskPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, renderDispatchSizeX, renderDispatchSizeY);
            }
            
            _interpolationPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, displayDispatchSizeX, displayDispatchSizeY);
            
            // Inpainting pyramid
            SetupSpdConstants(dispatchDescription.displaySize, out var dispatchThreadGroupCount);
            commandBuffer.SetBufferData(_spdConstantsBuffer, _spdConstantsArray);
            _inpaintingPyramidPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, dispatchThreadGroupCount.x, dispatchThreadGroupCount.y);
            
            _inpaintingPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, displayDispatchSizeX, displayDispatchSizeY);

            if ((dispatchDescription.flags & FrameInterpolation.DispatchFlags.DrawDebugView) != 0)
            {
                dispatchGameVectorFieldInpaintingPyramid();
                _debugViewPass.ScheduleDispatch(commandBuffer, dispatchDescription, doubleBufferId, displayDispatchSizeX, displayDispatchSizeY);
            }
            
            // Store current buffer
            commandBuffer.CopyTexture(FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, _resources.PreviousInterpolationSource);
            
            commandBuffer.EndSample(_sampler);
        }
        
        private Vector4 SetupDeviceDepthToViewSpaceDepthParams(FrameInterpolation.DispatchDescription dispatchParams)
        {
            bool inverted = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableDepthInverted) != 0;
            bool infinite = (_contextDescription.flags & FrameInterpolation.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 SetupSpdConstants(Vector2Int resolution, out Vector2Int dispatchThreadGroupCount)
        {
            const int resolutionMultiplier = 1;
            
            RectInt rectInfo = new RectInt(0, 0, resolution.x * resolutionMultiplier, resolution.y * resolutionMultiplier);
            SpdSetup(rectInfo, out dispatchThreadGroupCount, out var workGroupOffset, out var numWorkGroupsAndMips, 4);
            
            ref FrameInterpolation.InpaintingPyramidConstants spdConstants = ref SpdConstants;
            spdConstants.numWorkGroups = (uint)numWorkGroupsAndMips.x;
            spdConstants.mips = (uint)numWorkGroupsAndMips.y;
            spdConstants.workGroupOffsetX = (uint)workGroupOffset.x;
            spdConstants.workGroupOffsetY = (uint)workGroupOffset.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);
            }
        }

        private int _debugIndex = 0;
        private static readonly float[] DebugBarColorSequence =
        {
            0.0f, 1.0f, 1.0f, // teal
            1.0f, 0.42f, 0.0f, // orange
            0.0f, 0.16f, 1.0f, // blue
            0.74f, 1.0f, 0.0f, // lime
            0.68f, 0.0f, 1.0f, // purple
            0.0f, 1.0f, 0.1f, // green
            1.0f, 1.0f, 0.48f // bright yellow
        };
        
        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 DestroyPass(ref FrameInterpolationPass pass)
        {
            if (pass == null)
                return;
            
            pass.Dispose();
            pass = null;
        }
    }
}