Skip to content

Commit 8c15533

Browse files
mvaligurskyMartin Valigursky
mvaligursky
and
Martin Valigursky
authored
Add heatmap debug mode for compute GSplat rasterizer (#8599)
Adds GSPLAT_DEBUG_HEATMAP debug rendering mode that visualizes the average number of splats processed per pixel in each tile before the transmittance early-out, displayed as a blue-to-red color ramp at tile resolution using workgroup-shared atomics. Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
1 parent 7be3747 commit 8c15533

File tree

8 files changed

+119
-49
lines changed

8 files changed

+119
-49
lines changed

‎examples/src/examples/gaussian-splatting/lod-streaming-sh.controls.mjs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ export const controls = ({ observer, ReactPCUI, React, jsx, fragment }) => {
3535
options: [
3636
{ v: 0, t: 'None' },
3737
{ v: 1, t: 'LOD' },
38-
{ v: 2, t: 'SH Update' }
38+
{ v: 2, t: 'SH Update' },
39+
{ v: 3, t: 'Heatmap' }
3940
]
4041
})
4142
),

‎examples/src/examples/gaussian-splatting/lod-streaming.controls.mjs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,8 @@ export const controls = ({ observer, ReactPCUI, React, jsx, fragment }) => {
227227
options: [
228228
{ v: 0, t: 'None' },
229229
{ v: 1, t: 'LOD' },
230-
{ v: 2, t: 'SH Update' }
230+
{ v: 2, t: 'SH Update' },
231+
{ v: 3, t: 'Heatmap' }
231232
]
232233
})
233234
),

‎examples/src/examples/gaussian-splatting/world.controls.mjs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,8 @@ export const controls = ({ observer, ReactPCUI, React, jsx, fragment }) => {
5353
options: [
5454
{ v: 0, t: 'None' },
5555
{ v: 1, t: 'LOD' },
56-
{ v: 2, t: 'SH Update' }
56+
{ v: 2, t: 'SH Update' },
57+
{ v: 3, t: 'Heatmap' }
5758
]
5859
})
5960
),

‎src/scene/constants.js‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1271,3 +1271,13 @@ export const GSPLAT_DEBUG_LOD = 1;
12711271
* @category Graphics
12721272
*/
12731273
export const GSPLAT_DEBUG_SH_UPDATE = 2;
1274+
1275+
/**
1276+
* Debug heatmap rendering for the compute rasterizer. Visualizes the average number of splats
1277+
* processed per pixel in each tile as a blue-to-red color ramp. Only supported with
1278+
* {@link GSPLAT_RENDERER_COMPUTE}.
1279+
*
1280+
* @type {number}
1281+
* @category Graphics
1282+
*/
1283+
export const GSPLAT_DEBUG_HEATMAP = 3;

‎src/scene/gsplat-unified/gsplat-compute-local-renderer.js‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import {
1414
UNIFORMTYPE_MAT4,
1515
UNIFORMTYPE_UINT
1616
} from '../../platform/graphics/constants.js';
17-
import { GSPLAT_FORWARD, PROJECTION_ORTHOGRAPHIC, FOG_NONE } from '../constants.js';
17+
import { GSPLAT_FORWARD, PROJECTION_ORTHOGRAPHIC, FOG_NONE, GSPLAT_DEBUG_HEATMAP } from '../constants.js';
1818
import { Debug } from '../../core/debug.js';
1919
import { Color } from '../../core/math/color.js';
2020
import { Mat4 } from '../../core/math/mat4.js';
@@ -383,6 +383,7 @@ class GSplatComputeLocalRenderer extends GSplatRenderer {
383383
this._exposure = exposure ?? 1.0;
384384
this._fisheye = gsplat.fisheye;
385385
this._fogParams = fogParams ?? null;
386+
this._debugMode = gsplat.debug;
386387

387388
const formatHash = this.workBuffer.format.hash;
388389
if (formatHash !== this._formatHash) {
@@ -796,7 +797,8 @@ class GSplatComputeLocalRenderer extends GSplatRenderer {
796797

797798
const fogParams = this._fogParams;
798799
const fogType = (fogParams && fogParams.type !== FOG_NONE) ? fogParams.type : 'none';
799-
const rasterizeCompute = set.getRasterizeCompute(pickMode, useDepth, fogType);
800+
const heatmap = !pickMode && this._debugMode === GSPLAT_DEBUG_HEATMAP;
801+
const rasterizeCompute = set.getRasterizeCompute(pickMode, useDepth, fogType, heatmap);
800802

801803
rasterizeCompute.setParameter('screenWidth', width);
802804
rasterizeCompute.setParameter('screenHeight', height);

‎src/scene/gsplat-unified/gsplat-local-dispatch-set.js‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -281,15 +281,17 @@ class GSplatLocalDispatchSet {
281281
* @param {boolean} pickMode - Whether to use the pick variant.
282282
* @param {boolean} depthTest - Whether to enable depth testing against scene geometry.
283283
* @param {string} [fogType] - Fog type string: 'none', 'linear', 'exp', or 'exp2'.
284+
* @param {boolean} [heatmap] - Whether to enable heatmap debug visualization.
284285
* @returns {Compute} The cached Compute instance.
285286
*/
286-
getRasterizeCompute(pickMode, depthTest, fogType = 'none') {
287+
getRasterizeCompute(pickMode, depthTest, fogType = 'none', heatmap = false) {
287288
let key = pickMode ? 'pick' : 'color';
288289
if (depthTest) key += '-depth';
289290
if (fogType !== 'none') key += `-fog-${fogType}`;
291+
if (heatmap) key += '-heatmap';
290292
let variant = this._rasterizeVariants.get(key);
291293
if (!variant) {
292-
const { shader, bindGroupFormat } = this._createRasterizeShaderAndFormat(pickMode, depthTest, fogType);
294+
const { shader, bindGroupFormat } = this._createRasterizeShaderAndFormat(pickMode, depthTest, fogType, heatmap);
293295
const compute = new Compute(this.device, shader, `GSplatRasterize-${key}`);
294296
variant = { shader, bindGroupFormat, compute };
295297
this._rasterizeVariants.set(key, variant);
@@ -303,10 +305,11 @@ class GSplatLocalDispatchSet {
303305
* @param {boolean} pickMode - Whether to create the pick variant.
304306
* @param {boolean} depthTest - Whether to enable depth testing against scene geometry.
305307
* @param {string} [fogType] - Fog type string: 'none', 'linear', 'exp', or 'exp2'.
308+
* @param {boolean} [heatmap] - Whether to enable heatmap debug visualization.
306309
* @returns {{ shader: Shader, bindGroupFormat: BindGroupFormat }} The shader and format.
307310
* @private
308311
*/
309-
_createRasterizeShaderAndFormat(pickMode, depthTest = false, fogType = 'none') {
312+
_createRasterizeShaderAndFormat(pickMode, depthTest = false, fogType = 'none', heatmap = false) {
310313
const device = this.device;
311314
const hasFog = fogType !== 'none';
312315

@@ -354,6 +357,7 @@ class GSplatLocalDispatchSet {
354357
const cdefines = new Map();
355358
if (pickMode) cdefines.set('PICK_MODE', '');
356359
if (depthTest) cdefines.set('DEPTH_TEST', '');
360+
if (heatmap) cdefines.set('HEATMAP_MODE', '');
357361
cdefines.set('GAMMA', 'SRGB'); // assumes splat colors are in gamma space, will need to change when we get linear splats
358362
cdefines.set('FOG', hasFog ? fogType.toUpperCase() : 'NONE');
359363

‎src/scene/gsplat-unified/gsplat-params.js‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import {
99
GSPLATDATA_COMPACT,
1010
GSPLAT_RENDERER_AUTO, GSPLAT_RENDERER_RASTER_CPU_SORT,
1111
GSPLAT_RENDERER_RASTER_GPU_SORT, GSPLAT_RENDERER_COMPUTE,
12-
GSPLAT_DEBUG_NONE, GSPLAT_DEBUG_LOD, GSPLAT_DEBUG_SH_UPDATE
12+
GSPLAT_DEBUG_NONE, GSPLAT_DEBUG_LOD, GSPLAT_DEBUG_SH_UPDATE, GSPLAT_DEBUG_HEATMAP
1313
} from '../constants.js';
1414

1515
import glslCompactRead from '../shader-lib/glsl/chunks/gsplat/vert/formats/containerCompactRead.js';
@@ -219,6 +219,8 @@ class GSplatParams {
219219
* - {@link GSPLAT_DEBUG_LOD}: Colorize splats by their selected LOD level.
220220
* - {@link GSPLAT_DEBUG_SH_UPDATE}: Random color per SH update pass to visualize update
221221
* frequency.
222+
* - {@link GSPLAT_DEBUG_HEATMAP}: Heatmap visualization of average splats processed per
223+
* pixel in each tile. Only supported with {@link GSPLAT_RENDERER_COMPUTE}.
222224
*
223225
* Only one debug mode can be active at a time. Defaults to {@link GSPLAT_DEBUG_NONE}.
224226
*
@@ -229,7 +231,8 @@ class GSplatParams {
229231
const prev = this._debug;
230232
this._debug = value;
231233

232-
if (value === GSPLAT_DEBUG_LOD || prev === GSPLAT_DEBUG_LOD) {
234+
if (value === GSPLAT_DEBUG_LOD || prev === GSPLAT_DEBUG_LOD ||
235+
value === GSPLAT_DEBUG_HEATMAP || prev === GSPLAT_DEBUG_HEATMAP) {
233236
this.dirty = true;
234237
}
235238
}

‎src/scene/shader-lib/wgsl/chunks/gsplat/compute-gsplat-local-rasterize.js‎

Lines changed: 87 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,9 @@ struct Uniforms {
6060
6161
var<workgroup> sharedCenterScreen: array<vec2f, 64>;
6262
var<workgroup> sharedCoeffs: array<vec3f, 64>;
63+
#ifdef HEATMAP_MODE
64+
var<workgroup> sharedHeatCount: atomic<u32>;
65+
#endif
6366
6467
// Pick mode stores per-splat opacity, ID and depth; color mode stores packed RGBA.
6568
// Depth test mode also needs per-splat view depth for occlusion against scene geometry.
@@ -104,6 +107,20 @@ fn evalSplatPick(pixelCoord: vec2f, center: vec2f, coeffX: f32, coeffY: f32, coe
104107
}
105108
#endif
106109
110+
#ifdef HEATMAP_MODE
111+
fn heatmapColor(v: f32) -> vec3f {
112+
let t = saturate(v / 2000.0);
113+
if (t < 0.2) {
114+
return mix(vec3f(0.0, 0.0, 1.0), vec3f(0.0, 1.0, 1.0), t * 5.0);
115+
} else if (t < 0.4) {
116+
return mix(vec3f(0.0, 1.0, 1.0), vec3f(1.0, 1.0, 0.0), (t - 0.2) * 5.0);
117+
} else if (t < 0.6) {
118+
return mix(vec3f(1.0, 1.0, 0.0), vec3f(1.0, 0.0, 0.0), (t - 0.4) * 5.0);
119+
}
120+
return mix(vec3f(1.0, 0.0, 0.0), vec3f(0.15, 0.0, 0.0), (t - 0.6) * 2.5);
121+
}
122+
#endif
123+
107124
@compute @workgroup_size(8, 8)
108125
fn main(
109126
@builtin(local_invocation_id) lid: vec3u,
@@ -179,6 +196,12 @@ fn main(
179196
180197
let tileCount = tEnd - tStart;
181198
199+
#ifdef HEATMAP_MODE
200+
if (localIdx == 0u) { atomicStore(&sharedHeatCount, 0u); }
201+
workgroupBarrier();
202+
var processedCount: u32 = 0u;
203+
#endif
204+
182205
let numBatches = (tileCount + BATCH_SIZE - 1u) / BATCH_SIZE;
183206
var threadDone = false;
184207
@@ -288,6 +311,10 @@ fn main(
288311
c11 += splatColor.rgb * weight.w;
289312
T = select(T, newT, valid);
290313
314+
#ifdef HEATMAP_MODE
315+
processedCount += 1u;
316+
#endif
317+
291318
if (all(T < half4(ALPHA_THRESHOLD))) {
292319
threadDone = true;
293320
break;
@@ -299,44 +326,65 @@ fn main(
299326
workgroupBarrier();
300327
}
301328
302-
// Write results for the 2x2 pixel quad owned by this thread.
303-
// Pick mode: store the front-most pick ID and (accumulated depth, weight) per pixel.
304-
// Color mode: convert accumulated gamma-space color to linear via decodeGamma3 and store
305-
// to the rgba16float output texture; alpha holds total opacity (1 - transmittance).
306-
if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
307-
#ifdef PICK_MODE
308-
textureStore(pickIdTexture, basePixel, vec4u(pickId00, 0u, 0u, 0u));
309-
textureStore(pickDepthTexture, basePixel, vec4f(dAcc00, wAcc00, 0.0, 0.0));
310-
#else
311-
textureStore(outputTexture, basePixel, vec4f(decodeGamma3(vec3f(c00)), f32(half(1.0) - T.x)));
312-
#endif
313-
}
314-
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
315-
let px10 = vec2u(basePixel.x + 1u, basePixel.y);
316-
#ifdef PICK_MODE
317-
textureStore(pickIdTexture, px10, vec4u(pickId10, 0u, 0u, 0u));
318-
textureStore(pickDepthTexture, px10, vec4f(dAcc10, wAcc10, 0.0, 0.0));
319-
#else
320-
textureStore(outputTexture, px10, vec4f(decodeGamma3(vec3f(c10)), f32(half(1.0) - T.y)));
321-
#endif
322-
}
323-
if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
324-
let px01 = vec2u(basePixel.x, basePixel.y + 1u);
325-
#ifdef PICK_MODE
326-
textureStore(pickIdTexture, px01, vec4u(pickId01, 0u, 0u, 0u));
327-
textureStore(pickDepthTexture, px01, vec4f(dAcc01, wAcc01, 0.0, 0.0));
328-
#else
329-
textureStore(outputTexture, px01, vec4f(decodeGamma3(vec3f(c01)), f32(half(1.0) - T.z)));
330-
#endif
331-
}
332-
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
333-
let px11 = vec2u(basePixel.x + 1u, basePixel.y + 1u);
334-
#ifdef PICK_MODE
335-
textureStore(pickIdTexture, px11, vec4u(pickId11, 0u, 0u, 0u));
336-
textureStore(pickDepthTexture, px11, vec4f(dAcc11, wAcc11, 0.0, 0.0));
337-
#else
338-
textureStore(outputTexture, px11, vec4f(decodeGamma3(vec3f(c11)), f32(half(1.0) - T.w)));
339-
#endif
340-
}
329+
#ifdef HEATMAP_MODE
330+
atomicAdd(&sharedHeatCount, processedCount);
331+
workgroupBarrier();
332+
let avgCount = f32(atomicLoad(&sharedHeatCount)) / 64.0;
333+
let heatColor = vec4f(heatmapColor(avgCount), 1.0);
334+
if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
335+
textureStore(outputTexture, basePixel, heatColor);
336+
}
337+
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
338+
textureStore(outputTexture, vec2u(basePixel.x + 1u, basePixel.y), heatColor);
339+
}
340+
if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
341+
textureStore(outputTexture, vec2u(basePixel.x, basePixel.y + 1u), heatColor);
342+
}
343+
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
344+
textureStore(outputTexture, vec2u(basePixel.x + 1u, basePixel.y + 1u), heatColor);
345+
}
346+
#else
347+
348+
// Write results for the 2x2 pixel quad owned by this thread.
349+
// Pick mode: store the front-most pick ID and (accumulated depth, weight) per pixel.
350+
// Color mode: convert accumulated gamma-space color to linear via decodeGamma3 and store
351+
// to the rgba16float output texture; alpha holds total opacity (1 - transmittance).
352+
if (basePixel.x < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
353+
#ifdef PICK_MODE
354+
textureStore(pickIdTexture, basePixel, vec4u(pickId00, 0u, 0u, 0u));
355+
textureStore(pickDepthTexture, basePixel, vec4f(dAcc00, wAcc00, 0.0, 0.0));
356+
#else
357+
textureStore(outputTexture, basePixel, vec4f(decodeGamma3(vec3f(c00)), f32(half(1.0) - T.x)));
358+
#endif
359+
}
360+
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y < uniforms.screenHeight) {
361+
let px10 = vec2u(basePixel.x + 1u, basePixel.y);
362+
#ifdef PICK_MODE
363+
textureStore(pickIdTexture, px10, vec4u(pickId10, 0u, 0u, 0u));
364+
textureStore(pickDepthTexture, px10, vec4f(dAcc10, wAcc10, 0.0, 0.0));
365+
#else
366+
textureStore(outputTexture, px10, vec4f(decodeGamma3(vec3f(c10)), f32(half(1.0) - T.y)));
367+
#endif
368+
}
369+
if (basePixel.x < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
370+
let px01 = vec2u(basePixel.x, basePixel.y + 1u);
371+
#ifdef PICK_MODE
372+
textureStore(pickIdTexture, px01, vec4u(pickId01, 0u, 0u, 0u));
373+
textureStore(pickDepthTexture, px01, vec4f(dAcc01, wAcc01, 0.0, 0.0));
374+
#else
375+
textureStore(outputTexture, px01, vec4f(decodeGamma3(vec3f(c01)), f32(half(1.0) - T.z)));
376+
#endif
377+
}
378+
if (basePixel.x + 1u < uniforms.screenWidth && basePixel.y + 1u < uniforms.screenHeight) {
379+
let px11 = vec2u(basePixel.x + 1u, basePixel.y + 1u);
380+
#ifdef PICK_MODE
381+
textureStore(pickIdTexture, px11, vec4u(pickId11, 0u, 0u, 0u));
382+
textureStore(pickDepthTexture, px11, vec4f(dAcc11, wAcc11, 0.0, 0.0));
383+
#else
384+
textureStore(outputTexture, px11, vec4f(decodeGamma3(vec3f(c11)), f32(half(1.0) - T.w)));
385+
#endif
386+
}
387+
388+
#endif
341389
}
342390
`;

0 commit comments

Comments
 (0)