GSplat reorder and shader speedup (#6715) (#6716)

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(

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 };

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
  */
 
 /**

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;