Add SH shading to gaussian splat renderer

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

@@ -15,6 +15,72 @@ Shader "Builtin Gaussian Splat Utilities"
         float4 colorOpacity;
     };
 
+    struct GaussianSplatSHData
+    {
+        float coefficients[45];
+    };
+
+    static const float SH_C1 = 0.4886025;
+    static const float SH_C2[] = { 1.0925484, -1.0925484, 0.3153916, -1.0925484, 0.5462742 };
+    static const float SH_C3[] = { -0.5900436, 2.8906114, -0.4570458, 0.3731763, -0.4570458, 1.4453057, -0.5900436 };
+
+    float3 LoadSHCoefficientTriplet(GaussianSplatSHData data, uint coefficientIndex)
+    {
+        return float3(
+            data.coefficients[coefficientIndex + 0u],
+            data.coefficients[coefficientIndex + 15u],
+            data.coefficients[coefficientIndex + 30u]);
+    }
+
+    float3 ShadeGaussianSplatSH(float3 baseColor, GaussianSplatSHData data, float3 direction, uint shOrder)
+    {
+        direction *= -1.0;
+
+        const float x = direction.x;
+        const float y = direction.y;
+        const float z = direction.z;
+        float3 result = baseColor;
+
+        if (shOrder >= 1u)
+        {
+            result += SH_C1 * (
+                -LoadSHCoefficientTriplet(data, 0u) * y +
+                 LoadSHCoefficientTriplet(data, 1u) * z -
+                 LoadSHCoefficientTriplet(data, 2u) * x);
+
+            if (shOrder >= 2u)
+            {
+                const float xx = x * x;
+                const float yy = y * y;
+                const float zz = z * z;
+                const float xy = x * y;
+                const float yz = y * z;
+                const float xz = x * z;
+
+                result +=
+                    (SH_C2[0] * xy) * LoadSHCoefficientTriplet(data, 3u) +
+                    (SH_C2[1] * yz) * LoadSHCoefficientTriplet(data, 4u) +
+                    (SH_C2[2] * (2.0 * zz - xx - yy)) * LoadSHCoefficientTriplet(data, 5u) +
+                    (SH_C2[3] * xz) * LoadSHCoefficientTriplet(data, 6u) +
+                    (SH_C2[4] * (xx - yy)) * LoadSHCoefficientTriplet(data, 7u);
+
+                if (shOrder >= 3u)
+                {
+                    result +=
+                        (SH_C3[0] * y * (3.0 * xx - yy)) * LoadSHCoefficientTriplet(data, 8u) +
+                        (SH_C3[1] * xy * z) * LoadSHCoefficientTriplet(data, 9u) +
+                        (SH_C3[2] * y * (4.0 * zz - xx - yy)) * LoadSHCoefficientTriplet(data, 10u) +
+                        (SH_C3[3] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) * LoadSHCoefficientTriplet(data, 11u) +
+                        (SH_C3[4] * x * (4.0 * zz - xx - yy)) * LoadSHCoefficientTriplet(data, 12u) +
+                        (SH_C3[5] * z * (xx - yy)) * LoadSHCoefficientTriplet(data, 13u) +
+                        (SH_C3[6] * x * (xx - 3.0 * yy)) * LoadSHCoefficientTriplet(data, 14u);
+                }
+            }
+        }
+
+        return max(result, 0.0);
+    }
+
     uint FloatToSortableUint(float value)
     {
         const uint rawValue = asuint(value);
@@ -133,8 +199,10 @@ Shader "Builtin Gaussian Splat Utilities"
                 float4x4 gProjectionMatrix;
                 float4x4 gViewMatrix;
                 float4x4 gModelMatrix;
+                float4x4 gWorldToObjectMatrix;
                 float4 gCameraRight;
                 float4 gCameraUp;
+                float4 gCameraWorldPos;
                 float4 gScreenParams;
                 float4 gSplatParams;
             };
@@ -142,6 +210,7 @@ Shader "Builtin Gaussian Splat Utilities"
             StructuredBuffer<float4> GaussianSplatPositions;
             StructuredBuffer<GaussianSplatOtherData> GaussianSplatOther;
             StructuredBuffer<float4> GaussianSplatColor;
+            StructuredBuffer<GaussianSplatSHData> GaussianSplatSH;
             RWStructuredBuffer<uint> GaussianSplatSortDistances;
             RWStructuredBuffer<uint> GaussianSplatOrderBuffer;
             RWStructuredBuffer<GaussianSplatViewData> GaussianSplatViewDataBuffer;
@@ -170,6 +239,8 @@ Shader "Builtin Gaussian Splat Utilities"
                 const float3 localCenter = GaussianSplatPositions[index].xyz;
                 const GaussianSplatOtherData otherData = GaussianSplatOther[index];
                 const float4 colorOpacity = GaussianSplatColor[index];
+                const GaussianSplatSHData shData = GaussianSplatSH[index];
+                const uint shOrder = min((uint)gSplatParams.z, 3u);
 
                 const float3 worldCenter = mul(gModelMatrix, float4(localCenter, 1.0)).xyz;
                 const float3 viewCenter = mul(gViewMatrix, float4(worldCenter, 1.0)).xyz;
@@ -202,7 +273,16 @@ Shader "Builtin Gaussian Splat Utilities"
                     viewData.clipCenter = clipCenter;
                     viewData.ellipseAxisU = float4(axisU, 0.0, 0.0);
                     viewData.ellipseAxisV = float4(axisV, 0.0, 0.0);
-                    viewData.colorOpacity = colorOpacity;
+                    float3 shadedColor = colorOpacity.rgb;
+                    if (shOrder > 0u)
+                    {
+                        const float3 worldViewDirection = gCameraWorldPos.xyz - worldCenter;
+                        const float3 objectViewDirection = normalize(
+                            mul((float3x3)gWorldToObjectMatrix, worldViewDirection));
+                        shadedColor = ShadeGaussianSplatSH(colorOpacity.rgb, shData, objectViewDirection, shOrder);
+                    }
+
+                    viewData.colorOpacity = float4(shadedColor, colorOpacity.a);
                 }
 
                 GaussianSplatViewDataBuffer[index] = viewData;
@@ -222,8 +302,10 @@ Shader "Builtin Gaussian Splat Utilities"
                 float4x4 gProjectionMatrix;
                 float4x4 gViewMatrix;
                 float4x4 gModelMatrix;
+                float4x4 gWorldToObjectMatrix;
                 float4 gCameraRight;
                 float4 gCameraUp;
+                float4 gCameraWorldPos;
                 float4 gScreenParams;
                 float4 gSplatParams;
             };