Skip to content

Commit d0d9b96

Browse files
mvaligurskyMartin Valigursky
mvaligursky
and
Martin Valigursky
authored
Lightmapping WGSL chunks (#7627)
Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
1 parent 3d6f577 commit d0d9b96

File tree

19 files changed

+396
-60
lines changed

19 files changed

+396
-60
lines changed

‎src/framework/lightmapper/lightmap-filters.js

+22-7
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,25 @@
11
import { createShaderFromCode } from '../../scene/shader-lib/utils.js';
22
import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
3-
import { shaderChunksLightmapper } from '../../scene/shader-lib/chunks/chunks-lightmapper.js';
3+
import { shaderChunksWGSL } from '../../scene/shader-lib/chunks-wgsl/chunks-wgsl.js';
4+
import { SEMANTIC_POSITION, SHADERLANGUAGE_GLSL, SHADERLANGUAGE_WGSL } from '../../platform/graphics/constants.js';
45

56
// size of the kernel - needs to match the constant in the shader
67
const DENOISE_FILTER_SIZE = 15;
78

89
// helper class used by lightmapper, wrapping functionality of dilate and denoise shaders
910
class LightmapFilters {
11+
shaderDilate = [];
12+
13+
shaderDenoise = [];
14+
1015
constructor(device) {
1116
this.device = device;
12-
this.shaderDilate = createShaderFromCode(device, shaderChunks.fullscreenQuadVS, shaderChunksLightmapper.dilatePS, 'lmDilate');
1317

1418
this.constantTexSource = device.scope.resolve('source');
1519

1620
this.constantPixelOffset = device.scope.resolve('pixelOffset');
1721
this.pixelOffset = new Float32Array(2);
1822

19-
// denoise is optional and gets created only when needed
20-
this.shaderDenoise = [];
2123
this.sigmas = null;
2224
this.constantSigmas = null;
2325
this.kernel = null;
@@ -39,9 +41,18 @@ class LightmapFilters {
3941

4042
const index = bakeHDR ? 0 : 1;
4143
if (!this.shaderDenoise[index]) {
44+
const wgsl = this.device.isWebGPU;
45+
const chunks = wgsl ? shaderChunksWGSL : shaderChunks;
4246
const name = `lmBilateralDeNoise-${bakeHDR ? 'hdr' : 'rgbm'}`;
43-
const define = bakeHDR ? '#define HDR\n' : '';
44-
this.shaderDenoise[index] = createShaderFromCode(this.device, shaderChunks.fullscreenQuadVS, define + shaderChunksLightmapper.bilateralDeNoisePS, name);
47+
48+
const defines = new Map();
49+
defines.set('{MSIZE}', 15);
50+
if (bakeHDR) defines.set('HDR', '');
51+
52+
this.shaderDenoise[index] = createShaderFromCode(this.device, chunks.fullscreenQuadVS, chunks.bilateralDeNoisePS, name, { vertex_position: SEMANTIC_POSITION }, undefined, {
53+
shaderLanguage: wgsl ? SHADERLANGUAGE_WGSL : SHADERLANGUAGE_GLSL,
54+
fragmentDefines: defines
55+
});
4556
this.sigmas = new Float32Array(2);
4657
this.constantSigmas = this.device.scope.resolve('sigmas');
4758
this.constantKernel = this.device.scope.resolve('kernel[0]');
@@ -63,9 +74,13 @@ class LightmapFilters {
6374
getDilate(device, bakeHDR) {
6475
const index = bakeHDR ? 0 : 1;
6576
if (!this.shaderDilate[index]) {
77+
const wgsl = this.device.isWebGPU;
78+
const chunks = wgsl ? shaderChunksWGSL : shaderChunks;
6679
const name = `lmDilate-${bakeHDR ? 'hdr' : 'rgbm'}`;
6780
const define = bakeHDR ? '#define HDR\n' : '';
68-
this.shaderDilate[index] = createShaderFromCode(device, shaderChunks.fullscreenQuadVS, define + shaderChunksLightmapper.dilatePS, name);
81+
this.shaderDilate[index] = createShaderFromCode(device, chunks.fullscreenQuadVS, define + chunks.dilatePS, name, { vertex_position: SEMANTIC_POSITION }, undefined, {
82+
shaderLanguage: wgsl ? SHADERLANGUAGE_WGSL : SHADERLANGUAGE_GLSL
83+
});
6984
}
7085
return this.shaderDilate[index];
7186
}

‎src/framework/lightmapper/lightmapper.js

+15-25
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,6 @@ import {
3030
import { MeshInstance } from '../../scene/mesh-instance.js';
3131
import { LightingParams } from '../../scene/lighting/lighting-params.js';
3232
import { WorldClusters } from '../../scene/lighting/world-clusters.js';
33-
import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
34-
import { shaderChunksLightmapper } from '../../scene/shader-lib/chunks/chunks-lightmapper.js';
3533
import { Camera } from '../../scene/camera.js';
3634
import { GraphNode } from '../../scene/graph-node.js';
3735
import { StandardMaterial } from '../../scene/materials/standard-material.js';
@@ -232,41 +230,29 @@ class Lightmapper {
232230
}
233231
}
234232

235-
createMaterialForPass(device, scene, pass, addAmbient) {
233+
createMaterialForPass(scene, pass, addAmbient) {
236234
const material = new StandardMaterial();
237235
material.name = `lmMaterial-pass:${pass}-ambient:${addAmbient}`;
238236
material.chunks.APIVersion = CHUNKAPI_1_65;
239237
material.setDefine('UV1LAYOUT', ''); // draw into UV1 texture space
238+
material.setDefine('LIT_LIGHTMAP_BAKING', '');
240239

241240
if (pass === PASS_COLOR) {
242-
let bakeLmEndChunk = shaderChunksLightmapper.bakeLmEndPS; // encode to RGBM
241+
material.setDefine('LIT_LIGHTMAP_BAKING_COLOR', '');
243242
if (addAmbient) {
244-
// diffuse light stores accumulated AO, apply contrast and brightness to it
245-
// and multiply ambient light color by the AO
246-
bakeLmEndChunk = `
247-
dDiffuseLight = ((dDiffuseLight - 0.5) * max(${scene.ambientBakeOcclusionContrast.toFixed(1)} + 1.0, 0.0)) + 0.5;
248-
dDiffuseLight += vec3(${scene.ambientBakeOcclusionBrightness.toFixed(1)});
249-
dDiffuseLight = saturate(dDiffuseLight);
250-
dDiffuseLight *= dAmbientLight;
251-
${bakeLmEndChunk}
252-
`;
243+
material.setDefine('LIT_LIGHTMAP_BAKING_ADD_AMBIENT', '');
253244
} else {
254245
material.ambient = new Color(0, 0, 0); // don't bake ambient
255246
}
256-
material.chunks.basePS = shaderChunks.basePS + (this.bakeHDR ? '' : '\n#define LIGHTMAP_RGBM\n');
257-
material.chunks.endPS = bakeLmEndChunk;
247+
248+
if (!this.bakeHDR) material.setDefine('LIGHTMAP_RGBM', '');
249+
258250
material.lightMap = this.blackTex;
259251
} else {
260-
material.chunks.basePS = `
261-
#define STD_LIGHTMAP_DIR
262-
${shaderChunks.basePS}
263-
uniform float bakeDir;
264-
`;
265-
material.chunks.endPS = shaderChunksLightmapper.bakeDirLmEndPS;
252+
material.setDefine('LIT_LIGHTMAP_BAKING_DIR', '');
253+
material.setDefine('STD_LIGHTMAP_DIR', '');
266254
}
267255

268-
// avoid writing unrelated things to alpha
269-
material.chunks.outputAlphaPS = '\n';
270256
material.cull = CULLFACE_NONE;
271257
material.forceUv1 = true; // provide data to xformUv1
272258
material.update();
@@ -277,13 +263,13 @@ class Lightmapper {
277263
createMaterials(device, scene, passCount) {
278264
for (let pass = 0; pass < passCount; pass++) {
279265
if (!this.passMaterials[pass]) {
280-
this.passMaterials[pass] = this.createMaterialForPass(device, scene, pass, false);
266+
this.passMaterials[pass] = this.createMaterialForPass(scene, pass, false);
281267
}
282268
}
283269

284270
// material used on last render of ambient light to multiply accumulated AO in lightmap by ambient light
285271
if (!this.ambientAOMaterial) {
286-
this.ambientAOMaterial = this.createMaterialForPass(device, scene, 0, true);
272+
this.ambientAOMaterial = this.createMaterialForPass(scene, 0, true);
287273
this.ambientAOMaterial.onUpdateShader = function (options) {
288274
// mark LM as without ambient, to add it
289275
options.litOptions.lightMapWithoutAmbient = true;
@@ -694,6 +680,10 @@ class Lightmapper {
694680

695681
// apply scene settings
696682
this.renderer.setSceneConstants();
683+
684+
// uniforms
685+
this.device.scope.resolve('ambientBakeOcclusionContrast').setValue(this.scene.ambientBakeOcclusionContrast);
686+
this.device.scope.resolve('ambientBakeOcclusionBrightness').setValue(this.scene.ambientBakeOcclusionBrightness);
697687
}
698688

699689
restoreScene() {

‎src/index.js

-1
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,6 @@ export { createShader, createShaderFromCode } from './scene/shader-lib/utils.js'
227227
export { LitShaderOptions } from './scene/shader-lib/programs/lit-shader-options.js';
228228
export { ProgramLibrary } from './scene/shader-lib/program-library.js';
229229
export { shaderChunks } from './scene/shader-lib/chunks/chunks.js';
230-
export { shaderChunksLightmapper } from './scene/shader-lib/chunks/chunks-lightmapper.js';
231230
export { ChunkUtils } from './scene/shader-lib/chunk-utils.js';
232231

233232
// SCENE / SKY

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

+12-4
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@ import ambientPS from './lit/frag/ambient.js';
33
import aoPS from './standard/frag/ao.js';
44
import aoDiffuseOccPS from './lit/frag/aoDiffuseOcc.js';
55
import aoSpecOccPS from './lit/frag/aoSpecOcc.js';
6+
import bakeDirLmEndPS from './lightmapper/frag/bakeDirLmEnd.js';
7+
import bakeLmEndPS from './lightmapper/frag/bakeLmEnd.js';
68
import basePS from './lit/frag/base.js';
79
import baseNineSlicedPS from './lit/frag/baseNineSliced.js';
810
import baseNineSlicedTiledPS from './lit/frag/baseNineSlicedTiled.js';
911
import bayerPS from './common/frag/bayer.js';
12+
import bilateralDeNoisePS from './lightmapper/frag/bilateralDeNoise.js';
1013
import blurVSMPS from './lit/frag/blurVSM.js';
1114
import clearCoatPS from './standard/frag/clearCoat.js';
1215
import clearCoatGlossPS from './standard/frag/clearCoatGloss.js';
@@ -26,6 +29,7 @@ import debugOutputPS from './lit/frag/debug-output.js';
2629
import debugProcessFrontendPS from './lit/frag/debug-process-frontend.js';
2730
import decodePS from './common/frag/decode.js';
2831
import detailModesPS from './standard/frag/detailModes.js';
32+
import dilatePS from './lightmapper/frag/dilate.js';
2933
import diffusePS from './standard/frag/diffuse.js';
3034
import emissivePS from './standard/frag/emissive.js';
3135
import encodePS from './common/frag/encode.js';
@@ -65,8 +69,8 @@ import lightEvaluationPS from './lit/frag/lighting/lightEvaluation.js';
6569
import lightFunctionLightPS from './lit/frag/lighting/lightFunctionLight.js';
6670
import lightFunctionShadowPS from './lit/frag/lighting/lightFunctionShadow.js';
6771
import lightingPS from './lit/frag/lighting/lighting.js';
68-
// import lightmapAddPS from './lit/frag/lightmapAdd.js';
69-
// import lightmapPS from './standard/frag/lightmap.js';
72+
import lightmapAddPS from './lit/frag/lightmapAdd.js';
73+
import lightmapPS from './standard/frag/lightmap.js';
7074
import lightSpecularAnisoGGXPS from './lit/frag/lightSpecularAnisoGGX.js';
7175
import lightSpecularBlinnPS from './lit/frag/lightSpecularBlinn.js';
7276
import lightSheenPS from './lit/frag/lightSheen.js';
@@ -211,10 +215,13 @@ const shaderChunksWGSL = {
211215
aoPS,
212216
aoDiffuseOccPS,
213217
aoSpecOccPS,
218+
bakeDirLmEndPS,
219+
bakeLmEndPS,
214220
basePS,
215221
baseNineSlicedPS,
216222
baseNineSlicedTiledPS,
217223
bayerPS,
224+
bilateralDeNoisePS,
218225
blurVSMPS,
219226
clearCoatPS,
220227
clearCoatGlossPS,
@@ -233,6 +240,7 @@ const shaderChunksWGSL = {
233240
debugOutputPS,
234241
debugProcessFrontendPS,
235242
detailModesPS,
243+
dilatePS,
236244
diffusePS,
237245
decodePS,
238246
emissivePS,
@@ -274,8 +282,8 @@ const shaderChunksWGSL = {
274282
lightFunctionLightPS,
275283
lightFunctionShadowPS,
276284
lightingPS,
277-
// lightmapAddPS,
278-
// lightmapPS,
285+
lightmapAddPS,
286+
lightmapPS,
279287
lightSpecularAnisoGGXPS,
280288
lightSpecularBlinnPS,
281289
lightSheenPS,

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

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
export default /* wgsl */`
2+
let dirLm = textureSample(texture_dirLightMap, texture_dirLightMapSampler, vUv1);
3+
4+
if (uniform.bakeDir > 0.5) {
5+
if (dAtten > 0.00001) {
6+
let unpacked_dir = dirLm.xyz * 2.0 - vec3f(1.0);
7+
dAtten = clamp(dAtten, 0.0, 1.0);
8+
let combined_dir = dLightDirNormW.xyz * dAtten + unpacked_dir * dirLm.w;
9+
let finalRgb = normalize(combined_dir) * 0.5 + vec3f(0.5);
10+
let finalA = max(dirLm.w + dAtten, 1.0 / 255.0);
11+
output.color = vec4f(finalRgb, finalA);
12+
} else {
13+
output.color = dirLm;
14+
}
15+
} else {
16+
let alpha_min = select(0.0, 1.0 / 255.0, dAtten > 0.00001);
17+
let finalA = max(dirLm.w, alpha_min);
18+
output.color = vec4f(dirLm.rgb, finalA);
19+
}
20+
`;

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

+26
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
export default /* wgsl */`
2+
3+
#ifdef LIT_LIGHTMAP_BAKING_ADD_AMBIENT
4+
// diffuse light stores accumulated AO, apply contrast and brightness to it
5+
// and multiply ambient light color by the AO
6+
dDiffuseLight = ((dDiffuseLight - 0.5) * max(uniform.ambientBakeOcclusionContrast + 1.0, 0.0)) + 0.5;
7+
dDiffuseLight = dDiffuseLight + vec3f(uniform.ambientBakeOcclusionBrightness);
8+
dDiffuseLight = saturate3(dDiffuseLight);
9+
dDiffuseLight = dDiffuseLight * dAmbientLight;
10+
#endif
11+
12+
#ifdef LIGHTMAP_RGBM
13+
// encode to RGBM
14+
var temp_color_rgbm = vec4f(dDiffuseLight, 1.0);
15+
temp_color_rgbm = vec4f(pow(temp_color_rgbm.rgb, vec3f(0.5)), temp_color_rgbm.a);
16+
temp_color_rgbm = vec4f(temp_color_rgbm.rgb / 8.0, temp_color_rgbm.a);
17+
let max_g_b = max(temp_color_rgbm.g, max(temp_color_rgbm.b, 1.0 / 255.0));
18+
let max_rgb = max(temp_color_rgbm.r, max_g_b);
19+
temp_color_rgbm.a = clamp(max_rgb, 0.0, 1.0);
20+
temp_color_rgbm.a = ceil(temp_color_rgbm.a * 255.0) / 255.0;
21+
temp_color_rgbm = vec4f(temp_color_rgbm.rgb / temp_color_rgbm.a, temp_color_rgbm.a);
22+
output.color = temp_color_rgbm;
23+
#else
24+
output.color = vec4f(dDiffuseLight, 1.0);
25+
#endif
26+
`;

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

+120
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
export default /* wgsl */`
2+
// bilateral filter, based on https://www.shadertoy.com/view/4dfGDH# and
3+
// http://people.csail.mit.edu/sparis/bf_course/course_notes.pdf
4+
5+
// A bilateral filter is a non-linear, edge-preserving, and noise-reducing smoothing filter for images.
6+
// It replaces the intensity of each pixel with a weighted average of intensity values from nearby pixels.
7+
// This weight can be based on a Gaussian distribution. Crucially, the weights depend not only on
8+
// Euclidean distance of pixels, but also on the radiometric differences (e.g., range differences, such
9+
// as color intensity, depth distance, etc.). This preserves sharp edges.
10+
11+
fn normpdf3(v: vec3f, sigma: f32) -> f32 {
12+
return 0.39894 * exp(-0.5 * dot(v, v) / (sigma * sigma)) / sigma;
13+
}
14+
15+
fn decodeRGBM(rgbm: vec4f) -> vec3f {
16+
let color = (8.0 * rgbm.a) * rgbm.rgb;
17+
return color * color;
18+
}
19+
20+
fn saturate(x: f32) -> f32 {
21+
return clamp(x, 0.0, 1.0);
22+
}
23+
24+
fn encodeRGBM(color: vec3f) -> vec4f {
25+
var encoded: vec4f;
26+
let rgb_processed = pow(color.rgb, vec3f(0.5)) * (1.0 / 8.0);
27+
encoded = vec4f(rgb_processed, 0.0);
28+
29+
let max_g_b = max( encoded.g, max( encoded.b, 1.0 / 255.0 ) );
30+
let max_rgb = max( encoded.r, max_g_b );
31+
encoded.a = clamp(max_rgb, 0.0, 1.0);
32+
encoded.a = ceil(encoded.a * 255.0) / 255.0;
33+
34+
encoded = vec4f(encoded.rgb / encoded.a, encoded.a);
35+
return encoded;
36+
}
37+
38+
fn decode(pixel: vec4f) -> vec3f {
39+
#if HDR
40+
return pixel.rgb;
41+
#else
42+
return decodeRGBM(pixel);
43+
#endif
44+
}
45+
46+
fn isUsed(pixel: vec4f) -> bool {
47+
#if HDR
48+
return any(pixel.rgb > vec3f(0.0));
49+
#else
50+
return pixel.a > 0.0;
51+
#endif
52+
}
53+
54+
varying vUv0: vec2f;
55+
var source: texture_2d<f32>;
56+
var sourceSampler: sampler;
57+
uniform kernel: array<f32, {MSIZE}>;
58+
uniform pixelOffset: vec2f;
59+
uniform sigmas: vec2f;
60+
uniform bZnorm: f32;
61+
62+
@fragment
63+
fn fragmentMain(input: FragmentInput) -> FragmentOutput {
64+
var output: FragmentOutput;
65+
66+
let pixel = textureSampleLevel(source, sourceSampler, input.vUv0, 0.0);
67+
68+
// lightmap specific optimization - skip pixels that were not baked
69+
// this also allows dilate filter that work on the output of this to work correctly, as it depends on .a being zero
70+
// to dilate, which the following blur filter would otherwise modify
71+
if (!isUsed(pixel)) {
72+
output.color = pixel;
73+
return output;
74+
}
75+
76+
// range sigma - controls blurriness based on a pixel distance
77+
let sigma = uniform.sigmas.x;
78+
79+
// domain sigma - controls blurriness based on a pixel similarity (to preserve edges)
80+
let bSigma = uniform.sigmas.y;
81+
82+
let pixelHdr = decode(pixel);
83+
var accumulatedHdr = vec3f(0.0);
84+
var accumulatedFactor = 0.000001; // avoid division by zero
85+
86+
// read out the texels
87+
const kSize = ({MSIZE} - 1) / 2;
88+
for (var i: i32 = -kSize; i <= kSize; i = i + 1) {
89+
for (var j: i32 = -kSize; j <= kSize; j = j + 1) {
90+
91+
// sample the pixel with offset
92+
let coord = input.vUv0 + vec2f(f32(i), f32(j)) * uniform.pixelOffset;
93+
let pix = textureSampleLevel(source, sourceSampler, coord, 0.0);
94+
95+
// lightmap - only use baked pixels
96+
if (isUsed(pix)) {
97+
let hdr = decode(pix);
98+
99+
// bilateral factors
100+
var factor = uniform.kernel[u32(kSize + j)].element * uniform.kernel[u32(kSize + i)].element;
101+
factor = factor * normpdf3(hdr - pixelHdr, bSigma) * uniform.bZnorm;
102+
103+
// accumulate
104+
accumulatedHdr = accumulatedHdr + factor * hdr;
105+
accumulatedFactor = accumulatedFactor + factor;
106+
}
107+
}
108+
}
109+
110+
let finalHDR = accumulatedHdr / accumulatedFactor;
111+
112+
#if HDR
113+
output.color = vec4f(finalHDR, 1.0);
114+
#else
115+
output.color = encodeRGBM(finalHDR);
116+
#endif
117+
118+
return output;
119+
}
120+
`;

0 commit comments

Comments
 (0)