Lightmapping WGSL chunks

src/framework/lightmapper/lightmap-filters.js

@@ -1,23 +1,25 @@
 import { createShaderFromCode } from '../../scene/shader-lib/utils.js';
 import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
-import { shaderChunksLightmapper } from '../../scene/shader-lib/chunks/chunks-lightmapper.js';
+import { shaderChunksWGSL } from '../../scene/shader-lib/chunks-wgsl/chunks-wgsl.js';
+import { SEMANTIC_POSITION, SHADERLANGUAGE_GLSL, SHADERLANGUAGE_WGSL } from '../../platform/graphics/constants.js';
 
 // size of the kernel - needs to match the constant in the shader
 const DENOISE_FILTER_SIZE = 15;
 
 // helper class used by lightmapper, wrapping functionality of dilate and denoise shaders
 class LightmapFilters {
+    shaderDilate = [];
+
+    shaderDenoise = [];
+
     constructor(device) {
         this.device = device;
-        this.shaderDilate = createShaderFromCode(device, shaderChunks.fullscreenQuadVS, shaderChunksLightmapper.dilatePS, 'lmDilate');
 
         this.constantTexSource = device.scope.resolve('source');
 
         this.constantPixelOffset = device.scope.resolve('pixelOffset');
         this.pixelOffset = new Float32Array(2);
 
-        // denoise is optional and gets created only when needed
-        this.shaderDenoise = [];
         this.sigmas = null;
         this.constantSigmas = null;
         this.kernel = null;
@@ -39,9 +41,18 @@ class LightmapFilters {
 
         const index = bakeHDR ? 0 : 1;
         if (!this.shaderDenoise[index]) {
+            const wgsl = this.device.isWebGPU;
+            const chunks = wgsl ? shaderChunksWGSL : shaderChunks;
             const name = `lmBilateralDeNoise-${bakeHDR ? 'hdr' : 'rgbm'}`;
-            const define = bakeHDR ? '#define HDR\n' : '';
-            this.shaderDenoise[index] = createShaderFromCode(this.device, shaderChunks.fullscreenQuadVS, define + shaderChunksLightmapper.bilateralDeNoisePS, name);
+
+            const defines = new Map();
+            defines.set('{MSIZE}', 15);
+            if (bakeHDR) defines.set('HDR', '');
+
+            this.shaderDenoise[index] = createShaderFromCode(this.device, chunks.fullscreenQuadVS, chunks.bilateralDeNoisePS, name, { vertex_position: SEMANTIC_POSITION }, undefined, {
+                shaderLanguage: wgsl ? SHADERLANGUAGE_WGSL : SHADERLANGUAGE_GLSL,
+                fragmentDefines: defines
+            });
             this.sigmas = new Float32Array(2);
             this.constantSigmas = this.device.scope.resolve('sigmas');
             this.constantKernel = this.device.scope.resolve('kernel[0]');
@@ -63,9 +74,13 @@ class LightmapFilters {
     getDilate(device, bakeHDR) {
         const index = bakeHDR ? 0 : 1;
         if (!this.shaderDilate[index]) {
+            const wgsl = this.device.isWebGPU;
+            const chunks = wgsl ? shaderChunksWGSL : shaderChunks;
             const name = `lmDilate-${bakeHDR ? 'hdr' : 'rgbm'}`;
             const define = bakeHDR ? '#define HDR\n' : '';
-            this.shaderDilate[index] = createShaderFromCode(device, shaderChunks.fullscreenQuadVS, define + shaderChunksLightmapper.dilatePS, name);
+            this.shaderDilate[index] = createShaderFromCode(device, chunks.fullscreenQuadVS, define + chunks.dilatePS, name, { vertex_position: SEMANTIC_POSITION }, undefined, {
+                shaderLanguage: wgsl ? SHADERLANGUAGE_WGSL : SHADERLANGUAGE_GLSL
+            });
         }
         return this.shaderDilate[index];
     }

src/framework/lightmapper/lightmapper.js

@@ -30,8 +30,6 @@ import {
 import { MeshInstance } from '../../scene/mesh-instance.js';
 import { LightingParams } from '../../scene/lighting/lighting-params.js';
 import { WorldClusters } from '../../scene/lighting/world-clusters.js';
-import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
-import { shaderChunksLightmapper } from '../../scene/shader-lib/chunks/chunks-lightmapper.js';
 import { Camera } from '../../scene/camera.js';
 import { GraphNode } from '../../scene/graph-node.js';
 import { StandardMaterial } from '../../scene/materials/standard-material.js';
@@ -232,41 +230,29 @@ class Lightmapper {
         }
     }
 
-    createMaterialForPass(device, scene, pass, addAmbient) {
+    createMaterialForPass(scene, pass, addAmbient) {
         const material = new StandardMaterial();
         material.name = `lmMaterial-pass:${pass}-ambient:${addAmbient}`;
         material.chunks.APIVersion = CHUNKAPI_1_65;
         material.setDefine('UV1LAYOUT', '');    // draw into UV1 texture space
+        material.setDefine('LIT_LIGHTMAP_BAKING', '');
 
         if (pass === PASS_COLOR) {
-            let bakeLmEndChunk = shaderChunksLightmapper.bakeLmEndPS; // encode to RGBM
+            material.setDefine('LIT_LIGHTMAP_BAKING_COLOR', '');
             if (addAmbient) {
-                // diffuse light stores accumulated AO, apply contrast and brightness to it
-                // and multiply ambient light color by the AO
-                bakeLmEndChunk = `
-                    dDiffuseLight = ((dDiffuseLight - 0.5) * max(${scene.ambientBakeOcclusionContrast.toFixed(1)} + 1.0, 0.0)) + 0.5;
-                    dDiffuseLight += vec3(${scene.ambientBakeOcclusionBrightness.toFixed(1)});
-                    dDiffuseLight = saturate(dDiffuseLight);
-                    dDiffuseLight *= dAmbientLight;
-                    ${bakeLmEndChunk}
-                `;
+                material.setDefine('LIT_LIGHTMAP_BAKING_ADD_AMBIENT', '');
             } else {
                 material.ambient = new Color(0, 0, 0);    // don't bake ambient
             }
-            material.chunks.basePS = shaderChunks.basePS + (this.bakeHDR ? '' : '\n#define LIGHTMAP_RGBM\n');
-            material.chunks.endPS = bakeLmEndChunk;
+
+            if (!this.bakeHDR) material.setDefine('LIGHTMAP_RGBM', '');
+
             material.lightMap = this.blackTex;
         } else {
-            material.chunks.basePS = `
-                #define STD_LIGHTMAP_DIR
-                ${shaderChunks.basePS}
-                uniform float bakeDir;
-            `;
-            material.chunks.endPS = shaderChunksLightmapper.bakeDirLmEndPS;
+            material.setDefine('LIT_LIGHTMAP_BAKING_DIR', '');
+            material.setDefine('STD_LIGHTMAP_DIR', '');
         }
 
-        // avoid writing unrelated things to alpha
-        material.chunks.outputAlphaPS = '\n';
         material.cull = CULLFACE_NONE;
         material.forceUv1 = true; // provide data to xformUv1
         material.update();
@@ -277,13 +263,13 @@ class Lightmapper {
     createMaterials(device, scene, passCount) {
         for (let pass = 0; pass < passCount; pass++) {
             if (!this.passMaterials[pass]) {
-                this.passMaterials[pass] = this.createMaterialForPass(device, scene, pass, false);
+                this.passMaterials[pass] = this.createMaterialForPass(scene, pass, false);
             }
         }
 
         // material used on last render of ambient light to multiply accumulated AO in lightmap by ambient light
         if (!this.ambientAOMaterial) {
-            this.ambientAOMaterial = this.createMaterialForPass(device, scene, 0, true);
+            this.ambientAOMaterial = this.createMaterialForPass(scene, 0, true);
             this.ambientAOMaterial.onUpdateShader = function (options) {
                 // mark LM as without ambient, to add it
                 options.litOptions.lightMapWithoutAmbient = true;
@@ -694,6 +680,10 @@ class Lightmapper {
 
         // apply scene settings
         this.renderer.setSceneConstants();
+
+        // uniforms
+        this.device.scope.resolve('ambientBakeOcclusionContrast').setValue(this.scene.ambientBakeOcclusionContrast);
+        this.device.scope.resolve('ambientBakeOcclusionBrightness').setValue(this.scene.ambientBakeOcclusionBrightness);
     }
 
     restoreScene() {

src/index.js

@@ -227,7 +227,6 @@ export { createShader, createShaderFromCode } from './scene/shader-lib/utils.js'
 export { LitShaderOptions } from './scene/shader-lib/programs/lit-shader-options.js';
 export { ProgramLibrary } from './scene/shader-lib/program-library.js';
 export { shaderChunks } from './scene/shader-lib/chunks/chunks.js';
-export { shaderChunksLightmapper } from './scene/shader-lib/chunks/chunks-lightmapper.js';
 export { ChunkUtils } from './scene/shader-lib/chunk-utils.js';
 
 // SCENE / SKY

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

@@ -3,10 +3,13 @@ import ambientPS from './lit/frag/ambient.js';
 import aoPS from './standard/frag/ao.js';
 import aoDiffuseOccPS from './lit/frag/aoDiffuseOcc.js';
 import aoSpecOccPS from './lit/frag/aoSpecOcc.js';
+import bakeDirLmEndPS from './lightmapper/frag/bakeDirLmEnd.js';
+import bakeLmEndPS from './lightmapper/frag/bakeLmEnd.js';
 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 bilateralDeNoisePS from './lightmapper/frag/bilateralDeNoise.js';
 import blurVSMPS from './lit/frag/blurVSM.js';
 import clearCoatPS from './standard/frag/clearCoat.js';
 import clearCoatGlossPS from './standard/frag/clearCoatGloss.js';
@@ -26,6 +29,7 @@ import debugOutputPS from './lit/frag/debug-output.js';
 import debugProcessFrontendPS from './lit/frag/debug-process-frontend.js';
 import decodePS from './common/frag/decode.js';
 import detailModesPS from './standard/frag/detailModes.js';
+import dilatePS from './lightmapper/frag/dilate.js';
 import diffusePS from './standard/frag/diffuse.js';
 import emissivePS from './standard/frag/emissive.js';
 import encodePS from './common/frag/encode.js';
@@ -65,8 +69,8 @@ import lightEvaluationPS from './lit/frag/lighting/lightEvaluation.js';
 import lightFunctionLightPS from './lit/frag/lighting/lightFunctionLight.js';
 import lightFunctionShadowPS from './lit/frag/lighting/lightFunctionShadow.js';
 import lightingPS from './lit/frag/lighting/lighting.js';
-// import lightmapAddPS from './lit/frag/lightmapAdd.js';
-// import lightmapPS from './standard/frag/lightmap.js';
+import lightmapAddPS from './lit/frag/lightmapAdd.js';
+import lightmapPS from './standard/frag/lightmap.js';
 import lightSpecularAnisoGGXPS from './lit/frag/lightSpecularAnisoGGX.js';
 import lightSpecularBlinnPS from './lit/frag/lightSpecularBlinn.js';
 import lightSheenPS from './lit/frag/lightSheen.js';
@@ -211,10 +215,13 @@ const shaderChunksWGSL = {
     aoPS,
     aoDiffuseOccPS,
     aoSpecOccPS,
+    bakeDirLmEndPS,
+    bakeLmEndPS,
     basePS,
     baseNineSlicedPS,
     baseNineSlicedTiledPS,
     bayerPS,
+    bilateralDeNoisePS,
     blurVSMPS,
     clearCoatPS,
     clearCoatGlossPS,
@@ -233,6 +240,7 @@ const shaderChunksWGSL = {
     debugOutputPS,
     debugProcessFrontendPS,
     detailModesPS,
+    dilatePS,
     diffusePS,
     decodePS,
     emissivePS,
@@ -274,8 +282,8 @@ const shaderChunksWGSL = {
     lightFunctionLightPS,
     lightFunctionShadowPS,
     lightingPS,
-    // lightmapAddPS,
-    // lightmapPS,
+    lightmapAddPS,
+    lightmapPS,
     lightSpecularAnisoGGXPS,
     lightSpecularBlinnPS,
     lightSheenPS,

src/scene/shader-lib/chunks-wgsl/lightmapper/frag/bakeDirLmEnd.js

@@ -0,0 +1,20 @@
+export default /* wgsl */`
+    let dirLm = textureSample(texture_dirLightMap, texture_dirLightMapSampler, vUv1);
+
+    if (uniform.bakeDir > 0.5) {
+        if (dAtten > 0.00001) {
+            let unpacked_dir = dirLm.xyz * 2.0 - vec3f(1.0);
+            dAtten = clamp(dAtten, 0.0, 1.0);
+            let combined_dir = dLightDirNormW.xyz * dAtten + unpacked_dir * dirLm.w;
+            let finalRgb = normalize(combined_dir) * 0.5 + vec3f(0.5);
+            let finalA = max(dirLm.w + dAtten, 1.0 / 255.0);
+            output.color = vec4f(finalRgb, finalA);
+        } else {
+            output.color = dirLm;
+        }
+    } else {
+        let alpha_min = select(0.0, 1.0 / 255.0, dAtten > 0.00001);
+        let finalA = max(dirLm.w, alpha_min);
+        output.color = vec4f(dirLm.rgb, finalA);
+    }
+`;

src/scene/shader-lib/chunks-wgsl/lightmapper/frag/bakeLmEnd.js

@@ -0,0 +1,26 @@
+export default /* wgsl */`
+
+#ifdef LIT_LIGHTMAP_BAKING_ADD_AMBIENT
+    // diffuse light stores accumulated AO, apply contrast and brightness to it
+    // and multiply ambient light color by the AO
+    dDiffuseLight = ((dDiffuseLight - 0.5) * max(uniform.ambientBakeOcclusionContrast + 1.0, 0.0)) + 0.5;
+    dDiffuseLight = dDiffuseLight + vec3f(uniform.ambientBakeOcclusionBrightness);
+    dDiffuseLight = saturate3(dDiffuseLight);
+    dDiffuseLight = dDiffuseLight * dAmbientLight;
+#endif
+
+#ifdef LIGHTMAP_RGBM
+    // encode to RGBM
+    var temp_color_rgbm = vec4f(dDiffuseLight, 1.0);
+    temp_color_rgbm = vec4f(pow(temp_color_rgbm.rgb, vec3f(0.5)), temp_color_rgbm.a);
+    temp_color_rgbm = vec4f(temp_color_rgbm.rgb / 8.0, temp_color_rgbm.a);
+    let max_g_b = max(temp_color_rgbm.g, max(temp_color_rgbm.b, 1.0 / 255.0));
+    let max_rgb = max(temp_color_rgbm.r, max_g_b);
+    temp_color_rgbm.a = clamp(max_rgb, 0.0, 1.0);
+    temp_color_rgbm.a = ceil(temp_color_rgbm.a * 255.0) / 255.0;
+    temp_color_rgbm = vec4f(temp_color_rgbm.rgb / temp_color_rgbm.a, temp_color_rgbm.a);
+    output.color = temp_color_rgbm;
+#else
+    output.color = vec4f(dDiffuseLight, 1.0);
+#endif
+`;

src/scene/shader-lib/chunks-wgsl/lightmapper/frag/bilateralDeNoise.js

@@ -0,0 +1,120 @@
+export default /* wgsl */`
+// bilateral filter, based on https://www.shadertoy.com/view/4dfGDH# and
+// http://people.csail.mit.edu/sparis/bf_course/course_notes.pdf
+
+// A bilateral filter is a non-linear, edge-preserving, and noise-reducing smoothing filter for images.
+// It replaces the intensity of each pixel with a weighted average of intensity values from nearby pixels.
+// This weight can be based on a Gaussian distribution. Crucially, the weights depend not only on
+// Euclidean distance of pixels, but also on the radiometric differences (e.g., range differences, such
+// as color intensity, depth distance, etc.). This preserves sharp edges.
+
+fn normpdf3(v: vec3f, sigma: f32) -> f32 {
+    return 0.39894 * exp(-0.5 * dot(v, v) / (sigma * sigma)) / sigma;
+}
+
+fn decodeRGBM(rgbm: vec4f) -> vec3f {
+    let color = (8.0 * rgbm.a) * rgbm.rgb;
+    return color * color;
+}
+
+fn saturate(x: f32) -> f32 {
+    return clamp(x, 0.0, 1.0);
+}
+
+fn encodeRGBM(color: vec3f) -> vec4f {
+    var encoded: vec4f;
+    let rgb_processed = pow(color.rgb, vec3f(0.5)) * (1.0 / 8.0);
+    encoded = vec4f(rgb_processed, 0.0);
+
+    let max_g_b = max( encoded.g, max( encoded.b, 1.0 / 255.0 ) );
+    let max_rgb = max( encoded.r, max_g_b );
+    encoded.a = clamp(max_rgb, 0.0, 1.0);
+    encoded.a = ceil(encoded.a * 255.0) / 255.0;
+
+    encoded = vec4f(encoded.rgb / encoded.a, encoded.a);
+    return encoded;
+}
+
+fn decode(pixel: vec4f) -> vec3f {
+    #if HDR
+        return pixel.rgb;
+    #else
+        return decodeRGBM(pixel);
+    #endif
+}
+
+fn isUsed(pixel: vec4f) -> bool {
+    #if HDR
+        return any(pixel.rgb > vec3f(0.0));
+    #else
+        return pixel.a > 0.0;
+    #endif
+}
+
+varying vUv0: vec2f;
+var source: texture_2d<f32>;
+var sourceSampler: sampler;
+uniform kernel: array<f32, {MSIZE}>;
+uniform pixelOffset: vec2f;
+uniform sigmas: vec2f;
+uniform bZnorm: f32;
+
+@fragment
+fn fragmentMain(input: FragmentInput) -> FragmentOutput {
+    var output: FragmentOutput;
+
+    let pixel = textureSampleLevel(source, sourceSampler, input.vUv0, 0.0);
+
+    // lightmap specific optimization - skip pixels that were not baked
+    // this also allows dilate filter that work on the output of this to work correctly, as it depends on .a being zero
+    // to dilate, which the following blur filter would otherwise modify
+    if (!isUsed(pixel)) {
+        output.color = pixel;
+        return output;
+    }
+
+    // range sigma - controls blurriness based on a pixel distance
+    let sigma = uniform.sigmas.x;
+
+    // domain sigma - controls blurriness based on a pixel similarity (to preserve edges)
+    let bSigma = uniform.sigmas.y;
+
+    let pixelHdr = decode(pixel);
+    var accumulatedHdr = vec3f(0.0);
+    var accumulatedFactor = 0.000001;  // avoid division by zero
+
+    // read out the texels
+    const kSize = ({MSIZE} - 1) / 2;
+    for (var i: i32 = -kSize; i <= kSize; i = i + 1) {
+        for (var j: i32 = -kSize; j <= kSize; j = j + 1) {
+
+            // sample the pixel with offset
+            let coord = input.vUv0 + vec2f(f32(i), f32(j)) * uniform.pixelOffset;
+            let pix = textureSampleLevel(source, sourceSampler, coord, 0.0);
+
+            // lightmap - only use baked pixels
+            if (isUsed(pix)) {
+                let hdr = decode(pix);
+
+                // bilateral factors
+                var factor = uniform.kernel[u32(kSize + j)].element * uniform.kernel[u32(kSize + i)].element;
+                factor = factor * normpdf3(hdr - pixelHdr, bSigma) * uniform.bZnorm;
+
+                // accumulate
+                accumulatedHdr = accumulatedHdr + factor * hdr;
+                accumulatedFactor = accumulatedFactor + factor;
+            }
+        }
+    }
+
+    let finalHDR = accumulatedHdr / accumulatedFactor;
+
+    #if HDR
+        output.color = vec4f(finalHDR, 1.0);
+    #else
+        output.color = encodeRGBM(finalHDR);
+    #endif
+
+    return output;
+}
+`;