com.misaki.ao-volume/Runtime/CustomPass/AoVolumePass.cs

using Misaki.AoVolume;
using System.Buffers;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using UnityEngine.Rendering.HighDefinition;

internal class AoVolumePass : CustomPass
{
    private NativeArray<OrientedBoundingBox> _volumeBounds;
    private ComputeBuffer _volumeBoundsBuffer;
    private ComputeBuffer _visibleIndicesBuffer;
    private ComputeBuffer _visibleVolumeCountBuffer;

    private ComputeBuffer _volumeDataBuffer;

    private RTHandle _volumeBuffer;

    [ResourcePath("Packages/com.misaki.ao-volume/Runtime/Shader/HzCulling.compute")]
    public ComputeShader cullingShader;
    [ResourcePath("Packages/com.misaki.ao-volume/Runtime/Shader/AoVolume.compute")]
    public ComputeShader renderingShader;

    [Range(16, 256)]
    public int maxVolumeOnScreen = 64;

    private void ClearCounterBuffer()
    {
        var zeroBuffer = new NativeArray<uint>(1, Allocator.Temp);
        zeroBuffer[0] = 0u;
        _visibleVolumeCountBuffer.SetData(zeroBuffer);
        zeroBuffer.Dispose();
    }

    // It can be used to configure render targets and their clear state. Also to create temporary render target textures.
    // When empty this render pass will render to the active camera render target.
    // You should never call CommandBuffer.SetRenderTarget. Instead call <c>ConfigureTarget</c> and <c>ConfigureClear</c>.
    // The render pipeline will ensure target setup and clearing happens in an performance manner.
    protected override void Setup(ScriptableRenderContext renderContext, CommandBuffer cmd)
    {
        _volumeBounds = new NativeArray<OrientedBoundingBox>(maxVolumeOnScreen, Allocator.Persistent);
        _volumeBoundsBuffer = new ComputeBuffer(maxVolumeOnScreen, Marshal.SizeOf<OrientedBoundingBox>());
        _visibleIndicesBuffer = new ComputeBuffer(maxVolumeOnScreen, Marshal.SizeOf<uint>(), ComputeBufferType.Raw);
        _visibleVolumeCountBuffer = new ComputeBuffer(1, Marshal.SizeOf<uint>(), ComputeBufferType.Counter);

        _volumeDataBuffer = new ComputeBuffer(maxVolumeOnScreen, Marshal.SizeOf<VolumeData>());

        _volumeBuffer = RTHandles.Alloc(Vector2.one, useDynamicScale: true, dimension: TextureXR.dimension, enableRandomWrite: true, format: GraphicsFormat.R8_UNorm, name: "AO Volume Buffer");

        ClearCounterBuffer();
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static int GetDepthPyramidMaxMipLevel(HDCamera hDCamera)
    {
        return Mathf.FloorToInt(Mathf.Log(Mathf.Min(hDCamera.actualWidth, hDCamera.actualHeight), 2)) - 1;
    }

    private void HierarchicalZCulling(CustomPassContext ctx, int volumeCount)
    {
        for (var i = 0; i < volumeCount; i++)
        {
            _volumeBounds[i] = new OrientedBoundingBox(VolumeDatabase.Instance.VolumeDatas[i].worldMatrix);
        }
        _volumeBoundsBuffer.SetData(_volumeBounds);

        ctx.cmd.SetComputeBufferParam(cullingShader, 0, "_VolumeBounds", _volumeBoundsBuffer);
        ctx.cmd.SetComputeIntParam(cullingShader, "_FullVolumeCount", volumeCount);
        ctx.cmd.SetComputeIntParam(cullingShader, "_DepthPyramidMaxMip", GetDepthPyramidMaxMipLevel(ctx.hdCamera));

        ctx.cmd.SetComputeBufferParam(cullingShader, 0, "_VisibleVolumeCount", _visibleVolumeCountBuffer);
        ctx.cmd.SetComputeBufferParam(cullingShader, 0, "_VisibleVolumeIndices", _visibleIndicesBuffer);

        const int groupSize = 64;
        var threadGroup = (volumeCount + (groupSize - 1)) / groupSize;
        ctx.cmd.DispatchCompute(cullingShader, 0, threadGroup, 1, 1);
    }

    private void RenderVisibleVolumes(CustomPassContext ctx)
    {
        _volumeDataBuffer.SetData(VolumeDatabase.Instance.VolumeDatas);

        ctx.cmd.SetComputeBufferParam(renderingShader, 0, "_VolumeDatas", _volumeDataBuffer);
        ctx.cmd.SetComputeBufferParam(renderingShader, 0, "_VisibleVolumeCount", _visibleVolumeCountBuffer);
        ctx.cmd.SetComputeBufferParam(renderingShader, 0, "_VisibleVolumeIndices", _visibleIndicesBuffer);

        ctx.cmd.SetComputeTextureParam(renderingShader, 0, "_AOVolumeBuffer", _volumeBuffer);

        const int groupSizeX = 8;
        const int groupSizeY = 8;
        var threadGroupX = (ctx.hdCamera.actualWidth + (groupSizeX - 1)) / groupSizeX;
        var threadGroupY = (ctx.hdCamera.actualHeight + (groupSizeY - 1)) / groupSizeY;
        ctx.cmd.DispatchCompute(renderingShader, 0, threadGroupX, threadGroupY, _volumeBuffer.rt.volumeDepth);

        ctx.cmd.SetGlobalTexture("_AmbientOcclusionTexture", _volumeBuffer);
    }

    private void DebugVisibleVolumes()
    {
        var visibleVolumeCount = new int[1];
        _visibleVolumeCountBuffer.GetData(visibleVolumeCount);
        Debug.Log($"Visible Volume Count: {visibleVolumeCount[0]}");

        var visibleVolumeIndices = ArrayPool<uint>.Shared.Rent(64);
        _visibleIndicesBuffer.GetData(visibleVolumeIndices);
        for (var i = 0; i < visibleVolumeCount[0]; i++)
        {
            if (i >= 64)
            {
                break;
            }

            var volumeIndex = visibleVolumeIndices[i];
            Debug.Log($"Visible Volume Index: {volumeIndex}");
            //Debug.Log($"Volume Data: {VolumeDatabase.Instance.VolumeDatas[(int)volumeIndex].worldMatrix.GetColumn(3)}");
        }
        Debug.Log("End");
        ArrayPool<uint>.Shared.Return(visibleVolumeIndices);
    }

    protected override void Execute(CustomPassContext ctx)
    {
        if (cullingShader == null || renderingShader == null)
        {
            return;
        }

        // Worth it to allocate a new buffer?
        //if (Shader.GetGlobalTexture("_AmbientOcclusionTexture") is not RenderTexture gtaoBuffer)
        //{
        //    return;
        //}

        var volumeCount = VolumeDatabase.Instance.EntityCount;
        if (volumeCount <= 0)
        {
            return;
        }

        HierarchicalZCulling(ctx, volumeCount);
        RenderVisibleVolumes(ctx);
        //DebugVisibleVolumes();
        ClearCounterBuffer();
    }

    protected override void Cleanup()
    {
        _volumeBounds.Dispose();
        _volumeBoundsBuffer?.Dispose();
        _visibleIndicesBuffer?.Dispose();
        _visibleVolumeCountBuffer?.Dispose();

        _volumeDataBuffer?.Dispose();

        _volumeBuffer?.Release();
    }
}