Converted remaining Gsplat shader chunks to WGSL (#7650)

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>

Author: mvaligursky

examples/assets/splats/skull.ply

@@ -0,0 +1,141 @@
+import { deviceType, rootPath } from 'examples/utils';
+import * as pc from 'playcanvas';
+
+
+const canvas = /** @type {HTMLCanvasElement} */ (document.getElementById('application-canvas'));
+window.focus();
+
+const gfxOptions = {
+    deviceTypes: [deviceType],
+    glslangUrl: `${rootPath}/static/lib/glslang/glslang.js`,
+    twgslUrl: `${rootPath}/static/lib/twgsl/twgsl.js`,
+
+    // disable antialiasing as gaussian splats do not benefit from it and it's expensive
+    antialias: false
+};
+
+const device = await pc.createGraphicsDevice(canvas, gfxOptions);
+device.maxPixelRatio = Math.min(window.devicePixelRatio, 2);
+
+const createOptions = new pc.AppOptions();
+createOptions.graphicsDevice = device;
+createOptions.mouse = new pc.Mouse(document.body);
+createOptions.touch = new pc.TouchDevice(document.body);
+
+createOptions.componentSystems = [
+    pc.RenderComponentSystem,
+    pc.CameraComponentSystem,
+    pc.LightComponentSystem,
+    pc.ScriptComponentSystem,
+    pc.GSplatComponentSystem
+];
+createOptions.resourceHandlers = [pc.TextureHandler, pc.ContainerHandler, pc.ScriptHandler, pc.GSplatHandler];
+
+const app = new pc.AppBase(canvas);
+app.init(createOptions);
+
+// Set the canvas to fill the window and automatically change resolution to be the same as the canvas size
+app.setCanvasFillMode(pc.FILLMODE_FILL_WINDOW);
+app.setCanvasResolution(pc.RESOLUTION_AUTO);
+
+// Ensure canvas is resized when window changes size
+const resize = () => app.resizeCanvas();
+window.addEventListener('resize', resize);
+app.on('destroy', () => {
+    window.removeEventListener('resize', resize);
+});
+
+const assets = {
+    skull: new pc.Asset('gsplat', 'gsplat', { url: `${rootPath}/static/assets/splats/skull.ply` }),
+    orbit: new pc.Asset('script', 'script', { url: `${rootPath}/static/scripts/camera/orbit-camera.js` }),
+    hdri_street: new pc.Asset(
+        'hdri',
+        'texture',
+        { url: `${rootPath}/static/assets/hdri/st-peters-square.hdr` },
+        { mipmaps: false }
+    )
+};
+
+const assetListLoader = new pc.AssetListLoader(Object.values(assets), app.assets);
+assetListLoader.load(() => {
+    app.start();
+
+    // create a splat entity and place it in the world
+    const skull = new pc.Entity();
+    skull.addComponent('gsplat', {
+        asset: assets.skull,
+        castShadows: true
+    });
+    skull.setLocalPosition(-1.5, 0.05, 0);
+    skull.setLocalEulerAngles(180, 90, 0);
+    skull.setLocalScale(0.7, 0.7, 0.7);
+    app.root.addChild(skull);
+
+    // set alpha clip value, used by shadows and picking
+    skull.gsplat.material.setParameter('alphaClip', 0.4);
+    skull.gsplat.material.setParameter('alphaClip', 0.1);
+
+    // Create an Entity with a camera component
+    const camera = new pc.Entity();
+    camera.addComponent('camera', {
+        clearColor: new pc.Color(0.2, 0.2, 0.2),
+        toneMapping: pc.TONEMAP_ACES
+    });
+    camera.setLocalPosition(-2, 1.5, 2);
+
+    // add orbit camera script with a mouse and a touch support
+    camera.addComponent('script');
+    camera.script.create('orbitCamera', {
+        attributes: {
+            inertiaFactor: 0.2,
+            focusEntity: skull,
+            distanceMax: 60,
+            frameOnStart: false
+        }
+    });
+    camera.script.create('orbitCameraInputMouse');
+    camera.script.create('orbitCameraInputTouch');
+    app.root.addChild(camera);
+
+    // create ground to receive shadows
+    const material = new pc.StandardMaterial();
+    material.diffuse = new pc.Color(0.5, 0.5, 0.4);
+    material.gloss = 0.2;
+    material.metalness = 0.5;
+    material.useMetalness = true;
+    material.update();
+
+    const ground = new pc.Entity();
+    ground.addComponent('render', {
+        type: 'box',
+        material: material,
+        castShadows: false
+    });
+    ground.setLocalScale(10, 1, 10);
+    ground.setLocalPosition(0, -0.45, 0);
+    app.root.addChild(ground);
+
+    // shadow casting directional light
+    // Note: it does not affect gsplat, as lighting is not supported there currently
+    const directionalLight = new pc.Entity();
+    directionalLight.addComponent('light', {
+        type: 'directional',
+        color: pc.Color.WHITE,
+        castShadows: true,
+        intensity: 1,
+        shadowBias: 0.2,
+        normalOffsetBias: 0.05,
+        shadowDistance: 10,
+        shadowIntensity: 0.5,
+        shadowResolution: 2048,
+        shadowType: pc.SHADOW_PCSS_32F,
+        penumbraSize: 10,
+        penumbraFalloff: 4,
+        shadowSamples: 16,
+        shadowBlockerSamples: 16
+    });
+    directionalLight.setEulerAngles(55, 90, 20);
+    app.root.addChild(directionalLight);
+});
+
+export { app };

examples/src/examples/loaders/gsplat-sh.example.mjs

@@ -1,6 +1,7 @@
 import { CULLFACE_NONE, SEMANTIC_ATTR13, SEMANTIC_POSITION } from '../../platform/graphics/constants.js';
 import { BLEND_NONE, BLEND_PREMULTIPLIED, DITHER_NONE } from '../constants.js';
 import { ShaderMaterial } from '../materials/shader-material.js';
+import { shaderChunksWGSL } from '../shader-lib/chunks-wgsl/chunks-wgsl.js';
 import { shaderChunks } from '../shader-lib/chunks/chunks.js';
 
 /**
@@ -27,6 +28,8 @@ const createGSplatMaterial = (options = {}) => {
         uniqueName: 'SplatMaterial',
         vertexGLSL: options.vertex ?? shaderChunks.gsplatVS,
         fragmentGLSL: options.fragment ?? shaderChunks.gsplatPS,
+        vertexWGSL: shaderChunksWGSL.gsplatVS,
+        fragmentWGSL: shaderChunksWGSL.gsplatPS,
         attributes: {
             vertex_position: SEMANTIC_POSITION,
             vertex_id_attrib: SEMANTIC_ATTR13

examples/thumbnails/loaders_gsplat-sh_large.webp

@@ -44,17 +44,17 @@ import fresnelSchlickPS from './lit/frag/fresnelSchlick.js';
 import fullscreenQuadVS from './common/vert/fullscreenQuad.js';
 import gammaPS from './common/frag/gamma.js';
 import glossPS from './standard/frag/gloss.js';
-// import gsplatCenterVS from './gsplat/vert/gsplatCenter.js';
-// import gsplatColorVS from './gsplat/vert/gsplatColor.js';
+import gsplatCenterVS from './gsplat/vert/gsplatCenter.js';
+import gsplatColorVS from './gsplat/vert/gsplatColor.js';
 import gsplatCommonVS from './gsplat/vert/gsplatCommon.js';
-// import gsplatCompressedDataVS from './gsplat/vert/gsplatCompressedData.js';
-// import gsplatCompressedSHVS from './gsplat/vert/gsplatCompressedSH.js';
-// import gsplatCornerVS from './gsplat/vert/gsplatCorner.js';
-// import gsplatDataVS from './gsplat/vert/gsplatData.js';
-// import gsplatOutputVS from './gsplat/vert/gsplatOutput.js';
+import gsplatCompressedDataVS from './gsplat/vert/gsplatCompressedData.js';
+import gsplatCompressedSHVS from './gsplat/vert/gsplatCompressedSH.js';
+import gsplatCornerVS from './gsplat/vert/gsplatCorner.js';
+import gsplatDataVS from './gsplat/vert/gsplatData.js';
+import gsplatOutputVS from './gsplat/vert/gsplatOutput.js';
 import gsplatPS from './gsplat/frag/gsplat.js';
-// import gsplatSHVS from './gsplat/vert/gsplatSH.js';
-// import gsplatSourceVS from './gsplat/vert/gsplatSource.js';
+import gsplatSHVS from './gsplat/vert/gsplatSH.js';
+import gsplatSourceVS from './gsplat/vert/gsplatSource.js';
 import gsplatVS from './gsplat/vert/gsplat.js';
 import immediateLinePS from './internal/frag/immediateLine.js';
 import immediateLineVS from './internal/vert/immediateLine.js';
@@ -259,17 +259,17 @@ const shaderChunksWGSL = {
     fullscreenQuadVS,
     gammaPS,
     glossPS,
-    // gsplatCenterVS,
-    // gsplatCornerVS,
-    // gsplatColorVS,
+    gsplatCenterVS,
+    gsplatCornerVS,
+    gsplatColorVS,
     gsplatCommonVS,
-    // gsplatCompressedDataVS,
-    // gsplatCompressedSHVS,
-    // gsplatDataVS,
-    // gsplatOutputVS,
+    gsplatCompressedDataVS,
+    gsplatCompressedSHVS,
+    gsplatDataVS,
+    gsplatOutputVS,
     gsplatPS,
-    // gsplatSHVS,
-    // gsplatSourceVS,
+    gsplatSHVS,
+    gsplatSourceVS,
     gsplatVS,
     immediateLinePS,
     immediateLineVS,

examples/thumbnails/loaders_gsplat-sh_small.webp

@@ -0,0 +1,27 @@
+export default /* wgsl */`
+uniform matrix_model: mat4x4f;
+uniform matrix_view: mat4x4f;
+uniform matrix_projection: mat4x4f;
+
+// project the model space gaussian center to view and clip space
+fn initCenter(modelCenter: vec3f, center: ptr<function, SplatCenter>) -> bool {
+    let modelView: mat4x4f = uniform.matrix_view * uniform.matrix_model;
+    let centerView: vec4f = modelView * vec4f(modelCenter, 1.0);
+
+    // early out if splat is behind the camera
+    if (centerView.z > 0.0) {
+        return false;
+    }
+
+    var centerProj: vec4f = uniform.matrix_projection * centerView;
+
+    // ensure gaussians are not clipped by camera near and far
+    centerProj.z = clamp(centerProj.z, 0.0, abs(centerProj.w));
+
+    center.view = centerView.xyz / centerView.w;
+    center.proj = centerProj;
+    center.projMat00 = uniform.matrix_projection[0][0];
+    center.modelView = modelView;
+    return true;
+}
+`;

src/scene/gsplat/gsplat-material.js

@@ -0,0 +1,8 @@
+export default /* wgsl */`
+
+var splatColor: texture_2d<f32>;
+
+fn readColor(source: ptr<function, SplatSource>) -> vec4f {
+    return textureLoad(splatColor, source.uv, 0);
+}
+`;

src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js

@@ -0,0 +1,108 @@
+export default /* wgsl */`
+var packedTexture: texture_2d<u32>;
+var chunkTexture: texture_2d<f32>;
+
+// work values
+var<private> chunkDataA: vec4f;    // x: min_x, y: min_y, z: min_z, w: max_x
+var<private> chunkDataB: vec4f;    // x: max_y, y: max_z, z: scale_min_x, w: scale_min_y
+var<private> chunkDataC: vec4f;    // x: scale_min_z, y: scale_max_x, z: scale_max_y, w: scale_max_z
+var<private> chunkDataD: vec4f;    // x: min_r, y: min_g, z: min_b, w: max_r
+var<private> chunkDataE: vec4f;    // x: max_g, y: max_b, z: unused, w: unused
+var<private> packedData: vec4u;    // x: position bits, y: rotation bits, z: scale bits, w: color bits
+
+fn unpack111011(bits: u32) -> vec3f {
+    return (vec3f((vec3<u32>(bits) >> vec3<u32>(21u, 11u, 0u)) & vec3<u32>(0x7ffu, 0x3ffu, 0x7ffu))) / vec3f(2047.0, 1023.0, 2047.0);
+}
+
+fn unpack8888(bits: u32) -> vec4f {
+    return vec4f(
+        f32((bits >> 24u) & 0xffu),
+        f32((bits >> 16u) & 0xffu),
+        f32((bits >> 8u)  & 0xffu),
+        f32(bits         & 0xffu)
+    ) / 255.0;
+}
+
+const norm_const: f32 = 1.0 / (sqrt(2.0) * 0.5);
+
+fn unpackRotation(bits: u32) -> vec4f {
+    let a = (f32((bits >> 20u) & 0x3ffu) / 1023.0 - 0.5) * norm_const;
+    let b = (f32((bits >> 10u) & 0x3ffu) / 1023.0 - 0.5) * norm_const;
+    let c = (f32(bits & 0x3ffu) / 1023.0 - 0.5) * norm_const;
+    let m = sqrt(1.0 - (a * a + b * b + c * c));
+
+    let mode = bits >> 30u;
+    if (mode == 0u) { return vec4f(m, a, b, c); }
+    if (mode == 1u) { return vec4f(a, m, b, c); }
+    if (mode == 2u) { return vec4f(a, b, m, c); }
+    return vec4f(a, b, c, m);
+}
+
+fn quatToMat3(R: vec4f) -> mat3x3f {
+    let qw_scalar = R.x;
+    let qx_vec = R.y;
+    let qy_vec = R.z;
+    let qz_vec = R.w;
+
+    return mat3x3f(
+        vec3f(1.0 - 2.0 * (qy_vec * qy_vec + qz_vec * qz_vec),
+              2.0 * (qx_vec * qy_vec + qw_scalar * qz_vec),
+              2.0 * (qx_vec * qz_vec - qw_scalar * qy_vec)),
+
+        vec3f(2.0 * (qx_vec * qy_vec - qw_scalar * qz_vec),
+              1.0 - 2.0 * (qx_vec * qx_vec + qz_vec * qz_vec),
+              2.0 * (qy_vec * qz_vec + qw_scalar * qx_vec)),
+
+        vec3f(2.0 * (qx_vec * qz_vec + qw_scalar * qy_vec),
+              2.0 * (qy_vec * qz_vec - qw_scalar * qx_vec),
+              1.0 - 2.0 * (qx_vec * qx_vec + qy_vec * qy_vec))
+    );
+}
+
+// read center
+fn readCenter(source: ptr<function, SplatSource>) -> vec3f {
+    let tex_size_u = textureDimensions(chunkTexture, 0);
+    let w: u32 = tex_size_u.x / 5u;
+    let chunkId: u32 = source.id / 256u;
+    let chunkUV: vec2<i32> = vec2<i32>(i32((chunkId % w) * 5u), i32(chunkId / w));
+
+    // read chunk and packed compressed data
+    chunkDataA = textureLoad(chunkTexture, chunkUV + vec2<i32>(0, 0), 0);
+    chunkDataB = textureLoad(chunkTexture, chunkUV + vec2<i32>(1, 0), 0);
+    chunkDataC = textureLoad(chunkTexture, chunkUV + vec2<i32>(2, 0), 0);
+    chunkDataD = textureLoad(chunkTexture, chunkUV + vec2<i32>(3, 0), 0);
+    chunkDataE = textureLoad(chunkTexture, chunkUV + vec2<i32>(4, 0), 0);
+    packedData = textureLoad(packedTexture, source.uv, 0);
+
+    return mix(chunkDataA.xyz, vec3f(chunkDataA.w, chunkDataB.xy), unpack111011(packedData.x));
+}
+
+fn readColor(source: ptr<function, SplatSource>) -> vec4f {
+    let r = unpack8888(packedData.w);
+    return vec4f(mix(chunkDataD.xyz, vec3f(chunkDataD.w, chunkDataE.xy), r.rgb), r.w);
+}
+
+fn getRotation() -> vec4f {
+    return unpackRotation(packedData.y);
+}
+
+fn getScale() -> vec3f {
+    return exp(mix(vec3f(chunkDataB.zw, chunkDataC.x), chunkDataC.yzw, unpack111011(packedData.z)));
+}
+
+// given a rotation matrix and scale vector, compute 3d covariance A and B
+fn readCovariance(source: ptr<function, SplatSource>, covA_ptr: ptr<function, vec3f>, covB_ptr: ptr<function, vec3f>) {
+    let rot = quatToMat3(getRotation());
+    let scale = getScale();
+
+    // M = S * R
+    let M = transpose(mat3x3f(
+        scale.x * rot[0],
+        scale.y * rot[1],
+        scale.z * rot[2]
+    ));
+
+    *covA_ptr = vec3f(dot(M[0], M[0]), dot(M[0], M[1]), dot(M[0], M[2]));
+    *covB_ptr = vec3f(dot(M[1], M[1]), dot(M[1], M[2]), dot(M[2], M[2]));
+}
+`;