unity_native_render_plugin/Packages/com.unity.render-pipelines.universal@14.0.11/Runtime/OcclusionCulling/HierarchicalZFeature.cs

using System;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;

namespace X.Rendering.Feature
{
    public class HierarchicalZFeature : ScriptableRendererFeature
    {
        public enum EDepthPyramidFunc
        {
            CopyDepth,
            Compute,
            SPD
        }

        [Serializable]
        public class Settings
        {
            public EDepthPyramidFunc PyramidFunc = EDepthPyramidFunc.CopyDepth;
            public bool SkipThreeMip;
            // 留历史 depth，做 tow pass cull
            public bool UseThreeFrameReadback;
            public ComputeShader ComputeShader;
            public ComputeShader Spd;
            public Material CopyDepth;
        }

        [SerializeField]
        private Settings settings;

        private HizPass hizPass;

        public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData)
        {
            if (renderingData.cameraData.cameraType == UnityEngine.CameraType.Game)
            {
                renderer.EnqueuePass(hizPass);
            }
        }

        public override void Create()
        {
            hizPass = new(settings);
        }

        protected override void Dispose(bool disposing)
        {
            base.Dispose(disposing);
            if(hizPass != null)
            {
                hizPass.Dispose();
            }
        }

        internal class HizPass : ScriptableRenderPass, IDisposable
        {
            private readonly Settings settings;
            private ProfilingSampler profiler;

            private Vector2Int cachedHardwareTextureSize;
            private bool cachedSkipThreeMip;
            private Vector2Int mip0SizeNOP;
            private Vector2Int depthPyramidTextureSize;
            private int mipLevelCount;
            private Vector2Int[] mipLevelOffsets;
            private Vector2Int[] mipLevelSizes;
            private Vector4[] mipOffsetAndSizes = new Vector4[16];

            private RTHandle[] depthPyramidTexs = new RTHandle[3];
            private static string[] depthPyramidNames = new string[3] { "_DepthPyramidTextureA", "_DepthPyramidTextureB", "_DepthPyramidTextureC" };

            RenderTexture SpdAtomicCounter;

            public HizPass(Settings settings)
            {
                renderPassEvent = RenderPassEvent.BeforeRenderingPostProcessing;
                profiler = new("DepthPyramid");
                this.settings = settings;
                mipLevelOffsets = new Vector2Int[16];
                mipLevelSizes = new Vector2Int[16];
                SpdAtomicCounter = new RenderTexture(1, 1, 0, GraphicsFormat.R32_UInt) { name = "FSR2_SpdAtomicCounter", enableRandomWrite = true };
                SpdAtomicCounter.Create();
            }

            public static int HizIndex { get; private set; }

            private int GetHizIndex()
            {
                if(settings.UseThreeFrameReadback)
                {
                    HizIndex = Time.frameCount % 3;
                }
                else
                {
                    HizIndex = 0;
                }

                return HizIndex;
            }
            
            public void ComputePackedMipChainInfo(Vector2Int viewportSize)
            {
                if (cachedHardwareTextureSize == viewportSize && cachedSkipThreeMip == settings.SkipThreeMip)
                {
                    return;
                }

                cachedHardwareTextureSize = viewportSize;
                cachedSkipThreeMip = settings.SkipThreeMip;

                mip0SizeNOP = viewportSize;
                // PowerOfTwo 不会留缝隙
                int resizeX = settings.PyramidFunc == EDepthPyramidFunc.SPD || Mathf.IsPowerOfTwo(viewportSize.x) ? viewportSize.x : Mathf.NextPowerOfTwo(viewportSize.x);
                int resizeY = settings.PyramidFunc == EDepthPyramidFunc.SPD || Mathf.IsPowerOfTwo(viewportSize.y) ? viewportSize.y : Mathf.NextPowerOfTwo(viewportSize.y);
                Vector2Int hardwareTextureSize = new Vector2Int(resizeX, resizeY);
                mipLevelOffsets[0] = Vector2Int.zero;
                mipLevelSizes[0] = Vector2Int.zero;
                int mipLevel = 0;
                Vector2Int mipSize = hardwareTextureSize;
                Vector2Int texSize = Vector2Int.zero;
                do
                {
                    mipLevel++;

                    // Round up.
                    mipSize.x = System.Math.Max(1, (mipSize.x + 1) >> 1);
                    mipSize.y = System.Math.Max(1, (mipSize.y + 1) >> 1);

                    mipLevelSizes[mipLevel] = mipSize;

                    Vector2Int prevMipBegin = mipLevelOffsets[mipLevel - 1];
                    Vector2Int prevMipEnd = prevMipBegin + mipLevelSizes[mipLevel - 1];

                    Vector2Int mipBegin = new Vector2Int();

                    if ((mipLevel & 1) != 0)
                    {
                        mipBegin.x = prevMipBegin.x;
                        mipBegin.y = prevMipEnd.y;
                    }
                    else
                    {
                        mipBegin.x = prevMipEnd.x;
                        mipBegin.y = prevMipBegin.y;
                    }
                    if (settings.SkipThreeMip && mipLevel == 4) //跳过前面3级的Mip,从第四级开始算
                    {
                        mipBegin = Vector2Int.zero;
                        texSize = Vector2Int.zero;
                    }

                    mipLevelOffsets[mipLevel] = mipBegin;

                    texSize.x = System.Math.Max(texSize.x, mipBegin.x + mipSize.x);
                    texSize.y = System.Math.Max(texSize.y, mipBegin.y + mipSize.y);
                }
                while ((mipSize.x > 1) || (mipSize.y > 1));
                mipLevelSizes[0] = hardwareTextureSize;
                //RT实际大小
                depthPyramidTextureSize = new Vector2Int((int)Mathf.Ceil((float)texSize.x), (int)Mathf.Ceil((float)texSize.y));
                mipLevelCount = mipLevel + 1;
            }

            private static int[] depthMipId = null;

            private void DoCopyDepth(CommandBuffer cmd, Texture depthTex)
            {
                if (depthMipId == null)
                {
                    depthMipId = new int[32];
                    for (int i = 0; i < depthMipId.Length; i++)
                    {
                        depthMipId[i] = Shader.PropertyToID("_DepthMip" + i);
                    }
                }

                var sampleName = "DepthDownSample";
                cmd.BeginSample(sampleName);
                cmd.SetViewProjectionMatrices(Matrix4x4.identity, Matrix4x4.identity);
                cmd.SetGlobalTexture(HizShaderIds.DepthInputId, depthTex);
                for (int i = 1; i < mipLevelCount; i++)
                {
                    var index = i;
                    if (settings.SkipThreeMip && i < 3 + 1)
                    {
                        continue;
                    }
                    else if (settings.SkipThreeMip && i == 3 + 1)
                    {
                        index = 1;
                    }
                        
                    var mipSize = mipLevelSizes[i];
                    var inputMipSize = index == 1 ? mip0SizeNOP : mipLevelSizes[i - 1];
                    var texId = depthMipId[i];
                    cmd.SetGlobalVector(HizShaderIds.InputScaleAndMaxIndexId, new Vector4(inputMipSize.x / (float)mipSize.x, inputMipSize.y / (float)mipSize.y, inputMipSize.x - 1, inputMipSize.y - 1));
                    cmd.GetTemporaryRT(texId, mipSize.x, mipSize.y, 0, FilterMode.Point, RenderTextureFormat.RFloat);
                    cmd.SetRenderTarget(texId, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
                    cmd.DrawMesh(RenderingUtils.fullscreenMesh, Matrix4x4.identity, settings.CopyDepth, 0, 0);
                    cmd.SetGlobalTexture(HizShaderIds.DepthInputId, texId);
                }

                cmd.EndSample(sampleName);

                sampleName = "DepthCombine";
                cmd.BeginSample(sampleName);
                var hizIndex = GetHizIndex();
                RTHandle hizRt = depthPyramidTexs[hizIndex];
           
                RenderingUtils.ReAllocateIfNeeded(ref hizRt, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                }, filterMode : FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizRt;
                    
                cmd.SetRenderTarget(hizRt, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
                for (int i = 1; i < mipLevelCount; i++)
                {
                    if (settings.SkipThreeMip && i < 3 + 1)
                    {
                        continue;
                    }
                    
                    var texId = depthMipId[i];
                    cmd.SetGlobalTexture(HizShaderIds.DepthInputId, texId);
                    var mipSize = mipLevelSizes[i];
                    var mipOffset = mipLevelOffsets[i];
                    mipOffsetAndSizes[i] = new (mipOffset.x, mipOffset.y, mipSize.x, mipSize.y);
                    cmd.SetViewport(new Rect(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    cmd.DrawMesh(RenderingUtils.fullscreenMesh, Matrix4x4.identity, settings.CopyDepth, 0, 1);
                }
                
                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizRt);
                cmd.EndSample(sampleName);

                for (int i = 0; i < depthMipId.Length; i++)
                {
                    var texId = depthMipId[i];
                    cmd.ReleaseTemporaryRT(texId);
                }
            }

            private void DoComputeDepth(CommandBuffer cmd, Texture depthTex)
            {
                var hizIndex = GetHizIndex();
                RTHandle hizBuffer = depthPyramidTexs[hizIndex];

                RenderingUtils.ReAllocateIfNeeded(ref hizBuffer, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                    enableRandomWrite = true,
                }, filterMode: FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizBuffer;

                int outputMipCnt = mipLevelCount - 1;
                int dispatchCnt = Mathf.CeilToInt(outputMipCnt / 4f);
                int mipCnt = outputMipCnt;
                Matrix4x4 matrix4X4 = new Matrix4x4();
                for (int i = 0; i <= dispatchCnt; i++) 
                {
                    int  startMip = i * 4;
                    int mip1 = startMip + 1;
                    int mip2 = startMip + 2;
                    int mip3 = startMip + 3;
                    int mip4 = startMip + 4;
                    var outputMipSize = mipLevelSizes[mip1];
                    if(outputMipSize.x == 0 || outputMipSize.y == 0)
                    {
                        break;
                    }
                    int kernelId = 0;
                    if (i == 0)
                    {
                        kernelId = 1;
                        if (settings.SkipThreeMip)
                        {
                            kernelId = 2;
                        }
                    }

                    cmd.SetComputeTextureParam(settings.ComputeShader, kernelId, HizShaderIds.DepthInputId, i == 0 ? depthTex : hizBuffer);
                    cmd.SetComputeTextureParam(settings.ComputeShader, kernelId, "_DepthMipChain", hizBuffer);

                    int inputMipIndex = startMip;
                    var inputMipOffset = mipLevelOffsets[inputMipIndex];
                    var inputMipSize = inputMipIndex == 0 ? mip0SizeNOP : mipLevelSizes[inputMipIndex];

                    cmd.SetComputeVectorParam(settings.ComputeShader, "_InputMipOffsetAndSize", new Vector4(inputMipOffset.x, inputMipOffset.y, inputMipSize.x, inputMipSize.y));
                    cmd.SetComputeIntParam(settings.ComputeShader, "_MipCount", Mathf.Min(mipCnt, 4));

                    ref var mipOffset = ref mipLevelOffsets[mip1];
                    ref var mipSize = ref mipLevelSizes[mip1];
                    matrix4X4.SetRow(0, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip2];
                    mipSize = ref mipLevelSizes[mip2];
                    matrix4X4.SetRow(1, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip3];
                    mipSize = ref mipLevelSizes[mip3];
                    matrix4X4.SetRow(2, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip4];
                    mipSize = ref mipLevelSizes[mip4];
                    matrix4X4.SetRow(3, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));

                    cmd.SetComputeMatrixParam(settings.ComputeShader, "_MipOffsetAndSizeArray", matrix4X4);

                    // XXX: 可以测试是否有正面作用
                    // cmd.SetExecutionFlags(CommandBufferExecutionFlags.AsyncCompute);
                    cmd.DispatchCompute(settings.ComputeShader, kernelId, Mathf.CeilToInt(outputMipSize.x / 8f), Mathf.CeilToInt(outputMipSize.y / 8f), 1);
                    mipCnt = mipCnt - 4;
                }
                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizBuffer);
            }

            private void DoSpdDepth(CommandBuffer cmd, Texture depthTex)
            {
                var hizIndex = GetHizIndex();
                RTHandle hizBuffer = depthPyramidTexs[hizIndex];

                RenderingUtils.ReAllocateIfNeeded(ref hizBuffer, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                    enableRandomWrite = true,
                }, filterMode: FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizBuffer;

                var dispatchX = Mathf.CeilToInt(mip0SizeNOP.x / 64f);
                var dispatchY = Mathf.CeilToInt(mip0SizeNOP.y / 64f);
                cmd.SetComputeIntParam(settings.Spd, "mips", mipLevelCount);
                cmd.SetComputeIntParam(settings.Spd, "numWorkGroups", dispatchX);

                for (int i = 0; i < mipLevelCount; i++)
                {
                    var mipSize = mipLevelSizes[i];
                    var mipOffset = mipLevelOffsets[i];
                    mipOffsetAndSizes[i] = new(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y);
                }
                cmd.SetComputeVectorArrayParam(settings.Spd, "_MipOffsetAndSizeArray", mipOffsetAndSizes);

                cmd.SetComputeTextureParam(settings.Spd, 0, "_InputDepth", depthTex);
                cmd.SetComputeTextureParam(settings.Spd, 0, "_OutDepth", hizBuffer);
                cmd.SetComputeTextureParam(settings.Spd, 0, "rw_internal_global_atomic", SpdAtomicCounter);
                
                cmd.DispatchCompute(settings.Spd, 0, dispatchX, dispatchY, 1);
                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizBuffer);
            }

            public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
            {
                var cmd = renderingData.commandBuffer;
                using var soc = new ProfilingScope(cmd, profiler);
                Texture depthTex = Shader.GetGlobalTexture(HizShaderIds.CameraDepthTextureId);
                if(depthTex == null)
                {
                    return;
                }

                ComputePackedMipChainInfo(new Vector2Int(depthTex.width, depthTex.height));

                switch (settings.PyramidFunc)
                {
                    case EDepthPyramidFunc.CopyDepth:
                        DoCopyDepth(cmd, depthTex);
                        break;
                    case EDepthPyramidFunc.Compute:
                        DoComputeDepth(cmd, depthTex);
                        break;
                    case EDepthPyramidFunc.SPD:
                        DoSpdDepth(cmd, depthTex);
                        break;
                    default:
                        break;
                }

                cmd.SetGlobalVector(HizShaderIds.Mip0SizeId, new Vector4(mip0SizeNOP.x, mip0SizeNOP.y, 0, 0));
                cmd.SetGlobalVector(HizShaderIds.MipmapLevelMinMaxIndexId, new Vector4(1, mipLevelCount, 0, 0));
                cmd.SetGlobalVectorArray(HizShaderIds.MipOffsetAndSizeArrayId, mipOffsetAndSizes);
            }

            public void Dispose()
            {
                if (SpdAtomicCounter != null)
                {
                    SpdAtomicCounter.Release();
                    SpdAtomicCounter = null;
                }

                for (int i = 0; i < depthPyramidTexs.Length; i++)
                {
                    var rt = depthPyramidTexs[i];
                    rt?.Release();
                }
            }
        }
    }
}