XCEngine/engine/assets/builtin/shaders/gaussian-splat-utilities.shader

Shader "Builtin Gaussian Splat Utilities"
{
    HLSLINCLUDE
    struct GaussianSplatOtherData
    {
        float4 rotation;
        float4 scaleReserved;
    };

    struct GaussianSplatViewData
    {
        float4 clipCenter;
        float4 ellipseAxisU;
        float4 ellipseAxisV;
        float4 colorOpacity;
    };

    uint FloatToSortableUint(float value)
    {
        const uint rawValue = asuint(value);
        const uint mask = (rawValue & 0x80000000u) != 0u ? 0xffffffffu : 0x80000000u;
        return rawValue ^ mask;
    }

    float3x3 CalcMatrixFromRotationScale(float4 rotation, float3 scale)
    {
        const float x = rotation.x;
        const float y = rotation.y;
        const float z = rotation.z;
        const float w = rotation.w;

        const float3x3 rotationMatrix = float3x3(
            1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - w * z),       2.0 * (x * z + w * y),
            2.0 * (x * y + w * z),       1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - w * x),
            2.0 * (x * z - w * y),       2.0 * (y * z + w * x),       1.0 - 2.0 * (x * x + y * y));
        const float3x3 scaleMatrix = float3x3(
            scale.x, 0.0,    0.0,
            0.0,    scale.y, 0.0,
            0.0,    0.0,    scale.z);
        return mul(rotationMatrix, scaleMatrix);
    }

    void CalcCovariance3D(float3x3 rotationScaleMatrix, out float3 sigma0, out float3 sigma1)
    {
        const float3x3 sigma = mul(rotationScaleMatrix, transpose(rotationScaleMatrix));
        sigma0 = float3(sigma._m00, sigma._m01, sigma._m02);
        sigma1 = float3(sigma._m11, sigma._m12, sigma._m22);
    }

    float3 CalcCovariance2D(
        float3 viewPosition,
        float3 covariance3D0,
        float3 covariance3D1,
        float4x4 viewMatrix,
        float4x4 projectionMatrix,
        float2 screenSize)
    {
        if (abs(viewPosition.z) <= 1.0e-5)
        {
            return float3(0.3, 0.0, 0.3);
        }

        const float aspect = projectionMatrix[0][0] / projectionMatrix[1][1];
        const float tanFovX = rcp(projectionMatrix[0][0]);
        const float tanFovY = rcp(projectionMatrix[1][1] * aspect);
        const float limitX = 1.3 * tanFovX;
        const float limitY = 1.3 * tanFovY;

        float3 clampedViewPosition = viewPosition;
        clampedViewPosition.x =
            clamp(clampedViewPosition.x / clampedViewPosition.z, -limitX, limitX) * clampedViewPosition.z;
        clampedViewPosition.y =
            clamp(clampedViewPosition.y / clampedViewPosition.z, -limitY, limitY) * clampedViewPosition.z;

        const float focalLength = screenSize.x * projectionMatrix[0][0] * 0.5;
        const float3x3 jacobian = float3x3(
            focalLength / clampedViewPosition.z,
            0.0,
            -(focalLength * clampedViewPosition.x) / (clampedViewPosition.z * clampedViewPosition.z),
            0.0,
            focalLength / clampedViewPosition.z,
            -(focalLength * clampedViewPosition.y) / (clampedViewPosition.z * clampedViewPosition.z),
            0.0,
            0.0,
            0.0);
        const float3x3 worldToView = (float3x3)viewMatrix;
        const float3x3 transform = mul(jacobian, worldToView);
        const float3x3 covariance3D = float3x3(
            covariance3D0.x, covariance3D0.y, covariance3D0.z,
            covariance3D0.y, covariance3D1.x, covariance3D1.y,
            covariance3D0.z, covariance3D1.y, covariance3D1.z);
        float3x3 covariance2D = mul(transform, mul(covariance3D, transpose(transform)));

        covariance2D._m00 += 0.3;
        covariance2D._m11 += 0.3;
        return float3(covariance2D._m00, covariance2D._m01, covariance2D._m11);
    }

    void DecomposeCovariance(float3 covariance2D, out float2 axisU, out float2 axisV)
    {
        const float diagonal0 = covariance2D.x;
        const float diagonal1 = covariance2D.z;
        const float offDiagonal = covariance2D.y;
        const float mid = 0.5 * (diagonal0 + diagonal1);
        const float radius = length(float2((diagonal0 - diagonal1) * 0.5, offDiagonal));
        const float lambda0 = max(mid + radius, 0.1);
        const float lambda1 = max(mid - radius, 0.1);

        float2 basis = normalize(float2(offDiagonal, lambda0 - diagonal0));
        if (all(abs(basis) < 1.0e-5))
        {
            basis = float2(1.0, 0.0);
        }

        basis.y = -basis.y;
        const float maxAxisLength = 4096.0;
        axisU = min(sqrt(2.0 * lambda0), maxAxisLength) * basis;
        axisV = min(sqrt(2.0 * lambda1), maxAxisLength) * float2(basis.y, -basis.x);
    }
    ENDHLSL

    SubShader
    {
        Pass
        {
            Name "GaussianSplatPrepareOrder"
            HLSLPROGRAM
            #pragma target 4.5
            #pragma compute GaussianSplatPrepareOrderCS

            cbuffer PerObjectConstants
            {
                float4x4 gProjectionMatrix;
                float4x4 gViewMatrix;
                float4x4 gModelMatrix;
                float4 gCameraRight;
                float4 gCameraUp;
                float4 gScreenParams;
                float4 gSplatParams;
            };

            StructuredBuffer<float4> GaussianSplatPositions;
            StructuredBuffer<GaussianSplatOtherData> GaussianSplatOther;
            StructuredBuffer<float4> GaussianSplatColor;
            RWStructuredBuffer<uint> GaussianSplatSortDistances;
            RWStructuredBuffer<uint> GaussianSplatOrderBuffer;
            RWStructuredBuffer<GaussianSplatViewData> GaussianSplatViewDataBuffer;

            [numthreads(64, 1, 1)]
            void GaussianSplatPrepareOrderCS(uint3 dispatchThreadId : SV_DispatchThreadID)
            {
                const uint splatCount = (uint)gSplatParams.x;
                const uint index = dispatchThreadId.x;
                if (index >= splatCount)
                {
                    return;
                }

                GaussianSplatOrderBuffer[index] = index;

                GaussianSplatViewData viewData = (GaussianSplatViewData)0;
                const float3 localCenter = GaussianSplatPositions[index].xyz;
                const GaussianSplatOtherData otherData = GaussianSplatOther[index];
                const float4 colorOpacity = GaussianSplatColor[index];

                const float3 worldCenter = mul(gModelMatrix, float4(localCenter, 1.0)).xyz;
                const float3 viewCenter = mul(gViewMatrix, float4(worldCenter, 1.0)).xyz;
                GaussianSplatSortDistances[index] = FloatToSortableUint(viewCenter.z);

                const float4 clipCenter = mul(gProjectionMatrix, float4(viewCenter, 1.0));
                if (clipCenter.w > 0.0)
                {
                    const float3x3 modelLinear = (float3x3)gModelMatrix;
                    const float3x3 rotationScaleMatrix =
                        CalcMatrixFromRotationScale(otherData.rotation, otherData.scaleReserved.xyz);
                    const float3x3 worldRotationScale = mul(modelLinear, rotationScaleMatrix);

                    float3 covariance3D0 = 0.0;
                    float3 covariance3D1 = 0.0;
                    CalcCovariance3D(worldRotationScale, covariance3D0, covariance3D1);

                    const float3 covariance2D = CalcCovariance2D(
                        viewCenter,
                        covariance3D0,
                        covariance3D1,
                        gViewMatrix,
                        gProjectionMatrix,
                        gScreenParams.xy);

                    float2 axisU = 0.0;
                    float2 axisV = 0.0;
                    DecomposeCovariance(covariance2D, axisU, axisV);

                    viewData.clipCenter = clipCenter;
                    viewData.ellipseAxisU = float4(axisU, 0.0, 0.0);
                    viewData.ellipseAxisV = float4(axisV, 0.0, 0.0);
                    viewData.colorOpacity = colorOpacity;
                }

                GaussianSplatViewDataBuffer[index] = viewData;
            }
            ENDHLSL
        }
    }
}