fix: improve cpu sort throttling and webgl2 SH

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "bevy_gaussian_splatting"
 description = "bevy gaussian splatting render pipeline plugin"
-version = "1.0.0"
+version = "2.0.0"
 edition = "2021"
 authors = ["mosure <[email protected]>"]
 license = "MIT"

src/material/spherical_harmonics.rs

@@ -46,12 +46,11 @@ const fn num_sh_coefficients(degree: usize) -> usize {
 }
 
 
-// minimize download size on web
 #[cfg(feature = "web")]
 const SH_DEGREE: usize = 0;
 
 #[cfg(not(feature = "web"))]
-const SH_DEGREE: usize = 0;
+const SH_DEGREE: usize = 3;
 
 pub const SH_CHANNELS: usize = 3;
 pub const SH_COEFF_COUNT_PER_CHANNEL: usize = num_sh_coefficients(SH_DEGREE);

src/material/spherical_harmonics.wgsl

@@ -36,7 +36,7 @@ fn spherical_harmonics_lookup(
     ray_direction: vec3<f32>,
     sh: array<f32, #{SH_COEFF_COUNT}>,
 ) -> vec3<f32> {
-    var rds = ray_direction * ray_direction;
+    let rds = ray_direction * ray_direction;
     var color = vec3<f32>(0.5);
 
     color += shc[ 0] * vec3<f32>(sh[0], sh[1], sh[2]);

src/render/gaussian.wgsl

@@ -12,7 +12,6 @@
 #import bevy_gaussian_splatting::depth::{
     depth_to_rgb,
 }
-#import bevy_gaussian_splatting::spherical_harmonics::spherical_harmonics_lookup
 #import bevy_gaussian_splatting::transform::{
     world_to_clip,
     in_frustum,
@@ -21,7 +20,7 @@
 #ifdef PACKED
 #import bevy_gaussian_splatting::packed::{
     get_position,
-    get_spherical_harmonics,
+    get_color,
     get_rotation,
     get_scale,
     get_opacity,
@@ -32,7 +31,7 @@
 #ifdef BUFFER_STORAGE
 #import bevy_gaussian_splatting::planar::{
     get_position,
-    get_spherical_harmonics,
+    get_color,
     get_rotation,
     get_scale,
     get_opacity,
@@ -43,7 +42,7 @@
 #ifdef BUFFER_TEXTURE
 #import bevy_gaussian_splatting::texture::{
     get_position,
-    get_spherical_harmonics,
+    get_color,
     get_rotation,
     get_scale,
     get_opacity,
@@ -331,7 +330,7 @@ fn vs_points(
         max_distance,
     );
 #else
-    rgb = spherical_harmonics_lookup(ray_direction, get_spherical_harmonics(splat_index));
+    rgb = get_color(splat_index, ray_direction);
 #endif
 
     // TODO: precompute color, cov2d for every gaussian. cov2d only needs a single evaluation, while color needs to be evaluated every frame in SH degree > 0 mode

src/render/packed.wgsl

@@ -1,6 +1,10 @@
 #define_import_path bevy_gaussian_splatting::packed
 
 #import bevy_gaussian_splatting::bindings::points
+#import bevy_gaussian_splatting::spherical_harmonics::{
+    spherical_harmonics_lookup,
+    srgb_to_linear,
+}
 
 
 #ifdef PACKED_F32
@@ -9,6 +13,15 @@ fn get_position(index: u32) -> vec3<f32> {
     return points[index].position_visibility.xyz;
 }
 
+fn get_color(
+    index: u32,
+    ray_direction: vec3<f32>,
+) -> vec3<f32> {
+    let sh = get_spherical_harmonics(index);
+    let color = spherical_harmonics_lookup(sh, ray_direction);
+    return srgb_to_linear(color);
+}
+
 fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
     return points[index].sh;
 }

src/render/planar.wgsl

@@ -7,6 +7,10 @@
     rotation_scale_opacity,
     scale_opacity,
 };
+#import bevy_gaussian_splatting::spherical_harmonics::{
+    spherical_harmonics_lookup,
+    srgb_to_linear,
+}
 
 
 #ifdef PLANAR_F16
@@ -15,6 +19,15 @@ fn get_position(index: u32) -> vec3<f32> {
     return position_visibility[index].xyz;
 }
 
+fn get_color(
+    index: u32,
+    ray_direction: vec3<f32>,
+) -> vec3<f32> {
+    let sh = get_spherical_harmonics(index);
+    let color = spherical_harmonics_lookup(sh, ray_direction);
+    return srgb_to_linear(color);
+}
+
 fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
     var coefficients: array<f32, #{SH_COEFF_COUNT}>;
 

src/render/texture.wgsl

@@ -8,6 +8,10 @@
     rotation_scale_opacity,
     scale_opacity,
 };
+#import bevy_gaussian_splatting::spherical_harmonics::{
+    shc,
+    srgb_to_linear,
+}
 
 
 fn location(index: u32) -> vec2<i32> {
@@ -30,15 +34,199 @@ fn get_position(index: u32) -> vec3<f32> {
     return sample.xyz;
 }
 
+fn get_sh_vec(
+    index: u32,
+    plane: i32,
+) -> vec4<u32> {
+#if SH_VEC4_PLANES == 1
+    return textureLoad(
+        spherical_harmonics,
+        location(index),
+        plane,
+    );
+#else
+    return textureLoad(
+        spherical_harmonics,
+        location(index),
+        plane,
+        0,
+    );
+#endif
+}
+
+#ifdef WEBGL2
+fn get_color(
+    index: u32,
+    ray_direction: vec3<f32>,
+) -> vec3<f32> {
+    let s0 = get_sh_vec(index, 0);
+
+    let v0 = unpack2x16float(s0.x);
+    let v1 = unpack2x16float(s0.y);
+    let v2 = unpack2x16float(s0.z);
+    let v3 = unpack2x16float(s0.w);
+
+    let rds = ray_direction * ray_direction;
+    var color = vec3<f32>(0.5);
+
+    color += shc[ 0] * vec3<f32>(
+        v0.x,
+        v0.y,
+        v1.x,
+    );
+
+#if SH_COEFF_COUNT > 11
+    let r1 = vec3<f32>(
+        v1.y,
+        v2.x,
+        v2.y,
+    );
+
+    let s1 = get_sh_vec(index, 1);
+    let v4 = unpack2x16float(s1.x);
+    let v5 = unpack2x16float(s1.y);
+    let v6 = unpack2x16float(s1.z);
+    let v7 = unpack2x16float(s1.w);
+
+    let r2 = vec3<f32>(
+        v3.x,
+        v3.y,
+        v4.x,
+    );
+
+    let r3 = vec3<f32>(
+        v4.y,
+        v5.x,
+        v5.y,
+    );
+
+    color += shc[ 1] * r1 * ray_direction.y;
+    color += shc[ 2] * r2 * ray_direction.z;
+    color += shc[ 3] * r3 * ray_direction.x;
+#endif
+
+#if SH_COEFF_COUNT > 26
+    let r4 = vec3<f32>(
+        v6.x,
+        v6.y,
+        v7.x,
+    );
+
+    let s2 = get_sh_vec(index, 2);
+    let v8 = unpack2x16float(s2.x);
+    let v9 = unpack2x16float(s2.y);
+    let v10 = unpack2x16float(s2.z);
+    let v11 = unpack2x16float(s2.w);
+
+    let r5 = vec3<f32>(
+        v7.y,
+        v8.x,
+        v8.y,
+    );
+
+    let r6 = vec3<f32>(
+        v9.x,
+        v9.y,
+        v10.x,
+    );
+
+    let r7 = vec3<f32>(
+        v10.y,
+        v11.x,
+        v11.y,
+    );
+
+    let s3 = get_sh_vec(index, 3);
+    let v12 = unpack2x16float(s3.x);
+    let v13 = unpack2x16float(s3.y);
+    let v14 = unpack2x16float(s3.z);
+    let v15 = unpack2x16float(s3.w);
+
+    let r8 = vec3<f32>(
+        v12.x,
+        v12.y,
+        v13.x,
+    );
+
+    color += shc[ 4] * r4 * ray_direction.x * ray_direction.y;
+    color += shc[ 5] * r5 * ray_direction.y * ray_direction.z;
+    color += shc[ 6] * r6 * (2.0 * rds.z - rds.x - rds.y);
+    color += shc[ 7] * r7 * ray_direction.x * ray_direction.z;
+    color += shc[ 8] * r8 * (rds.x - rds.y);
+#endif
+
+#if SH_COEFF_COUNT > 47
+    let r9 = vec3<f32>(
+        v13.y,
+        v14.x,
+        v14.y,
+    );
+
+    let s4 = get_sh_vec(index, 4);
+    let v16 = unpack2x16float(s4.x);
+    let v17 = unpack2x16float(s4.y);
+    let v18 = unpack2x16float(s4.z);
+    let v19 = unpack2x16float(s4.w);
+
+    let r10 = vec3<f32>(
+        v15.x,
+        v15.y,
+        v16.x,
+    );
+
+    let r11 = vec3<f32>(
+        v16.y,
+        v17.x,
+        v17.y,
+    );
+
+    let r12 = vec3<f32>(
+        v18.x,
+        v18.y,
+        v19.x,
+    );
+
+    let s5 = get_sh_vec(index, 5);
+    let v20 = unpack2x16float(s5.x);
+    let v21 = unpack2x16float(s5.y);
+    let v22 = unpack2x16float(s5.z);
+    let v23 = unpack2x16float(s5.w);
+
+    let r13 = vec3<f32>(
+        v19.y,
+        v20.x,
+        v20.y,
+    );
+
+    let r14 = vec3<f32>(
+        v21.x,
+        v21.y,
+        v22.x,
+    );
+
+    let r15 = vec3<f32>(
+        v22.y,
+        v23.x,
+        v23.y,
+    );
+
+    color += shc[ 9] * r9 * ray_direction.y * (3.0 * rds.x - rds.y);
+    color += shc[10] * r10 * ray_direction.x * ray_direction.y * ray_direction.z;
+    color += shc[11] * r11 * ray_direction.y * (4.0 * rds.z - rds.x - rds.y);
+    color += shc[12] * r12 * ray_direction.z * (2.0 * rds.z - 3.0 * rds.x - 3.0 * rds.y);
+    color += shc[13] * r13 * ray_direction.x * (4.0 * rds.z - rds.x - rds.y);
+    color += shc[14] * r14 * ray_direction.z * (rds.x - rds.y);
+    color += shc[15] * r15 * ray_direction.x * (rds.x - 3.0 * rds.y);
+#endif
+
+    return srgb_to_linear(color);
+}
+#else
 fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
     var coefficients: array<f32, #{SH_COEFF_COUNT}>;
 
     for (var i = 0u; i < #{SH_VEC4_PLANES}u; i = i + 1u) {
-        let sample = textureLoad(
-            spherical_harmonics,
-            location(index),
-            i32(i),
-        );
+        let sample = get_sh_vec(index, i32(i));
 
         let v0 = unpack2x16float(sample.x);
         let v1 = unpack2x16float(sample.y);
@@ -62,6 +250,16 @@ fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
     return coefficients;
 }
 
+fn get_color(
+    index: u32,
+    ray_direction: vec3<f32>,
+) -> vec3<f32> {
+    let sh = get_spherical_harmonics(index);
+    let color = spherical_harmonics_lookup(sh, ray_direction);
+    return srgb_to_linear(color);
+}
+#endif
+
 fn get_rotation(index: u32) -> vec4<f32> {
     let sample = textureLoad(
         rotation_scale_opacity,

src/sort/mod.rs

@@ -43,7 +43,7 @@ pub mod radix;
 pub mod rayon;
 
 #[cfg(feature = "sort_std")]
-pub mod std;
+pub mod std; // rename to std_sort.rs to avoid name conflict with std crate
 
 
 assert_cfg!(