o3de/Gems/Atom/Feature/Common/Assets/Shaders/Reflections/ReflectionScreenSpaceDownsampleDepthLinear.azsl

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <Atom/Features/PostProcessing/FullscreenVertex.azsli>
#include <Atom/Features/SrgSemantics.azsli>

ShaderResourceGroup PassSrg : SRG_PerPass
{
    Texture2D<float2> m_input;

    float m_invOutputScale;
    uint  m_mipLevel;
    bool  m_halfResolution;
    uint  m_previousMipLevel;
    uint  m_previousMipWidth;
    uint  m_previousMipHeight;

    Sampler PointSampler
    {
        MinFilter = Point;
        MagFilter = Point;
        MipFilter = Point;
        AddressU  = Clamp;
        AddressV  = Clamp;
        AddressW  = Clamp;
    };
}

struct PSOutput
{    
    float2 m_result : SV_Target0;
};

PSOutput MainPS(VSOutput IN)
{
    uint2 coords = IN.m_position.xy;
    uint2 previousMipBaseCoords = IN.m_position.xy * 2;

    float minDepth = 1e27f;
    float sigma = 0.0f;
    float2 depthSamples[9]; // maximum of 3x3 samples depending on image width/height
    uint numSamples = 0;

    if (PassSrg::m_invOutputScale <= 1.0f)
    {
        // populate mip0 at 1:1 scale
        minDepth = PassSrg::m_input.Load(uint3(coords, 0)).r;
    }
    else
    {
        [unroll]
        for (int y = 0; y < 2; ++y)
        {
            [unroll]
            for (int x = 0; x < 2; ++x)
            {
                float2 previousMipCoords = float2(previousMipBaseCoords + float2(x, y));
                float2 previousMipUV = previousMipCoords / float2(PassSrg::m_previousMipWidth, PassSrg::m_previousMipHeight);
                depthSamples[numSamples++] = PassSrg::m_input.SampleLevel(PassSrg::PointSampler, previousMipUV, PassSrg::m_previousMipLevel).rg;
            }
        }

        bool oddWidth  = (PassSrg::m_previousMipWidth & 1) != 0;
        bool oddHeight = (PassSrg::m_previousMipHeight & 1) != 0;

        if (oddWidth)
        {
            // add extra column
            [unroll]
            for (int y = 0; y < 2; ++y)
            {
                float2 previousMipCoords = float2(previousMipBaseCoords + float2(2, y));
                float2 previousMipUV = previousMipCoords / float2(PassSrg::m_previousMipWidth, PassSrg::m_previousMipHeight);
                depthSamples[numSamples++] = PassSrg::m_input.SampleLevel(PassSrg::PointSampler, previousMipUV, PassSrg::m_previousMipLevel).rg;
            }

            if (oddHeight)
            {
                // both width and height are odd, add corner value
                float2 previousMipCoords = float2(previousMipBaseCoords + float2(2, 2));
                float2 previousMipUV = previousMipCoords / float2(PassSrg::m_previousMipWidth, PassSrg::m_previousMipHeight);
                depthSamples[numSamples++] = PassSrg::m_input.SampleLevel(PassSrg::PointSampler, previousMipUV, PassSrg::m_previousMipLevel).rg;
            }
        }

        if (oddHeight)
        {
            // add extra row
            [unroll]
            for (int x = 0; x < 2; ++x)
            {
                float2 previousMipCoords = float2(previousMipBaseCoords + float2(x, 2));
                float2 previousMipUV = previousMipCoords / float2(PassSrg::m_previousMipWidth, PassSrg::m_previousMipHeight);
                depthSamples[numSamples++] = PassSrg::m_input.SampleLevel(PassSrg::PointSampler, previousMipUV, PassSrg::m_previousMipLevel).rg;
            }
        }
    }

    if (numSamples > 0)
    {
        // find the minimum depth and compute the mean
        float mean = 0.0f;
        for (uint i = 0; i < numSamples; ++i)
        {
            if (depthSamples[i].r < minDepth)
            {
                minDepth = depthSamples[i].r;
            }

            mean += depthSamples[i].r;
        }
        mean /= numSamples;
                   
        // compute variance from depth samples
        float variance = 0.0f;
        for (uint i = 0; i < numSamples; ++i)
        {
            variance += pow(depthSamples[i].r - mean, 2);
        }
        variance /= numSamples;

        // compute standard deviation
        sigma = sqrt(variance);
        
        // find the max sigma for mip levels after the first downsampled mip
        // Note: this allows handling of depth discontinuities across samples as well as inside a single sample
        bool firstDownsampledMip = PassSrg::m_halfResolution ? (PassSrg::m_mipLevel == 0) : (PassSrg::m_mipLevel == 1);
        if (!firstDownsampledMip)
        {
            for (uint i = 0; i < numSamples; ++i)
            {
                if (depthSamples[i].g > sigma)
                {
                    sigma = depthSamples[i].g;
                }
            }                
        }
    }

    PSOutput OUT;
    OUT.m_result = float2(minDepth, sigma);
    return OUT;
}