From 523c320741c3412e97602f3692b9bfc05b419bd6 Mon Sep 17 00:00:00 2001 From: Donovan Hutchence Date: Tue, 18 Jun 2024 16:26:36 +0100 Subject: [PATCH] GSplat reorder and shader speedup (#6715) (#6716) --- src/framework/parsers/ply.js | 22 +++++- src/scene/gsplat/gsplat-data.js | 82 ++++++++++++--------- src/scene/gsplat/gsplat-material.js | 2 + src/scene/gsplat/shader-generator-gsplat.js | 17 ++++- 4 files changed, 79 insertions(+), 44 deletions(-) diff --git a/src/framework/parsers/ply.js b/src/framework/parsers/ply.js index 18822f21293..796b81efdf9 100644 --- a/src/framework/parsers/ply.js +++ b/src/framework/parsers/ply.js @@ -1,5 +1,6 @@ import { GSplatData } from '../../scene/gsplat/gsplat-data.js'; import { GSplatResource } from './gsplat-resource.js'; +import { Mat4 } from '../../core/math/mat4.js'; const magicBytes = new Uint8Array([112, 108, 121, 10]); // ply\n const endHeaderBytes = new Uint8Array([10, 101, 110, 100, 95, 104, 101, 97, 100, 101, 114, 10]); // \nend_header\n @@ -258,6 +259,8 @@ const defaultElements = [ const defaultElementsSet = new Set(defaultElements); const defaultElementFilter = val => defaultElementsSet.has(val); +const mat4 = new Mat4(); + class PlyParser { /** @type {import('../../platform/graphics/graphics-device.js').GraphicsDevice} */ device; @@ -295,10 +298,21 @@ class PlyParser { readPly(response.body.getReader(), asset.data.elementFilter ?? defaultElementFilter) .then((response) => { // construct the GSplatData object - const gsplatData = new GSplatData(response, { - performZScale: asset.data.performZScale, - reorder: asset.data.reorder - }); + const gsplatData = new GSplatData(response); + + if (!gsplatData.isCompressed) { + + // perform Z scale + if (asset.data.performZScale ?? true) { + mat4.setScale(-1, -1, 1); + gsplatData.transform(mat4); + } + + // reorder data + if (asset.data.reorder ?? true) { + gsplatData.reorderData(); + } + } // construct the resource const resource = new GSplatResource( diff --git a/src/scene/gsplat/gsplat-data.js b/src/scene/gsplat/gsplat-data.js index 061a11f5097..1af09c988b6 100644 --- a/src/scene/gsplat/gsplat-data.js +++ b/src/scene/gsplat/gsplat-data.js @@ -182,28 +182,10 @@ class GSplatData { // /** // * @param {import('./ply-reader').PlyElement[]} elements - The elements. - // * @param {boolean} [performZScale] - Whether to perform z scaling. - // * @param {object} [options] - The options. - // * @param {boolean} [options.performZScale] - Whether to perform z scaling. - // * @param {boolean} [options.reorder] - Whether to reorder the data. // */ - constructor(elements, options = {}) { + constructor(elements) { this.elements = elements; - this.numSplats = this.getElement('vertex').count; - - if (!this.isCompressed) { - if (options.performZScale ?? true) { - mat4.setScale(-1, -1, 1); - this.transform(mat4); - } - - // reorder uncompressed splats in morton order for better memory access - // efficiency during rendering - if (options.reorder ?? true) { - this.reorderData(); - } - } } /** @@ -585,10 +567,7 @@ class GSplatData { storage: data[name] }; }) - }], { - performZScale: false, - reorder: false - }); + }]); } calcMortonOrder() { @@ -624,40 +603,66 @@ class GSplatData { const { min: minY, max: maxY } = calcMinMax(y); const { min: minZ, max: maxZ } = calcMinMax(z); - const sizeX = 1024 / (maxX - minX); - const sizeY = 1024 / (maxY - minY); - const sizeZ = 1024 / (maxZ - minZ); + const sizeX = minX === maxX ? 0 : 1024 / (maxX - minX); + const sizeY = minY === maxY ? 0 : 1024 / (maxY - minY); + const sizeZ = minZ === maxZ ? 0 : 1024 / (maxZ - minZ); - const morton = new Uint32Array(this.numSplats); + const codes = new Map(); for (let i = 0; i < this.numSplats; i++) { const ix = Math.floor((x[i] - minX) * sizeX); const iy = Math.floor((y[i] - minY) * sizeY); const iz = Math.floor((z[i] - minZ) * sizeZ); - morton[i] = encodeMorton3(ix, iy, iz); + const code = encodeMorton3(ix, iy, iz); + + const val = codes.get(code); + if (val) { + val.push(i); + } else { + codes.set(code, [i]); + } } - // generate indices + const keys = Array.from(codes.keys()).sort((a, b) => a - b); const indices = new Uint32Array(this.numSplats); - for (let i = 0; i < this.numSplats; i++) { - indices[i] = i; + let idx = 0; + + for (let i = 0; i < keys.length; ++i) { + const val = codes.get(keys[i]); + for (let j = 0; j < val.length; ++j) { + indices[idx++] = val[j]; + } } - // order splats by morton code - indices.sort((a, b) => morton[a] - morton[b]); return indices; } // reorder the splat data to aid in better gpu memory access at render time - reorderData() { - const order = this.calcMortonOrder(); + reorder(order) { + const cache = new Map(); + + const getStorage = (size) => { + if (cache.has(size)) { + const buffer = cache.get(size); + cache.delete(size); + return buffer; + } + + return new ArrayBuffer(size); + }; + + const returnStorage = (buffer) => { + cache.set(buffer.byteLength, buffer); + }; const reorder = (data) => { - const result = new data.constructor(data.length); + const result = new data.constructor(getStorage(data.byteLength)); for (let i = 0; i < order.length; i++) { result[i] = data[order[i]]; } + returnStorage(data.buffer); + return result; }; @@ -669,6 +674,11 @@ class GSplatData { }); }); } + + // reorder the splat data to aid in better gpu memory access at render time + reorderData() { + this.reorder(this.calcMortonOrder()); + } } export { GSplatData }; diff --git a/src/scene/gsplat/gsplat-material.js b/src/scene/gsplat/gsplat-material.js index 9a12632cb6a..adcce779f4f 100644 --- a/src/scene/gsplat/gsplat-material.js +++ b/src/scene/gsplat/gsplat-material.js @@ -27,6 +27,8 @@ const splatMainFS = ` * @property {string} [vertex] - Custom vertex shader, see SPLAT MANY example. * @property {string} [fragment] - Custom fragment shader, see SPLAT MANY example. * @property {string} [dither] - Opacity dithering enum. + * + * @ignore */ /** diff --git a/src/scene/gsplat/shader-generator-gsplat.js b/src/scene/gsplat/shader-generator-gsplat.js index 206a0193eef..4859c05e67a 100644 --- a/src/scene/gsplat/shader-generator-gsplat.js +++ b/src/scene/gsplat/shader-generator-gsplat.js @@ -16,7 +16,10 @@ const splatCoreVS = /* glsl */ ` varying vec2 texCoord; varying vec4 color; - varying float id; + + #ifndef DITHER_NONE + varying float id; + #endif // width, numSplats uniform vec2 tex_params; @@ -88,7 +91,10 @@ const splatCoreVS = /* glsl */ ` return vec4(0.0, 0.0, 2.0, 1.0); } - id = float(splatId); + #ifndef DITHER_NONE + id = float(splatId); + #endif + color = getColor(); mat3 Vrk = mat3( @@ -138,10 +144,13 @@ const splatCoreVS = /* glsl */ ` } `; -const splatCoreFS = /* glsl_ */ ` +const splatCoreFS = /* glsl */ ` varying vec2 texCoord; varying vec4 color; - varying float id; + + #ifndef DITHER_NONE + varying float id; + #endif #ifdef PICK_PASS uniform vec4 uColor;