Gsplat sorting update

src/scene/gsplat/gsplat-compressed-data.js

@@ -221,6 +221,29 @@ class GSplatCompressedData {
         }
     }
 
+    getChunks(result) {
+        const { chunkData, numChunks, chunkSize } = this;
+
+        let mx, my, mz, Mx, My, Mz;
+
+        for (let c = 0; c < numChunks; ++c) {
+            const off = c * chunkSize;
+            mx = chunkData[off + 0];
+            my = chunkData[off + 1];
+            mz = chunkData[off + 2];
+            Mx = chunkData[off + 3];
+            My = chunkData[off + 4];
+            Mz = chunkData[off + 5];
+
+            result[c * 6 + 0] = mx;
+            result[c * 6 + 1] = my;
+            result[c * 6 + 2] = mz;
+            result[c * 6 + 3] = Mx;
+            result[c * 6 + 4] = My;
+            result[c * 6 + 5] = Mz;
+        }
+    }
+
     /**
      * @param {Vec3} result - The result.
      */

src/scene/gsplat/gsplat-compressed.js

@@ -69,6 +69,9 @@ class GSplatCompressed {
         this.centers = new Float32Array(numSplats * 3);
         gsplatData.getCenters(this.centers);
 
+        this.chunks = new Float32Array(numChunks * 6);
+        gsplatData.getChunks(this.chunks);
+
         // initialize packed data
         this.packedTexture = this.createTexture('packedData', PIXELFORMAT_RGBA32U, this.evalTextureSize(numSplats), vertexData);
 

src/scene/gsplat/gsplat-instance.js

@@ -132,12 +132,13 @@ class GSplatInstance {
         this.meshInstance.instancingCount = 0;
 
         // clone centers to allow multiple instances of sorter
-        this.centers = new Float32Array(splat.centers);
+        const centers = splat.centers.slice();
+        const chunks = splat.chunks?.slice();
 
         // create sorter
         if (!options.dither || options.dither === DITHER_NONE) {
             this.sorter = new GSplatSorter();
-            this.sorter.init(this.orderTexture, this.centers);
+            this.sorter.init(this.orderTexture, centers, chunks);
             this.sorter.on('updated', (count) => {
                 // limit splat render count to exclude those behind the camera
                 this.meshInstance.instancingCount = Math.ceil(count / splatInstanceSize);

src/scene/gsplat/gsplat-sorter.js

@@ -5,6 +5,7 @@ import { TEXTURELOCK_READ } from '../../platform/graphics/constants.js';
 function SortWorker() {
     let order;
     let centers;
+    let chunks;
     let mapping;
     let cameraPosition;
     let cameraDirection;
@@ -20,6 +21,12 @@ function SortWorker() {
     let distances;
     let countBuffer;
 
+    // could be increased, but this seems a good compromise between stability and performance
+    const numBins = 32;
+    const binCount = new Array(numBins).fill(0);
+    const binBase = new Array(numBins).fill(0);
+    const binDivider = new Array(numBins).fill(0);
+
     const binarySearch = (m, n, compare_fn) => {
         while (m <= n) {
             const k = (n + m) >> 1;
@@ -99,21 +106,83 @@ function SortWorker() {
             countBuffer.fill(0);
         }
 
-        // generate per vertex distance to camera
         const range = maxDist - minDist;
-        const divider = (range < 1e-6) ? 0 : 1 / range * (2 ** compareBits);
-        for (let i = 0; i < numVertices; ++i) {
-            const istride = i * 3;
-            const x = centers[istride + 0];
-            const y = centers[istride + 1];
-            const z = centers[istride + 2];
-            const d = x * dx + y * dy + z * dz;
-            const sortKey = Math.floor((d - minDist) * divider);
 
-            distances[i] = sortKey;
+        if (range < 1e-6) {
+            // all points are at the same distance
+            for (let i = 0; i < numVertices; ++i) {
+                distances[i] = 0;
+                countBuffer[0]++;
+            }
+        } else if (chunks) {
+            // handle sort with compressed chunks
+            const numChunks = chunks.length / 6;
+
+            // calculate a histogram of chunk distances to camera
+            binCount.fill(0);
+            for (let i = 0; i < numChunks; ++i) {
+                const x = chunks[i * 6 + 0];
+                const y = chunks[i * 6 + 1];
+                const z = chunks[i * 6 + 2];
+                const r = chunks[i * 6 + 3];
+                const d = x * dx + y * dy + z * dz - minDist;
+
+                const binMin = Math.max(0, Math.floor((d - r) * numBins / range));
+                const binMax = Math.min(numBins, Math.ceil((d + r) * numBins / range));
+
+                for (let j = binMin; j < binMax; ++j) {
+                    binCount[j]++;
+                }
+            }
+
+            // count total number of histogram bin entries
+            const binTotal = binCount.reduce((a, b) => a + b, 0);
 
-            // count occurrences of each distance
-            countBuffer[sortKey]++;
+            // calculate per-bin base and divider
+            for (let i = 0; i < numBins; ++i) {
+                binDivider[i] = Math.ceil(binCount[i] / binTotal * bucketCount);
+            }
+            for (let i = 0; i < numBins; ++i) {
+                binBase[i] = i === 0 ? 0 : binBase[i - 1] + binDivider[i - 1];
+            }
+
+            // generate per vertex distance key using histogram to distribute bits
+            const binRange = range / numBins;
+            let ii = 0;
+            for (let i = 0; i < numVertices; ++i) {
+                const x = centers[ii++];
+                const y = centers[ii++];
+                const z = centers[ii++];
+                const d = (x * dx + y * dy + z * dz - minDist) / binRange;
+                const bin = d >>> 0;
+                const sortKey = (binBase[bin] + binDivider[bin] * (d - bin)) >>> 0;
+
+                if (sortKey < 0 || sortKey >= bucketCount) {
+                    console.log(`i=${i} d=${d} bin=${bin} sortKey=${sortKey} bucketCount=${bucketCount}`);
+                }
+
+                distances[i] = sortKey;
+
+                // count occurrences of each distance
+                countBuffer[sortKey]++;
+            }
+        } else {
+            // generate per vertex distance to camera for uncompressed data
+            const divider = (2 ** compareBits) / range;
+            let ii = 0;
+            for (let i = 0; i < numVertices; ++i) {
+                const x = centers[ii++];
+                const y = centers[ii++];
+                const z = centers[ii++];
+
+                const d = (x * dx + y * dy + z * dz - minDist) * divider;
+                const sortKey = d >>> 0;
+
+                distances[i] = sortKey;
+
+                // count occurrences of each distance
+                countBuffer[sortKey]++;
+            }
         }
 
         // Change countBuffer[i] so that it contains actual position of this digit in outputArray
@@ -128,13 +197,17 @@ function SortWorker() {
             order[destIndex] = i;
         }
 
-        // find splat with distance 0 to limit rendering behind the camera
-        const tmp = -px * dx - py * dy - pz * dz;
-        const dist = i => distances[order[i]] / divider + minDist + tmp;
+        // Find splat with distance 0 to limit rendering behind the camera
+        const cameraDist = px * dx + py * dy + pz * dz;
+        const dist = (i) => {
+            let o = order[i] * 3;
+            return centers[o++] * dx + centers[o++] * dy + centers[o] * dz - cameraDist;
+        };
         const findZero = () => {
             const result = binarySearch(0, numVertices - 1, i => -dist(i));
             return Math.min(numVertices, Math.abs(result));
         };
+
         const count = dist(numVertices - 1) >= 0 ? findZero() : numVertices;
 
         // apply mapping
@@ -159,6 +232,7 @@ function SortWorker() {
         }
         if (message.data.centers) {
             centers = new Float32Array(message.data.centers);
+            forceUpdate = true;
 
             // calculate bounds
             let initialized = false;
@@ -196,8 +270,25 @@ function SortWorker() {
             if (!initialized) {
                 boundMin.x = boundMax.x = boundMin.y = boundMax.y = boundMin.z = boundMax.z = 0;
             }
-
+        }
+        if (message.data.chunks) {
+            chunks = new Float32Array(message.data.chunks);
             forceUpdate = true;
+
+            // convert chunk min/max to center/radius
+            for (let i = 0; i < chunks.length / 6; ++i) {
+                const mx = chunks[i * 6 + 0];
+                const my = chunks[i * 6 + 1];
+                const mz = chunks[i * 6 + 2];
+                const Mx = chunks[i * 6 + 3];
+                const My = chunks[i * 6 + 4];
+                const Mz = chunks[i * 6 + 5];
+
+                chunks[i * 6 + 0] = (mx + Mx) * 0.5;
+                chunks[i * 6 + 1] = (my + My) * 0.5;
+                chunks[i * 6 + 2] = (mz + Mz) * 0.5;
+                chunks[i * 6 + 3] = Math.sqrt((Mx - mx) ** 2 + (My - my) ** 2 + (Mz - mz) ** 2) * 0.5;
+            }
         }
         if (message.data.hasOwnProperty('mapping')) {
             mapping = message.data.mapping ? new Uint32Array(message.data.mapping) : null;
@@ -247,27 +338,31 @@ class GSplatSorter extends EventHandler {
         this.worker = null;
     }
 
-    init(orderTexture, centers) {
+    init(orderTexture, centers, chunks) {
         this.orderTexture = orderTexture;
         this.centers = centers.slice();
 
         // get the texture's storage buffer and make a copy
         const orderBuffer = this.orderTexture.lock({
             mode: TEXTURELOCK_READ
-        }).buffer.slice();
+        }).slice();
+        this.orderTexture.unlock();
 
         // initialize order data
         for (let i = 0; i < orderBuffer.length; ++i) {
             orderBuffer[i] = i;
         }
 
-        this.orderTexture.unlock();
+        const obj = {
+            order: orderBuffer.buffer,
+            centers: centers.buffer,
+            chunks: chunks?.buffer
+        };
+
+        const transfer = [orderBuffer.buffer, centers.buffer].concat(chunks ? [chunks.buffer] : []);
 
         // send the initial buffer to worker
-        this.worker.postMessage({
-            order: orderBuffer,
-            centers: centers.buffer
-        }, [orderBuffer, centers.buffer]);
+        this.worker.postMessage(obj, transfer);
     }
 
     setMapping(mapping) {