Wgsl GS chunk updates(#7664)

Author: slimbuck

src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplatCommon.js

@@ -25,6 +25,20 @@ struct SplatCorner {
     #endif
 }
 
+fn quatToMat3(R: vec4<f32>) -> mat3x3<f32> {
+    let R2: vec4<f32> = R + R;
+    let X: f32       = R2.x * R.w;
+    let Y: vec4<f32> = R2.y * R;
+    let Z: vec4<f32> = R2.z * R;
+    let W: f32       = R2.w * R.w;
+
+    return mat3x3<f32>(
+        1.0 - Z.z - W,  Y.z + X,      Y.w - Z.x,
+        Y.z - X,        1.0 - Y.y - W, Z.w + Y.x,
+        Y.w + Z.x,      Z.w - Y.x,     1.0 - Y.y - Z.z
+    );
+}
+
 #if GSPLAT_COMPRESSED_DATA == true
     #include "gsplatCompressedDataVS"
     #include "gsplatCompressedSHVS"

src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplatCompressedData.js

@@ -38,27 +38,6 @@ fn unpackRotation(bits: u32) -> vec4f {
     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);

src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplatData.js

@@ -13,34 +13,6 @@ fn readCenter(source: ptr<function, SplatSource>) -> vec3f {
     return bitcast<vec3f>(tA.xyz);
 }
 
-fn quatToMat3(R: vec4f) -> mat3x3f {
-    let x: f32 = R.w;
-    let y: f32 = R.x;
-    let z: f32 = R.y;
-    let w: f32 = R.z;
-
-    // pre-calculate squares and products
-    let xx: f32 = x * x; let yy: f32 = y * y; let zz: f32 = z * z; let ww: f32 = w * w;
-    let xy: f32 = x * y; let xz: f32 = x * z; let xw: f32 = x * w;
-    let yz: f32 = y * z; let yw: f32 = y * w; let zw: f32 = z * w;
-
-    return mat3x3f(
-        vec3f(
-            1.0 - 2.0 * (zz + ww),
-            2.0 * (yz + xw),
-            2.0 * (yw - xz)
-        ), vec3f(
-            2.0 * (yz - xw),
-            1.0 - 2.0 * (yy + ww),
-            2.0 * (zw + xy)
-        ), vec3f(
-            2.0 * (yw + xz),
-            2.0 * (zw - xy),
-            1.0 - 2.0 * (yy + zz)
-        )
-    );
-}
-
 fn unpackRotation(packed: vec3f) -> vec4f {
     return vec4f(packed.xyz, sqrt(max(0.0, 1.0 - dot(packed, packed))));
 }
@@ -49,7 +21,7 @@ fn unpackRotation(packed: vec3f) -> vec4f {
 fn readCovariance(source: ptr<function, SplatSource>, covA_ptr: ptr<function, vec3f>, covB_ptr: ptr<function, vec3f>) {
     let tB: vec4f = textureLoad(transformB, source.uv, 0);
 
-    let rot: mat3x3f = quatToMat3(unpackRotation(vec3f(unpack2x16float(bitcast<u32>(tAw)), tB.w)));
+    let rot: mat3x3f = quatToMat3(unpackRotation(vec3f(unpack2x16float(bitcast<u32>(tAw)), tB.w)).wxyz);
     let scale: vec3f = tB.xyz;
 
     // M = S * R

src/scene/shader-lib/chunks-wgsl/gsplat/vert/gsplatSH.js

@@ -7,41 +7,73 @@ fn unpack111011s(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)) * 2.0 - 1.0;
 }
 
+struct ScaleAndSH {
+    scale: f32,
+    a: vec3f,
+    b: vec3f,
+    c: vec3f
+};
+
 // fetch quantized spherical harmonic coefficients
-fn fetchScale(t_in: vec4<u32>, scale: ptr<function, f32>, a: ptr<function, vec3f>, b: ptr<function, vec3f>, c: ptr<function, vec3f>) {
-    *scale = bitcast<f32>(t_in.x);
-    *a = unpack111011s(t_in.y);
-    *b = unpack111011s(t_in.z);
-    *c = unpack111011s(t_in.w);
+fn fetchScale(t_in: vec4<u32>) -> ScaleAndSH {
+    var result: ScaleAndSH;
+    result.scale = bitcast<f32>(t_in.x);
+    result.a = unpack111011s(t_in.y);
+    result.b = unpack111011s(t_in.z);
+    result.c = unpack111011s(t_in.w);
+    return result;
 }
 
+struct SH {
+    a: vec3f,
+    b: vec3f,
+    c: vec3f,
+    d: vec3f
+};
+
 // fetch quantized spherical harmonic coefficients
-fn fetch4(t_in: vec4<u32>, a: ptr<function, vec3f>, b: ptr<function, vec3f>, c: ptr<function, vec3f>, d: ptr<function, vec3f>) {
-    *a = unpack111011s(t_in.x);
-    *b = unpack111011s(t_in.y);
-    *c = unpack111011s(t_in.z);
-    *d = unpack111011s(t_in.w);
+fn fetch4(t_in: vec4<u32>) -> SH {
+    var result: SH;
+    result.a = unpack111011s(t_in.x);
+    result.b = unpack111011s(t_in.y);
+    result.c = unpack111011s(t_in.z);
+    result.d = unpack111011s(t_in.w);
+    return result;
 }
 
-fn fetch1(t_in: u32, a: ptr<function, vec3f>) {
-    *a = unpack111011s(t_in);
+fn fetch1(t_in: u32) -> vec3f {
+    return unpack111011s(t_in);
 }
 
 #if SH_BANDS == 1
     var splatSH_1to3: texture_2d<u32>;
 
     fn readSHData(source: ptr<function, SplatSource>, sh: ptr<function, array<vec3f, 3>>, scale: ptr<function, f32>) {
-        fetchScale(textureLoad(splatSH_1to3, source.uv, 0), scale, &sh[0], &sh[1], &sh[2]);
+        let result = fetchScale(textureLoad(splatSH_1to3, source.uv, 0));
+        *scale = result.scale;
+        sh[0] = result.a;
+        sh[1] = result.b;
+        sh[2] = result.c;
     }
 #elif SH_BANDS == 2
     var splatSH_1to3: texture_2d<u32>;
     var splatSH_4to7: texture_2d<u32>;
     var splatSH_8to11: texture_2d<u32>;
 
     fn readSHData(source: ptr<function, SplatSource>, sh: ptr<function, array<vec3f, 8>>, scale: ptr<function, f32>) {
-        fetchScale(textureLoad(splatSH_1to3, source.uv, 0), scale, &sh[0], &sh[1], &sh[2]);
-        fetch4(textureLoad(splatSH_4to7, source.uv, 0), &sh[3], &sh[4], &sh[5], &sh[6]);
-        fetch1(textureLoad(splatSH_8to11, source.uv, 0).x, &sh[7]);
+        let first: ScaleAndSH = fetchScale(textureLoad(splatSH_1to3, source.uv, 0));
+        *scale = first.scale;
+        sh[0] = first.a;
+        sh[1] = first.b;
+        sh[2] = first.c;
+
+        let second: SH = fetch4(textureLoad(splatSH_4to7, source.uv, 0));
+        sh[3] = second.a;
+        sh[4] = second.b;
+        sh[5] = second.c;
+        sh[6] = second.d;
+
+        sh[7] = fetch1(textureLoad(splatSH_8to11, source.uv, 0).x);
     }
 #else
     var splatSH_1to3: texture_2d<u32>;
@@ -50,10 +82,29 @@ fn fetch1(t_in: u32, a: ptr<function, vec3f>) {
     var splatSH_12to15: texture_2d<u32>;
 
     fn readSHData(source: ptr<function, SplatSource>, sh: ptr<function, array<vec3f, 15>>, scale: ptr<function, f32>) {
-        fetchScale(textureLoad(splatSH_1to3, source.uv, 0), scale, &sh[0], &sh[1], &sh[2]);
-        fetch4(textureLoad(splatSH_4to7, source.uv, 0), &sh[3], &sh[4], &sh[5], &sh[6]);
-        fetch4(textureLoad(splatSH_8to11, source.uv, 0), &sh[7], &sh[8], &sh[9], &sh[10]);
-        fetch4(textureLoad(splatSH_12to15, source.uv, 0), &sh[11], &sh[12], &sh[13], &sh[14]);
+        let first: ScaleAndSH = fetchScale(textureLoad(splatSH_1to3, source.uv, 0));
+        *scale = first.scale;
+        sh[0] = first.a;
+        sh[1] = first.b;
+        sh[2] = first.c;
+
+        let second: SH = fetch4(textureLoad(splatSH_4to7, source.uv, 0));
+        sh[3] = second.a;
+        sh[4] = second.b;
+        sh[5] = second.c;
+        sh[6] = second.d;
+
+        let third: SH = fetch4(textureLoad(splatSH_8to11, source.uv, 0));
+        sh[7] = third.a;
+        sh[8] = third.b;
+        sh[9] = third.c;
+        sh[10] = third.d;
+
+        let fourth: SH = fetch4(textureLoad(splatSH_12to15, source.uv, 0));
+        sh[11] = fourth.a;
+        sh[12] = fourth.b;
+        sh[13] = fourth.c;
+        sh[14] = fourth.d;
     }
 #endif