WGSL implementation of VSM shadows

src/scene/renderer/shadow-renderer.js

@@ -4,7 +4,10 @@ import { Color } from '../../core/math/color.js';
 import { Mat4 } from '../../core/math/mat4.js';
 import { Vec3 } from '../../core/math/vec3.js';
 import { Vec4 } from '../../core/math/vec4.js';
-import { SHADERSTAGE_FRAGMENT, SHADERSTAGE_VERTEX, UNIFORMTYPE_MAT4, UNIFORM_BUFFER_DEFAULT_SLOT_NAME } from '../../platform/graphics/constants.js';
+import {
+    SEMANTIC_POSITION, SHADERLANGUAGE_GLSL, SHADERLANGUAGE_WGSL, SHADERSTAGE_FRAGMENT, SHADERSTAGE_VERTEX,
+    UNIFORMTYPE_MAT4, UNIFORM_BUFFER_DEFAULT_SLOT_NAME
+} from '../../platform/graphics/constants.js';
 import { DebugGraphics } from '../../platform/graphics/debug-graphics.js';
 import { drawQuadWithShader } from '../graphics/quad-render-utils.js';
 import {
@@ -21,6 +24,7 @@ import { LightCamera } from './light-camera.js';
 import { UniformBufferFormat, UniformFormat } from '../../platform/graphics/uniform-buffer-format.js';
 import { BindUniformBufferFormat, BindGroupFormat } from '../../platform/graphics/bind-group-format.js';
 import { BlendState } from '../../platform/graphics/blend-state.js';
+import { shaderChunksWGSL } from '../shader-lib/chunks-wgsl/chunks-wgsl.js';
 
 /**
  * @import { Camera } from '../camera.js'
@@ -88,7 +92,8 @@ class ShadowRenderer {
         this.sourceId = scope.resolve('source');
         this.pixelOffsetId = scope.resolve('pixelOffset');
         this.weightId = scope.resolve('weight[0]');
-        this.blurVsmShaderCode = [shaderChunks.blurVSMPS, `#define GAUSS\n${shaderChunks.blurVSMPS}`];
+        const chunks = this.device.isWebGPU ? shaderChunksWGSL : shaderChunks;
+        this.blurVsmShaderCode = [chunks.blurVSMPS, `#define GAUSS\n${chunks.blurVSMPS}`];
 
         // cache for vsm blur shaders
         this.blurVsmShader = [{}, {}];
@@ -490,11 +495,14 @@ class ShadowRenderer {
         if (!blurShader) {
             this.blurVsmWeights[filterSize] = gaussWeights(filterSize);
 
-            const blurVS = shaderChunks.fullscreenQuadVS;
-            let blurFS = `#define SAMPLES ${filterSize}\n`;
+            const chunks = this.device.isWebGPU ? shaderChunksWGSL : shaderChunks;
+            const blurVS = chunks.fullscreenQuadVS;
+            let blurFS = `#define {SAMPLES} ${filterSize}\n`;
             blurFS += this.blurVsmShaderCode[blurMode];
             const blurShaderName = `blurVsm${blurMode}${filterSize}`;
-            blurShader = createShaderFromCode(this.device, blurVS, blurFS, blurShaderName);
+            blurShader = createShaderFromCode(this.device, blurVS, blurFS, blurShaderName, { vertex_position: SEMANTIC_POSITION }, undefined, {
+                shaderLanguage: this.device.isWebGPU ? SHADERLANGUAGE_WGSL : SHADERLANGUAGE_GLSL
+            });
             cache[blurMode][filterSize] = blurShader;
         }
 

src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js

@@ -7,7 +7,7 @@ import basePS from './lit/frag/base.js';
 // import baseNineSlicedPS from './lit/frag/baseNineSliced.js';
 // import baseNineSlicedTiledPS from './lit/frag/baseNineSlicedTiled.js';
 import bayerPS from './common/frag/bayer.js';
-// import blurVSMPS from './lit/frag/blurVSM.js';
+import blurVSMPS from './lit/frag/blurVSM.js';
 import clearCoatPS from './standard/frag/clearCoat.js';
 import clearCoatGlossPS from './standard/frag/clearCoatGloss.js';
 import clearCoatNormalPS from './standard/frag/clearCoatNormal.js';
@@ -38,7 +38,7 @@ import falloffLinearPS from './lit/frag/falloffLinear.js';
 import fogPS from './common/frag/fog.js';
 import fresnelSchlickPS from './lit/frag/fresnelSchlick.js';
 // import fullscreenQuadPS from './common/frag/fullscreenQuad.js';
-// import fullscreenQuadVS from './common/vert/fullscreenQuad.js';
+import fullscreenQuadVS from './common/vert/fullscreenQuad.js';
 import gammaPS from './common/frag/gamma.js';
 import glossPS from './standard/frag/gloss.js';
 // import gsplatCenterVS from './gsplat/vert/gsplatCenter.js';
@@ -159,7 +159,7 @@ import reprojectVS from './internal/vert/reproject.js';
 // import sampleCatmullRomPS from './common/frag/sampleCatmullRom.js';
 // import screenDepthPS from './common/frag/screenDepth.js';
 import shadowCascadesPS from './lit/frag/lighting/shadowCascades.js';
-// import shadowEVSMPS from './lit/frag/lighting/shadowEVSM.js';
+import shadowEVSMPS from './lit/frag/lighting/shadowEVSM.js';
 import shadowPCF1PS from './lit/frag/lighting/shadowPCF1.js';
 import shadowPCF3PS from './lit/frag/lighting/shadowPCF3.js';
 import shadowPCF5PS from './lit/frag/lighting/shadowPCF5.js';
@@ -217,7 +217,7 @@ const shaderChunksWGSL = {
     // baseNineSlicedPS,
     // baseNineSlicedTiledPS,
     bayerPS,
-    // blurVSMPS,
+    blurVSMPS,
     clearCoatPS,
     clearCoatGlossPS,
     clearCoatNormalPS,
@@ -248,7 +248,7 @@ const shaderChunksWGSL = {
     fogPS,
     fresnelSchlickPS,
     // fullscreenQuadPS,
-    // fullscreenQuadVS,
+    fullscreenQuadVS,
     gammaPS,
     glossPS,
     // gsplatCenterVS,
@@ -370,7 +370,7 @@ const shaderChunksWGSL = {
     // sampleCatmullRomPS,
     // screenDepthPS,
     shadowCascadesPS,
-    // shadowEVSMPS,
+    shadowEVSMPS,
     shadowPCF1PS,
     shadowPCF3PS,
     shadowPCF5PS,

src/scene/shader-lib/chunks-wgsl/common/vert/fullscreenQuad.js

@@ -0,0 +1,13 @@
+export default /* wgsl */`
+attribute vertex_position: vec2f;
+
+varying vUv0: vec2f;
+
+@vertex
+fn vertexMain(input: VertexInput) -> VertexOutput {
+    var output: VertexOutput;
+    output.position = vec4f(input.vertex_position, 0.5, 1.0);
+    output.vUv0 = input.vertex_position.xy * 0.5 + vec2f(0.5);
+    return output;
+}
+`;

src/scene/shader-lib/chunks-wgsl/lit/frag/blurVSM.js

@@ -0,0 +1,35 @@
+export default /* wgsl */`
+varying vUv0: vec2f;
+
+var source: texture_2d<f32>;
+var sourceSampler: sampler;
+
+#ifdef GAUSS
+    uniform weight: array<f32, {SAMPLES}>;
+#endif
+uniform pixelOffset: vec2f;
+
+@fragment
+fn fragmentMain(input: FragmentInput) -> FragmentOutput {
+    var output: FragmentOutput;
+    var moments: vec3f = vec3f(0.0);
+    let uv: vec2f = input.vUv0 - uniform.pixelOffset * (f32({SAMPLES}) * 0.5);
+
+    for (var i: i32 = 0; i < {SAMPLES}; i = i + 1) {
+        let c: vec4f = textureSample(source, sourceSampler, uv + uniform.pixelOffset * f32(i));
+
+        #ifdef GAUSS
+            moments = moments + c.xyz * uniform.weight[i].element;
+        #else
+            moments = moments + c.xyz;
+        #endif
+    }
+
+    #ifndef GAUSS
+        moments = moments * (1.0 / f32({SAMPLES}));
+    #endif
+
+    output.color = vec4f(moments, 1.0);
+    return output;
+}
+`;

src/scene/shader-lib/chunks-wgsl/lit/frag/lighting/lightDeclaration.js

@@ -81,7 +81,7 @@ export default /* wgsl */`
                 var light{i}_shadowMap: texture_depth_2d;
                 var light{i}_shadowMapSampler: sampler_comparison;
             #else
-                var light{i}_shadowMap texture_2D<f32>;
+                var light{i}_shadowMap: texture_2d<f32>;
                 var light{i}_shadowMapSampler: sampler;
             #endif
         #endif

src/scene/shader-lib/chunks-wgsl/lit/frag/lighting/shadowEVSM.js

@@ -0,0 +1,88 @@
+export default /* wgsl */`
+
+// ------ VSM Shared ------
+
+fn linstep(a: f32, b: f32, v: f32) -> f32 {
+    // WGSL saturate -> clamp
+    return clamp((v - a) / (b - a), 0.0, 1.0);
+}
+
+fn reduceLightBleeding(pMax: f32, amount: f32) -> f32 {
+   // Remove the [0, amount] tail and linearly rescale (amount, 1].
+   return linstep(amount, 1.0, pMax);
+}
+
+fn chebyshevUpperBound(moments: vec2f, mean: f32, minVariance: f32, lightBleedingReduction: f32) -> f32 {
+    // Compute variance
+    var variance: f32 = moments.y - (moments.x * moments.x);
+    variance = max(variance, minVariance);
+
+    // Compute probabilistic upper bound
+    let d: f32 = mean - moments.x;
+    var pMax: f32 = variance / (variance + (d * d));
+
+    pMax = reduceLightBleeding(pMax, lightBleedingReduction);
+
+    // One-tailed Chebyshev
+    return select(pMax, 1.0, mean <= moments.x);
+}
+
+fn calculateEVSM(moments_in: vec3f, Z_in: f32, vsmBias: f32, exponent: f32) -> f32 {
+    let Z: f32 = 2.0 * Z_in - 1.0;
+    let warpedDepth: f32 = exp(exponent * Z);
+
+    let moments: vec2f = moments_in.xy + vec2f(warpedDepth, warpedDepth*warpedDepth) * (1.0 - moments_in.z);
+
+    let VSMBias: f32 = vsmBias;//0.01 * 0.25;
+    let depthScale: f32 = VSMBias * exponent * warpedDepth;
+    let minVariance1: f32 = depthScale * depthScale;
+    return chebyshevUpperBound(moments, warpedDepth, minVariance1, 0.1);
+}
+// ------ VSM 16 ------
+
+fn VSM16(tex: texture_2d<f32>, texSampler: sampler, texCoords: vec2f, resolution: f32, Z: f32, vsmBias: f32, exponent: f32) -> f32 {
+    let moments: vec3f = textureSampleLevel(tex, texSampler, texCoords, 0.0).xyz;
+    return calculateEVSM(moments, Z, vsmBias, exponent);
+}
+
+fn getShadowVSM16(shadowMap: texture_2d<f32>, shadowMapSampler: sampler, shadowCoord: vec3f, shadowParams: vec4f, exponent: f32) -> f32 {
+    return VSM16(shadowMap, shadowMapSampler, shadowCoord.xy, shadowParams.x, shadowCoord.z, shadowParams.y, exponent);
+}
+
+fn getShadowSpotVSM16(shadowMap: texture_2d<f32>, shadowMapSampler: sampler, shadowCoord: vec3f, shadowParams: vec4f, exponent: f32, lightDir: vec3f) -> f32 {
+    let Z: f32 = length(lightDir) * shadowParams.w + shadowParams.z;
+    return VSM16(shadowMap, shadowMapSampler, shadowCoord.xy, shadowParams.x, Z, shadowParams.y, exponent);
+}
+
+// ------ VSM 32 ------
+
+fn VSM32(tex: texture_2d<f32>, texSampler: sampler, texCoords_in: vec2f, resolution: f32, Z: f32, vsmBias: f32, exponent: f32) -> f32 {
+
+    #ifdef CAPS_TEXTURE_FLOAT_FILTERABLE
+        var moments: vec3f = textureSampleLevel(tex, texSampler, texCoords_in, 0.0).xyz;
+    #else
+        // manual bilinear filtering
+        var pixelSize : f32 = 1.0 / resolution;
+        let texCoords: vec2f = texCoords_in - vec2f(pixelSize);
+        let s00: vec3f = textureSampleLevel(tex, texSampler, texCoords, 0.0).xyz;
+        let s10: vec3f = textureSampleLevel(tex, texSampler, texCoords + vec2f(pixelSize, 0.0), 0.0).xyz;
+        let s01: vec3f = textureSampleLevel(tex, texSampler, texCoords + vec2f(0.0, pixelSize), 0.0).xyz;
+        let s11: vec3f = textureSampleLevel(tex, texSampler, texCoords + vec2f(pixelSize), 0.0).xyz;
+        let fr: vec2f = fract(texCoords * resolution);
+        let h0: vec3f = mix(s00, s10, fr.x);
+        let h1: vec3f = mix(s01, s11, fr.x);
+        var moments: vec3f = mix(h0, h1, fr.y);
+    #endif
+
+    return calculateEVSM(moments, Z, vsmBias, exponent);
+}
+
+fn getShadowVSM32(shadowMap: texture_2d<f32>, shadowMapSampler: sampler, shadowCoord: vec3f, shadowParams: vec4f, exponent: f32) -> f32 {
+    return VSM32(shadowMap, shadowMapSampler, shadowCoord.xy, shadowParams.x, shadowCoord.z, shadowParams.y, exponent);
+}
+
+fn getShadowSpotVSM32(shadowMap: texture_2d<f32>, shadowMapSampler: sampler, shadowCoord: vec3f, shadowParams: vec4f, exponent: f32, lightDir: vec3f) -> f32 {
+    let Z: f32 = length(lightDir) * shadowParams.w + shadowParams.z;
+    return VSM32(shadowMap, shadowMapSampler, shadowCoord.xy, shadowParams.x, Z, shadowParams.y, exponent);
+}
+`;

src/scene/shader-lib/chunks/common/vert/fullscreenQuad.js

@@ -6,6 +6,6 @@ varying vec2 vUv0;
 void main(void)
 {
     gl_Position = vec4(vertex_position, 0.5, 1.0);
-    vUv0 = vertex_position.xy*0.5+0.5;
+    vUv0 = vertex_position.xy * 0.5 + 0.5;
 }
 `;

src/scene/shader-lib/chunks/lit/frag/blurVSM.js

@@ -5,24 +5,24 @@ uniform sampler2D source;
 uniform vec2 pixelOffset;
 
 #ifdef GAUSS
-uniform float weight[SAMPLES];
+    uniform float weight[{SAMPLES}];
 #endif
 
 void main(void) {
     vec3 moments = vec3(0.0);
-    vec2 uv = vUv0 - pixelOffset * (float(SAMPLES) * 0.5);
-    for (int i=0; i<SAMPLES; i++) {
+    vec2 uv = vUv0 - pixelOffset * (float({SAMPLES}) * 0.5);
+    for (int i = 0; i < {SAMPLES}; i++) {
         vec4 c = texture2D(source, uv + pixelOffset * float(i));
 
         #ifdef GAUSS
-        moments += c.xyz * weight[i];
+            moments += c.xyz * weight[i];
         #else
-        moments += c.xyz;
+            moments += c.xyz;
         #endif
     }
 
     #ifndef GAUSS
-    moments /= float(SAMPLES);
+        moments *= 1.0 / float({SAMPLES});
     #endif
 
     gl_FragColor = vec4(moments.x, moments.y, moments.z, 1.0);

src/scene/shader-lib/chunks/lit/frag/lighting/shadowEVSM.js

@@ -81,6 +81,7 @@ float getShadowVSM32(TEXTURE_ACCEPT(shadowMap), vec3 shadowCoord, vec4 shadowPar
 }
 
 float getShadowSpotVSM32(TEXTURE_ACCEPT(shadowMap), vec3 shadowCoord, vec4 shadowParams, float exponent, vec3 lightDir) {
-    return VSM32(TEXTURE_PASS(shadowMap), shadowCoord.xy, shadowParams.x, length(lightDir) * shadowParams.w + shadowParams.z, shadowParams.y, exponent);
+    float Z = length(lightDir) * shadowParams.w + shadowParams.z;
+    return VSM32(TEXTURE_PASS(shadowMap), shadowCoord.xy, shadowParams.x, Z, shadowParams.y, exponent);
 }
 `;