Update GS rendering v2.0 (playcanvas#6357)

src/framework/parsers/ply.js

@@ -294,7 +294,11 @@ class PlyParser {
         } else {
             readPly(response.body.getReader(), asset.data.elementFilter ?? defaultElementFilter)
                 .then((response) => {
-                    callback(null, new GSplatResource(this.device, new GSplatData(response)));
+                    const gsplatData = new GSplatData(response, {
+                        performZScale: asset.data.performZScale,
+                        reorder: asset.data.reorder
+                    });
+                    callback(null, new GSplatResource(this.device, gsplatData));
                 })
                 .catch((err) => {
                     callback(err, null);

src/platform/graphics/shader-utils.js

@@ -209,7 +209,7 @@ class ShaderUtils {
 
         if (!device.isWebGPU) {
 
-            code = `precision ${precision} float;\n`;
+            code = `precision ${precision} float;\nprecision ${precision} int;`;
 
             if (device.isWebGL2) {
                 code += `precision ${precision} sampler2DShadow;\n`;

src/platform/graphics/texture.js

@@ -781,7 +781,7 @@ class Texture {
      * - {@link TEXTURELOCK_READ}
      * - {@link TEXTURELOCK_WRITE}
      * Defaults to {@link TEXTURELOCK_WRITE}.
-     * @returns {Uint8Array|Uint16Array|Float32Array} A typed array containing the pixel data of
+     * @returns {Uint8Array|Uint16Array|Uint32Array|Float32Array} A typed array containing the pixel data of
      * the locked mip level.
      */
     lock(options = {}) {

src/scene/gsplat/gsplat-data.js

@@ -80,14 +80,25 @@ 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, performZScale = true) {
+    constructor(elements, options = {}) {
         this.elements = elements;
         this.vertexElement = elements.find(element => element.name === 'vertex');
 
-        if (!this.isCompressed && performZScale) {
-            mat4.setScale(-1, -1, 1);
-            this.transform(mat4);
+        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();
+            }
         }
     }
 
@@ -406,7 +417,87 @@ class GSplatData {
                     storage: data[name]
                 };
             })
-        }], false);
+        }], {
+            performZScale: false,
+            reorder: false
+        });
+    }
+
+    calcMortonOrder() {
+        const calcMinMax = (arr) => {
+            let min = arr[0];
+            let max = arr[0];
+            for (let i = 1; i < arr.length; i++) {
+                if (arr[i] < min) min = arr[i];
+                if (arr[i] > max) max = arr[i];
+            }
+            return { min, max };
+        };
+
+        // https://fgiesen.wordpress.com/2009/12/13/decoding-morton-codes/
+        const encodeMorton3 = (x, y, z) => {
+            const Part1By2 = (x) => {
+                x &= 0x000003ff;
+                x = (x ^ (x << 16)) & 0xff0000ff;
+                x = (x ^ (x <<  8)) & 0x0300f00f;
+                x = (x ^ (x <<  4)) & 0x030c30c3;
+                x = (x ^ (x <<  2)) & 0x09249249;
+                return x;
+            };
+
+            return (Part1By2(z) << 2) + (Part1By2(y) << 1) + Part1By2(x);
+        };
+
+        const x = this.getProp('x');
+        const y = this.getProp('y');
+        const z = this.getProp('z');
+
+        const { min: minX, max: maxX } = calcMinMax(x);
+        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 morton = new Uint32Array(this.numSplats);
+        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);
+        }
+
+        // generate indices
+        const indices = new Uint32Array(this.numSplats);
+        for (let i = 0; i < this.numSplats; i++) {
+            indices[i] = i;
+        }
+        // order splats by morton code
+        indices.sort((a, b) => morton[a] - morton[b]);
+
+        return indices;
+    }
+
+    reorderData() {
+        // calculate splat morton order
+        const order = this.calcMortonOrder();
+
+        const reorder = (data) => {
+            const result = new data.constructor(data.length);
+
+            for (let i = 0; i < order.length; i++) {
+                result[i] = data[order[i]];
+            }
+
+            return result;
+        };
+
+        this.elements.forEach((element) => {
+            element.properties.forEach((property) => {
+                property.storage = reorder(property.storage);
+            });
+        });
     }
 }
 

src/scene/gsplat/gsplat-instance.js

@@ -1,11 +1,13 @@
 import { Mat4 } from '../../core/math/mat4.js';
 import { Vec3 } from '../../core/math/vec3.js';
-import { SEMANTIC_POSITION, TYPE_UINT32 } from '../../platform/graphics/constants.js';
+import { BUFFER_STATIC, PIXELFORMAT_R32U, SEMANTIC_ATTR13, TYPE_UINT32 } from '../../platform/graphics/constants.js';
 import { DITHER_NONE } from '../constants.js';
 import { MeshInstance } from '../mesh-instance.js';
 import { Mesh } from '../mesh.js';
 import { createGSplatMaterial } from './gsplat-material.js';
 import { GSplatSorter } from './gsplat-sorter.js';
+import { VertexFormat } from '../../platform/graphics/vertex-format.js';
+import { VertexBuffer } from '../../platform/graphics/vertex-buffer.js';
 
 const mat = new Mat4();
 const cameraPosition = new Vec3();
@@ -26,8 +28,8 @@ class GSplatInstance {
     /** @type {import('../materials/material.js').Material} */
     material;
 
-    /** @type {import('../../platform/graphics/vertex-buffer.js').VertexBuffer} */
-    vb;
+    /** @type {import('../../platform/graphics/texture.js').Texture} */
+    orderTexture;
 
     options = {};
 
@@ -58,53 +60,81 @@ class GSplatInstance {
 
         // not supported on WebGL1
         const device = splat.device;
-        if (device.isWebGL1)
-            return;
+
+        // create the order texture
+        this.orderTexture = this.splat.createTexture(
+            'splatOrder',
+            PIXELFORMAT_R32U,
+            this.splat.evalTextureSize(this.splat.numSplats)
+        );
 
         // material
         this.createMaterial(options);
 
-        const numSplats = splat.numSplats;
-        const indices = new Uint32Array(numSplats * 6);
-        const ids = new Uint32Array(numSplats * 4);
-
-        for (let i = 0; i < numSplats; ++i) {
-            const base = i * 4;
-
-            // 4 vertices
-            ids[base + 0] = i;
-            ids[base + 1] = i;
-            ids[base + 2] = i;
-            ids[base + 3] = i;
-
-            // 2 triangles
-            const triBase = i * 6;
-            indices[triBase + 0] = base;
-            indices[triBase + 1] = base + 1;
-            indices[triBase + 2] = base + 2;
-            indices[triBase + 3] = base;
-            indices[triBase + 4] = base + 2;
-            indices[triBase + 5] = base + 3;
+        // number of quads to combine into a single instance. this is to increase occupancy
+        // in the vertex shader.
+        const splatInstanceSize = 128;
+        const numSplats = Math.ceil(splat.numSplats / splatInstanceSize) * splatInstanceSize;
+        const numSplatInstances = numSplats / splatInstanceSize;
+
+        // specify the base splat index per instance
+        const indexData = new Uint32Array(numSplatInstances);
+        for (let i = 0; i < numSplatInstances; ++i) {
+            indexData[i] = i * splatInstanceSize;
+        }
+
+        const vertexFormat = new VertexFormat(device, [
+            { semantic: SEMANTIC_ATTR13, components: 1, type: TYPE_UINT32, asInt: true }
+        ]);
+
+        const indicesVB = new VertexBuffer(device, vertexFormat, numSplatInstances, {
+            usage: BUFFER_STATIC,
+            data: indexData.buffer
+        });
+
+        // build the instance mesh
+        const meshPositions = new Float32Array(12 * splatInstanceSize);
+        const meshIndices = new Uint32Array(6 * splatInstanceSize);
+        for (let i = 0; i < splatInstanceSize; ++i) {
+            meshPositions.set([
+                -2, -2, i,
+                2, -2, i,
+                2, 2, i,
+                -2, 2, i
+            ], i * 12);
+
+            const b = i * 4;
+            meshIndices.set([
+                0 + b, 1 + b, 2 + b, 0 + b, 2 + b, 3 + b
+            ], i * 6);
         }
 
-        // mesh
         const mesh = new Mesh(device);
-        mesh.setVertexStream(SEMANTIC_POSITION, ids, 1, numSplats * 4, TYPE_UINT32, false, !device.isWebGL1);
-        mesh.setIndices(indices);
+        mesh.setPositions(meshPositions, 3);
+        mesh.setIndices(meshIndices);
         mesh.update();
+
         this.mesh = mesh;
         this.mesh.aabb.copy(splat.aabb);
 
         this.meshInstance = new MeshInstance(this.mesh, this.material);
+        this.meshInstance.setInstancing(indicesVB, true);
         this.meshInstance.gsplatInstance = this;
+        this.meshInstance.instancingCount = numSplatInstances;
 
         // clone centers to allow multiple instances of sorter
         this.centers = new Float32Array(splat.centers);
 
         // create sorter
         if (!options.dither || options.dither === DITHER_NONE) {
             this.sorter = new GSplatSorter();
-            this.sorter.init(mesh.vertexBuffer, this.centers, !this.splat.device.isWebGL1);
+            this.sorter.init(this.orderTexture, this.centers);
+            this.sorter.on('updated', (count) => {
+                // limit splat render count to exclude those behind the camera.
+                // NOTE: the last instance rendered may include non-existant splat
+                // data. this should be ok though as the data is filled with 0's.
+                this.meshInstance.instancingCount = Math.ceil(count / splatInstanceSize);
+            });
         }
     }
 
@@ -120,6 +150,7 @@ class GSplatInstance {
 
     createMaterial(options) {
         this.material = createGSplatMaterial(options);
+        this.material.setParameter('splatOrder', this.orderTexture);
         this.splat.setupMaterial(this.material);
         if (this.meshInstance) {
             this.meshInstance.material = this.material;