Source geometry from work buffer for SOG spherical harmonics updates

Color-only (spherical harmonics) work buffer updates for SOG resources
now reconstruct geometry from the work buffer's stored world-space
transform data instead of re-reading and dequantizing the source SOG
textures (means_l/means_u/quats/scales).

- Add gsplatWorkBufferGeometry chunk (GLSL + WGSL) that decodes the
  stored center/rotation/scale (both compact and packed work buffer
  encodings) and converts them back to local space, preserving the
  getCenter/getRotation/getScale API used by setWorkBufferModifier code.
- Add a work-buffer-geometry path to gsplatCopyToWorkbuffer and unify
  the shared color/SH tail across both paths.
- Add a dataFormat layout tag to GSplatFormat (set by GSplatParams) so
  the decode variant is selected from the format identity rather than by
  inspecting stream pixel formats.
- Gate behind GSplatResourceBase#supportsWorkBufferGeometry (SOG only).

Author: Martin Valigursky

src/scene/gsplat-unified/gsplat-params.js

@@ -120,6 +120,7 @@ class GSplatParams {
         }
 
         format.allowStreamRemoval = true;
+        format.dataFormat = dataFormat;
         return format;
     }
 

src/scene/gsplat-unified/gsplat-work-buffer-render-pass.js

@@ -17,6 +17,7 @@ import { TextureUtils } from '../../platform/graphics/texture-utils.js';
  */
 
 const _viewMat = new Mat4();
+const _invModelMat = new Mat4();
 const _modelScale = new Vec3();
 const _modelRotation = new Quat();
 const _tmpSize = new Vec2();
@@ -58,6 +59,9 @@ class GSplatWorkBufferRenderPass extends RenderPass {
     /** @type {Float32Array} */
     _modelRotationData = new Float32Array(4);
 
+    /** @type {Float32Array} */
+    _cameraPositionData = new Float32Array(3);
+
     /** @type {Int32Array} */
     _textureSize = new Int32Array(2);
 
@@ -204,6 +208,19 @@ class GSplatWorkBufferRenderPass extends RenderPass {
         const viewMat = _viewMat.copy(viewInvMat).invert();
         device.scope.resolve('matrix_view').setValue(viewMat.data);
 
+        // work-buffer-sourced geometry inputs for color-only (SH) updates
+        // (used by shaders compiled with GSPLAT_WORKBUFFER_GEOMETRY)
+        if (this.colorOnly) {
+            device.scope.resolve('uWorkBufferTransformA').setValue(this.workBuffer.getTexture('dataTransformA'));
+            device.scope.resolve('uWorkBufferTransformB').setValue(this.workBuffer.getTexture('dataTransformB'));
+
+            const cameraPos = cameraNode.getPosition();
+            this._cameraPositionData[0] = cameraPos.x;
+            this._cameraPositionData[1] = cameraPos.y;
+            this._cameraPositionData[2] = cameraPos.z;
+            device.scope.resolve('uCameraPosition').setValue(this._cameraPositionData);
+        }
+
         // render each splat info
         for (let i = 0; i < splats.length; i++) {
             const count = _partialData[i * 2 + 1];
@@ -281,6 +298,13 @@ class GSplatWorkBufferRenderPass extends RenderPass {
         scope.resolve('model_scale').setValue(this._modelScaleData);
         scope.resolve('model_rotation').setValue(this._modelRotationData);
 
+        // inverse model matrix, used by GSPLAT_WORKBUFFER_GEOMETRY shaders to convert stored
+        // world-space data back to local space
+        if (this.colorOnly) {
+            _invModelMat.copy(worldTransform).invert();
+            scope.resolve('matrix_model_inverse').setValue(_invModelMat.data);
+        }
+
         // Set placement ID for picking (unconditionally - cheap even if shader doesn't use it)
         scope.resolve('uId').setValue(splatInfo.placementId);
 

src/scene/gsplat/gsplat-format.js

@@ -111,6 +111,17 @@ class GSplatFormat {
      */
     allowStreamRemoval = false;
 
+    /**
+     * Work buffer data layout identifier (one of GSPLATDATA_*), or null for resource formats.
+     * Set by {@link GSplatParams} when creating a work buffer format. Identifies how transform
+     * data is encoded, so consumers (e.g. the work-buffer-sourced spherical harmonics update)
+     * can select the matching decode without inspecting individual stream pixel formats.
+     *
+     * @type {string|null}
+     * @ignore
+     */
+    dataFormat = null;
+
     /**
      * Extra streams added via addExtraStreams(). For resource formats, streams can only be
      * added, never removed. For work buffer formats (where {@link allowStreamRemoval} is true),

src/scene/gsplat/gsplat-resource-base.js

@@ -1,5 +1,6 @@
 import { Debug } from '../../core/debug.js';
 import { BoundingBox } from '../../core/shape/bounding-box.js';
+import { GSPLATDATA_COMPACT } from '../constants.js';
 import { Mesh } from '../mesh.js';
 import { ShaderMaterial } from '../materials/shader-material.js';
 import { WorkBufferRenderInfo } from '../gsplat-unified/gsplat-work-buffer.js';
@@ -241,6 +242,19 @@ class GSplatResourceBase {
         }
     }
 
+    /**
+     * True when this resource's color-only work buffer updates (spherical harmonics refresh) can
+     * source geometry from the work buffer itself instead of re-reading the source textures. The
+     * resource format's read chunk must compile out getCenter/getRotation/getScale when
+     * GSPLAT_WORKBUFFER_GEOMETRY is defined (see gsplatWorkBufferGeometryPS chunk).
+     *
+     * @type {boolean}
+     * @ignore
+     */
+    get supportsWorkBufferGeometry() {
+        return false;
+    }
+
     /**
      * Get or create a QuadRender for rendering to work buffer.
      *
@@ -257,7 +271,17 @@ class GSplatResourceBase {
         // configure defines to fetch cached data
         this.configureMaterialDefines(tempMap);
         tempMap.set('GSPLAT_LOD', '');
-        if (colorOnly) tempMap.set('GSPLAT_COLOR_ONLY', '');
+        if (colorOnly) {
+            tempMap.set('GSPLAT_COLOR_ONLY', '');
+
+            // source geometry from the work buffer instead of source textures
+            if (this.supportsWorkBufferGeometry) {
+                tempMap.set('GSPLAT_WORKBUFFER_GEOMETRY', '');
+                if (workBufferFormat.dataFormat === GSPLATDATA_COMPACT) {
+                    tempMap.set('GSPLAT_WORKBUFFER_COMPACT', '');
+                }
+            }
+        }
 
         let definesKey = '';
         for (const [k, v] of tempMap) {

src/scene/gsplat/gsplat-sog-resource.js

@@ -147,6 +147,12 @@ class GSplatSogResource extends GSplatResourceBase {
             defines.set('SOG_V2', '');
         }
     }
+
+    // SOG geometry getters compile out under GSPLAT_WORKBUFFER_GEOMETRY (see sog.js chunk),
+    // letting color-only (SH) updates skip the means/quats/scales source reads
+    get supportsWorkBufferGeometry() {
+        return true;
+    }
 }
 
 export { GSplatSogResource };

src/scene/shader-lib/glsl/chunks/gsplat/frag/gsplatCopyToWorkbuffer.js

@@ -3,6 +3,11 @@ export default /* glsl */`
 
 #define GSPLAT_CENTER_NOPROJ
 
+// pre-computed model matrix decomposition (declared before includes, used by
+// gsplatWorkBufferGeometryPS)
+uniform vec3 model_scale;
+uniform vec4 model_rotation;  // (x,y,z,w) format
+
 #include "gsplatHelpersVS"
 #include "gsplatFormatVS"
 #include "gsplatStructsVS"
@@ -11,6 +16,7 @@ export default /* glsl */`
 #include "gsplatEvalSHVS"
 #include "gsplatQuatToMat3VS"
 #include "gsplatReadVS"
+#include "gsplatWorkBufferGeometryPS"
 #include "gsplatWorkBufferOutputVS"
 #include "gsplatWriteVS"
 #include "gsplatModifyVS"
@@ -20,10 +26,6 @@ flat varying ivec4 vSubDraw;
 
 uniform vec3 uColorMultiply;
 
-// pre-computed model matrix decomposition
-uniform vec3 model_scale;
-uniform vec4 model_rotation;  // (x,y,z,w) format
-
 #ifdef GSPLAT_ID
     uniform uint uId;
 #endif
@@ -37,43 +39,71 @@ void main(void) {
     // Initialize global splat for format read functions
     setSplat(originalIndex);
 
-    // read center in local space
-    vec3 modelCenter = getCenter();
+    // World-space geometry to store. Rotation/scale default to identity and are only computed on
+    // the full-render path; under GSPLAT_COLOR_ONLY they are ignored by writeSplat (DCE'd).
+    vec3 worldCenter;
+    vec4 worldRotation = vec4(0.0, 0.0, 0.0, 1.0);
+    vec3 worldScale = vec3(1.0);
+    #if SH_BANDS > 0
+        vec3 dir;   // model-space view direction
+    #endif
+
+    #ifdef GSPLAT_WORKBUFFER_GEOMETRY
 
-    // compute world-space center for storage
-    vec3 worldCenter = (matrix_model * vec4(modelCenter, 1.0)).xyz;
-    SplatCenter center;
-    initCenter(modelCenter, center);
+        // Color-only update sourcing geometry from previously written work buffer data at this
+        // destination pixel. The stored center already includes the model transform and
+        // modifySplatCenter / modifySplatRotationScale, so neither is re-applied here.
+        initWorkBufferGeometry(ivec2(gl_FragCoord.xy));
+        worldCenter = workBufferWorldCenter();
 
-    // Get source rotation and scale
-    // getRotation() returns (w,x,y,z) format, convert to (x,y,z,w) for quatMul
-    vec4 srcRotation = getRotation().yzwx;
-    vec3 srcScale = getScale();
+        #if SH_BANDS > 0
+            // model-space view direction (matches the source path up to non-uniform model scale)
+            dir = normalize(quatRotateInv(model_rotation, worldCenter - uCameraPosition));
+        #endif
 
-    // Combine: world = model * source (both in x,y,z,w format)
-    vec4 worldRotation = quatMul(model_rotation, srcRotation);
-    // Ensure w is positive so sqrt() reconstruction works correctly
-    // (quaternions q and -q represent the same rotation)
-    if (worldRotation.w < 0.0) {
-        worldRotation = -worldRotation;
-    }
-    vec3 worldScale = model_scale * srcScale;
+    #else
 
-    // Apply custom center modification
-    vec3 originalCenter = worldCenter;
-    modifySplatCenter(worldCenter);
+        // read center in local space
+        vec3 modelCenter = getCenter();
 
-    // Apply custom rotation/scale modification
-    modifySplatRotationScale(originalCenter, worldCenter, worldRotation, worldScale);
+        // compute world-space center for storage
+        worldCenter = (matrix_model * vec4(modelCenter, 1.0)).xyz;
+        SplatCenter center;
+        initCenter(modelCenter, center);
+
+        // Get source rotation and scale
+        // getRotation() returns (w,x,y,z) format, convert to (x,y,z,w) for quatMul
+        vec4 srcRotation = getRotation().yzwx;
+        vec3 srcScale = getScale();
+
+        // Combine: world = model * source (both in x,y,z,w format)
+        worldRotation = quatMul(model_rotation, srcRotation);
+        // Ensure w is positive so sqrt() reconstruction works correctly
+        // (quaternions q and -q represent the same rotation)
+        if (worldRotation.w < 0.0) {
+            worldRotation = -worldRotation;
+        }
+        worldScale = model_scale * srcScale;
+
+        // Apply custom center modification
+        vec3 originalCenter = worldCenter;
+        modifySplatCenter(worldCenter);
+
+        // Apply custom rotation/scale modification
+        modifySplatRotationScale(originalCenter, worldCenter, worldRotation, worldScale);
+
+        #if SH_BANDS > 0
+            // calculate the model-space view direction
+            dir = normalize(center.view * mat3(center.modelView));
+        #endif
+
+    #endif
 
     // read color
     vec4 color = getColor();
 
     // evaluate spherical harmonics
     #if SH_BANDS > 0
-        // calculate the model-space view direction
-        vec3 dir = normalize(center.view * mat3(center.modelView));
-
         // read sh coefficients
         vec3 sh[SH_COEFFS];
         float scale;
@@ -88,7 +118,8 @@ void main(void) {
 
     color.xyz *= uColorMultiply;
 
-    // write color + transform using format-specific encoding
+    // write color + transform using format-specific encoding (rotation/scale ignored under
+    // GSPLAT_COLOR_ONLY)
     writeSplat(worldCenter, worldRotation, worldScale, color);
 
     #ifdef GSPLAT_ID

src/scene/shader-lib/glsl/chunks/gsplat/frag/gsplatWorkBufferGeometry.js

@@ -0,0 +1,98 @@
+// Work-buffer-backed geometry for color-only (spherical harmonics) work buffer updates.
+// Instead of re-reading geometry from the source format textures, getCenter/getRotation/getScale
+// read back the world-space data previously written to the work buffer at the destination pixel,
+// and convert it to the splat's local space so user modifier code keeps working. The stored data
+// already includes the model transform and modifySplatCenter / modifySplatRotationScale, so
+// neither is re-applied by this path.
+export default /* glsl */`
+#ifdef GSPLAT_WORKBUFFER_GEOMETRY
+
+    // world-space transform data previously written to the work buffer (see gsplatWriteVS)
+    uniform highp usampler2D uWorkBufferTransformA;
+    uniform highp usampler2D uWorkBufferTransformB;
+
+    // inverse of matrix_model, to convert stored world-space data back to local space
+    uniform mat4 matrix_model_inverse;
+
+    // world-space camera position
+    uniform vec3 uCameraPosition;
+
+    ivec2 wbCoord;
+    uvec4 wbTransformA;
+
+    // cache transformA at the destination pixel; must be called before any getters
+    void initWorkBufferGeometry(ivec2 coord) {
+        wbCoord = coord;
+        wbTransformA = texelFetch(uWorkBufferTransformA, coord, 0);
+    }
+
+    vec3 workBufferWorldCenter() {
+        return vec3(uintBitsToFloat(wbTransformA.x), uintBitsToFloat(wbTransformA.y), uintBitsToFloat(wbTransformA.z));
+    }
+
+    // world-space rotation (x,y,z,w), decoded to match the work buffer write encoding
+    vec4 workBufferWorldRotation() {
+        #ifdef GSPLAT_WORKBUFFER_COMPACT
+            // half-angle projected quaternion, 11+11+10 bits (see containerCompactWrite)
+            uint data = texelFetch(uWorkBufferTransformB, wbCoord, 0).x;
+            vec3 p = vec3(
+                float(data & 0x7FFu) / 2047.0 * 2.0 - 1.0,
+                float((data >> 11u) & 0x7FFu) / 2047.0 * 2.0 - 1.0,
+                float((data >> 22u) & 0x3FFu) / 1023.0 * 2.0 - 1.0
+            );
+            float d = dot(p, p);
+            return vec4(sqrt(max(0.0, 2.0 - d)) * p, 1.0 - d);
+        #else
+            // rotation.xy in transformA.w, rotation.z in transformB.x (see containerPackedWrite)
+            vec2 rotXY = unpackHalf2x16(wbTransformA.w);
+            vec3 r = vec3(rotXY, unpackHalf2x16(texelFetch(uWorkBufferTransformB, wbCoord, 0).x).x);
+            return vec4(r, sqrt(max(0.0, 1.0 - dot(r, r))));
+        #endif
+    }
+
+    vec3 workBufferWorldScale() {
+        #ifdef GSPLAT_WORKBUFFER_COMPACT
+            // log-encoded scale, 3x8 bits: 0 = true zero, 1-255 maps to e^-12..e^9 (see containerCompactWrite)
+            uint data = wbTransformA.w;
+            float sx = float(data & 0xFFu);
+            float sy = float((data >> 8u) & 0xFFu);
+            float sz = float((data >> 16u) & 0xFFu);
+            const float logRange = 21.0 / 255.0;
+            const float logMin = -12.0;
+            return vec3(
+                sx == 0.0 ? 0.0 : exp(sx * logRange + logMin),
+                sy == 0.0 ? 0.0 : exp(sy * logRange + logMin),
+                sz == 0.0 ? 0.0 : exp(sz * logRange + logMin)
+            );
+        #else
+            uvec2 b = texelFetch(uWorkBufferTransformB, wbCoord, 0).xy;
+            return vec3(unpackHalf2x16(b.x).y, unpackHalf2x16(b.y));
+        #endif
+    }
+
+    // rotate vector by the inverse of unit quaternion q (x,y,z,w)
+    vec3 quatRotateInv(vec4 q, vec3 v) {
+        vec3 t = -q.xyz;
+        return v + 2.0 * cross(t, cross(t, v) + q.w * v);
+    }
+
+    // Source-format-compatible getters for user modifier code: local-space values reconstructed
+    // from the stored world-space data (quantized by the work buffer format, so rotation and
+    // scale are approximate).
+    vec3 getCenter() {
+        return (matrix_model_inverse * vec4(workBufferWorldCenter(), 1.0)).xyz;
+    }
+
+    // returns (w,x,y,z) to match the source format getRotation convention
+    vec4 getRotation() {
+        vec4 worldRotation = workBufferWorldRotation();
+        vec4 localRotation = quatMul(vec4(-model_rotation.xyz, model_rotation.w), worldRotation);
+        return localRotation.wxyz;
+    }
+
+    vec3 getScale() {
+        return workBufferWorldScale() / model_scale;
+    }
+
+#endif
+`;

src/scene/shader-lib/glsl/chunks/gsplat/vert/formats/sog.js

@@ -24,6 +24,9 @@ const float norm = sqrt(2.0);
     float lutShN(int b)    { return texelFetch(sogCodebook, ivec2(b, 0), 0).b; }
 #endif
 
+// geometry getters are provided by gsplatWorkBufferGeometryPS when sourcing from the work buffer
+#ifndef GSPLAT_WORKBUFFER_GEOMETRY
+
 // read the model-space center of the gaussian (16-bit per-axis, low + high byte)
 vec3 getCenter() {
     vec3 l = texelFetch(means_l, splat.uv, 0).xyz;
@@ -55,6 +58,8 @@ vec3 getScale() {
     return exp(logS);
 }
 
+#endif
+
 vec4 getColor() {
     vec4 c = texelFetch(sh0, splat.uv, 0);
     #ifdef SOG_V2