feat: improve gcloud load times (#28)

* feat: flexbuffer gcloud format

* chore: allow bugfix version update

* docs: credits

* refactor: codec abstraction and benchmarks

* chore: correct workflow name

* fix: tests and bench

* chore: trigger bench on PR synchronize

* fix: disable libtest #[bench]

* fix: failure due to main comparison

* feat: bincode2 + flate2 format

* refactor: optional ply loading

.github/workflows/bench.yml

@@ -0,0 +1,41 @@
+name: bench
+
+on:
+  pull_request:
+    types: [ labeled, synchronize ]
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: 1
+
+jobs:
+  bench:
+    if: contains(github.event.pull_request.labels.*.name, 'bench')
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [windows-latest, macos-latest]
+
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/cache@v3
+      with:
+        path: |
+          ~/.cargo/bin/
+          ~/.cargo/registry/index/
+          ~/.cargo/registry/cache/
+          ~/.cargo/git/db/
+          target/
+        key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }}
+
+    - name: io benchmark
+      uses: boa-dev/[email protected]
+      with:
+        benchName: "io"
+        branchName: ${{ github.base_ref }}
+        token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/test.yml

@@ -33,9 +33,15 @@ jobs:
           target/
         key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }}
 
+    - name: check
+      run: cargo check
+
     - name: build
       run: cargo build
 
+    - name: build_tools
+      run: cargo build --bin ply_to_gcloud
+
     - name: lint
       run: cargo clippy
 

Cargo.toml

@@ -14,22 +14,31 @@ exclude = [".devcontainer", ".github", "docs", "dist", "build", "assets", "credi
 default-run = "viewer"
 
 
-# TODO: separate dev-dependencies
+[features]
+default = ["io_flexbuffers", "io_ply"]
+
+io_bincode2 = ["bincode2", "flate2"]
+io_flexbuffers = ["flexbuffers"]
+io_ply = ["ply-rs"]
+
+
 [dependencies]
 bevy-inspector-egui = "0.21"
-bevy_panorbit_camera = "0.9.0"
-bincode2 = "2.0.1"
-bytemuck = "1.14.0"
-flate2 = "1.0.28"
-ply-rs = "0.1.3"
-rand = "0.8.5"
-serde = "1.0.189"
+bevy_panorbit_camera = "0.9"
+bincode2 = { version = "2.0", optional = true }
+byte-unit = "4.0"
+bytemuck = "1.14"
+flate2  = { version = "1.0", optional = true }
+flexbuffers = { version = "2.0", optional = true }
+ply-rs = { version = "0.1", optional = true }
+rand = "0.8"
+serde = "1.0"
 wgpu = "0.17.1"
 
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
-console_error_panic_hook = "0.1.7"
-wasm-bindgen = "0.2.87"
+console_error_panic_hook = "0.1"
+wasm-bindgen = "0.2"
 
 
 # TODO: use minimal bevy features
@@ -49,6 +58,10 @@ features = [
 ]
 
 
+[dev-dependencies]
+criterion = { version = "0.5", features = ["html_reports"] }
+
+
 [profile.dev.package."*"]
 opt-level = 3
 
@@ -77,3 +90,9 @@ path = "viewer/viewer.rs"
 [[bin]]
 name = "ply_to_gcloud"
 path = "tools/ply_to_gcloud.rs"
+required-features = ["io_ply"]
+
+
+[[bench]]
+name = "io"
+harness = false

README.md

@@ -27,6 +27,7 @@ bevy gaussian splatting render pipeline plugin
 - [ ] skeletons
 - [ ] volume masks
 - [ ] level of detail
+- [ ] lighting and shadows
 - [ ] gaussian cloud particle effects (accelerated spatial queries)
 - [ ] bevy_openxr support
 - [ ] bevy 3D camera to gaussian cloud pipeline
@@ -94,6 +95,7 @@ to build wasm run:
 - [making gaussian splats smaller](https://aras-p.info/blog/2023/09/13/Making-Gaussian-Splats-smaller/)
 - [masked-spacetime-hashing](https://github.com/masked-spacetime-hashing/msth)
 - [onesweep](https://arxiv.org/ftp/arxiv/papers/2206/2206.01784.pdf)
+- [pasture](https://github.com/Mortano/pasture)
 - [point-visualizer](https://github.com/mosure/point-visualizer)
 - [rusty-automata](https://github.com/mosure/rusty-automata)
 - [splat](https://github.com/antimatter15/splat)

benches/io.rs

@@ -0,0 +1,47 @@
+use criterion::{
+    BenchmarkId,
+    criterion_group,
+    criterion_main,
+    Criterion,
+    Throughput,
+};
+
+use bevy_gaussian_splatting::{
+    Gaussian,
+    GaussianCloud,
+    io::codec::GaussianCloudCodec,
+    random_gaussians,
+};
+
+
+const GAUSSIAN_COUNTS: [usize; 4] = [
+    1000,
+    10000,
+    84_348,
+    1_244_819,
+    // 6_131_954,
+];
+
+fn gaussian_cloud_decode_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("encode gaussian clouds");
+    for count in GAUSSIAN_COUNTS.iter() {
+        group.throughput(Throughput::Bytes(*count as u64 * std::mem::size_of::<Gaussian>() as u64));
+        group.bench_with_input(
+            BenchmarkId::new("decode", count),
+            &count,
+            |b, &count| {
+                let gaussians = random_gaussians(*count);
+                let bytes = gaussians.encode();
+
+                b.iter(|| GaussianCloud::decode(bytes.as_slice()));
+            },
+        );
+    }
+}
+
+criterion_group!{
+    name = io_benches;
+    config = Criterion::default().sample_size(10);
+    targets = gaussian_cloud_decode_benchmark
+}
+criterion_main!(io_benches);

src/gaussian.rs

@@ -1,40 +1,26 @@
-use rand::{seq::SliceRandom, prelude::Distribution, Rng};
-use std::{
-    io::{
-        BufReader,
-        Cursor,
-        ErrorKind,
-    },
-    marker::Copy,
+use rand::{
+    seq::SliceRandom,
+    prelude::Distribution,
+    Rng,
 };
+use std::marker::Copy;
 
 use bevy::{
     prelude::*,
-    asset::{
-        AssetLoader,
-        AsyncReadExt,
-        LoadContext,
-        io::Reader,
-    },
     reflect::TypeUuid,
     render::render_resource::ShaderType,
-    utils::BoxedFuture,
 };
-use bincode2::deserialize_from;
 use bytemuck::{
     Pod,
     Zeroable,
 };
-use flate2::read::GzDecoder;
 use serde::{
     Deserialize,
     Serialize,
     Serializer,
     ser::SerializeTuple,
 };
 
-use crate::ply::parse_ply;
-
 
 const fn num_sh_coefficients(degree: usize) -> usize {
     if degree == 0 {
@@ -51,6 +37,7 @@ pub const MAX_SH_COEFF_COUNT: usize = MAX_SH_COEFF_COUNT_PER_CHANNEL * SH_CHANNE
     Clone,
     Copy,
     Debug,
+    PartialEq,
     Reflect,
     ShaderType,
     Pod,
@@ -121,6 +108,7 @@ pub const MAX_SIZE_VARIANCE: f32 = 5.0;
     Debug,
     Default,
     Copy,
+    PartialEq,
     Reflect,
     ShaderType,
     Pod,
@@ -141,6 +129,7 @@ pub struct Gaussian {
     Asset,
     Clone,
     Debug,
+    PartialEq,
     Reflect,
     TypeUuid,
     Serialize,
@@ -231,53 +220,6 @@ impl Default for GaussianCloudSettings {
     }
 }
 
-
-#[derive(Default)]
-pub struct GaussianCloudLoader;
-
-impl AssetLoader for GaussianCloudLoader {
-    type Asset = GaussianCloud;
-    type Settings = ();
-    type Error = std::io::Error;
-
-    fn load<'a>(
-        &'a self,
-        reader: &'a mut Reader,
-        _settings: &'a Self::Settings,
-        load_context: &'a mut LoadContext,
-    ) -> BoxedFuture<'a, Result<Self::Asset, Self::Error>> {
-
-        Box::pin(async move {
-            let mut bytes = Vec::new();
-            reader.read_to_end(&mut bytes).await?;
-
-            match load_context.path().extension() {
-                Some(ext) if ext == "ply" => {
-                    let cursor = Cursor::new(bytes);
-                    let mut f = BufReader::new(cursor);
-
-                    let ply_cloud = parse_ply(&mut f)?;
-                    let cloud = GaussianCloud(ply_cloud);
-
-                    Ok(cloud)
-                },
-                Some(ext) if ext == "gcloud" => {
-                    let decompressed = GzDecoder::new(bytes.as_slice());
-                    let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud");
-
-                    Ok(cloud)
-                },
-                _ => Err(std::io::Error::new(ErrorKind::Other, "only .ply and .gcloud supported")),
-            }
-        })
-    }
-
-    fn extensions(&self) -> &[&str] {
-        &["ply", "gcloud"]
-    }
-}
-
-
 impl Distribution<Gaussian> for rand::distributions::Standard {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Gaussian {
         Gaussian {

src/io/codec.rs

@@ -0,0 +1,6 @@
+
+// TODO: support streamed codecs
+pub trait GaussianCloudCodec {
+    fn encode(&self) -> Vec<u8>;
+    fn decode(data: &[u8]) -> Self;
+}

src/io/gcloud/bincode2.rs

@@ -0,0 +1,35 @@
+use bincode2::{
+    deserialize_from,
+    serialize_into,
+};
+use flate2::{
+    Compression,
+    read::GzDecoder,
+    write::GzEncoder,
+};
+
+use crate::{
+    gaussian::GaussianCloud,
+    io::codec::GaussianCloudCodec,
+};
+
+
+impl GaussianCloudCodec for GaussianCloud {
+    fn encode(&self) -> Vec<u8> {
+        let mut output = Vec::new();
+
+        {
+            let mut gz_encoder = GzEncoder::new(&mut output, Compression::default());
+            serialize_into(&mut gz_encoder, &self).expect("failed to encode cloud");
+        }
+
+        output
+    }
+
+    fn decode(data: &[u8]) -> Self {
+        let decompressed = GzDecoder::new(data);
+        let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud");
+
+        cloud
+    }
+}

src/io/gcloud/flexbuffers.rs

@@ -0,0 +1,30 @@
+use flexbuffers::{
+    FlexbufferSerializer,
+    Reader,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
+
+use crate::{
+    GaussianCloud,
+    io::codec::GaussianCloudCodec,
+};
+
+
+impl GaussianCloudCodec for GaussianCloud {
+    fn encode(&self) -> Vec<u8> {
+        let mut serializer = FlexbufferSerializer::new();
+        self.serialize(&mut serializer).expect("failed to serialize cloud");
+
+        serializer.view().to_vec()
+    }
+
+    fn decode(data: &[u8]) -> Self {
+        let reader = Reader::get_root(data).expect("failed to read flexbuffer");
+        let cloud = GaussianCloud::deserialize(reader).expect("deserialization failed");
+
+        cloud
+    }
+}

src/io/gcloud/mod.rs

@@ -0,0 +1,5 @@
+#[cfg(feature = "io_bincode2")]
+pub mod bincode2;
+
+#[cfg(feature = "io_flexbuffers")]
+pub mod flexbuffers;