support instancing and multiple cameras

[github-actions]

separate entry_buffer_a binding into unique a bind group to optimize buffer updates

bevy_gaussian_splatting/src/sort/mod.rs

Line 214 in 5c9a20a

// TODO: support instancing and multiple cameras

use bevy::{
    prelude::*,
    asset::LoadState,
    ecs::system::{
        lifetimeless::SRes,
        SystemParamItem,
    },
    reflect::TypeUuid,
    render::{
        render_asset::{
            RenderAsset,
            RenderAssetPlugin,
            PrepareAssetError,
        },
        render_resource::*,
        renderer::RenderDevice,
    },
};
use bytemuck::{
    Pod,
    Zeroable,
};
use static_assertions::assert_cfg;

use crate::{
    GaussianCloud,
    GaussianCloudSettings,
};


#[cfg(feature = "sort_radix")]
pub mod radix;

#[cfg(feature = "sort_rayon")]
pub mod rayon;


assert_cfg!(
    any(
        feature = "sort_radix",
        feature = "sort_rayon",
    ),
    "no sort mode enabled",
);


#[derive(
    Component,
    Debug,
    Clone,
    PartialEq,
    Reflect,
)]
pub enum SortMode {
    None,

    #[cfg(feature = "sort_radix")]
    Radix,

    #[cfg(feature = "sort_rayon")]
    Rayon,
}

impl Default for SortMode {
    #[allow(unreachable_code)]
    fn default() -> Self {
        #[cfg(feature = "sort_rayon")]
        return Self::Rayon;

        #[cfg(feature = "sort_radix")]
        return Self::Radix;

        Self::None
    }
}


#[derive(Default)]
pub struct SortPlugin;

impl Plugin for SortPlugin {
    fn build(&self, app: &mut App) {
        #[cfg(feature = "sort_radix")]
        app.add_plugins(radix::RadixSortPlugin);

        #[cfg(feature = "sort_rayon")]
        app.add_plugins(rayon::RayonSortPlugin);


        app.register_type::<SortedEntries>();
        app.init_asset::<SortedEntries>();
        app.register_asset_reflect::<SortedEntries>();

        app.add_plugins(RenderAssetPlugin::<SortedEntries>::default());

        app.add_systems(Update, auto_insert_sorted_entries);
    }
}


fn auto_insert_sorted_entries(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    gaussian_clouds_res: Res<Assets<GaussianCloud>>,
    mut sorted_entries_res: ResMut<Assets<SortedEntries>>,
    gaussian_clouds: Query<(
        Entity,
        &Handle<GaussianCloud>,
        &GaussianCloudSettings,
        Without<Handle<SortedEntries>>,
    )>,
) {
    for (
        entity,
        gaussian_cloud_handle,
        _settings,
        _,
    ) in gaussian_clouds.iter() {
        // // TODO: specialize vertex shader for sort mode (e.g. draw_indirect but no sort indirection)
        // if settings.sort_mode == SortMode::None {
        //     continue;
        // }

        if Some(LoadState::Loading) == asset_server.get_load_state(gaussian_cloud_handle) {
            continue;
        }

        let cloud = gaussian_clouds_res.get(gaussian_cloud_handle);
        if cloud.is_none() {
            continue;
        }
        let cloud = cloud.unwrap();

        // TODO: move gaussian_cloud and sorted_entry assets into an asset bundle
        let sorted_entries = sorted_entries_res.add(SortedEntries {
            sorted: (0..cloud.gaussians.len())
                .map(|idx| {
                    SortEntry {
                        key: 1,
                        index: idx as u32,
                    }
                })
                .collect(),
        });

        commands.entity(entity)
            .insert(sorted_entries);
    }
}


#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    PartialEq,
    Reflect,
    ShaderType,
    Pod,
    Zeroable,
)]
#[repr(C)]
pub struct SortEntry {
    pub key: u32,
    pub index: u32,
}

// TODO: add RenderAssetPlugin for SortedEntries & auto-insert to GaussianCloudBundles if their sort mode is not None
// supports pre-sorting or CPU sorting in main world, initializes the sorting_entry_buffer
#[derive(
    Clone,
    Asset,
    Debug,
    Default,
    PartialEq,
    Reflect,
    TypeUuid,
)]
#[uuid = "ac2f08eb-fa13-ccdd-ea11-51571ea332d5"]
pub struct SortedEntries {
    pub sorted: Vec<SortEntry>,
}

impl RenderAsset for SortedEntries {
    type ExtractedAsset = SortedEntries;
    type PreparedAsset = GpuSortedEntry;
    type Param = SRes<RenderDevice>;

    fn extract_asset(&self) -> Self::ExtractedAsset {
        self.clone()
    }

    fn prepare_asset(
        sorted_entries: Self::ExtractedAsset,
        render_device: &mut SystemParamItem<Self::Param>,
    ) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>> {
        let sorted_entry_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
            label: Some("sorted_entry_buffer"),
            contents: bytemuck::cast_slice(sorted_entries.sorted.as_slice()),
            usage: BufferUsages::COPY_SRC | BufferUsages::COPY_DST | BufferUsages::STORAGE,
        });

        let count = sorted_entries.sorted.len();

        Ok(GpuSortedEntry {
            sorted_entry_buffer,
            count,
        })
    }
}


// TODO: support instancing and multiple cameras
//       separate entry_buffer_a binding into unique a bind group to optimize buffer updates
#[derive(Debug, Clone)]
pub struct GpuSortedEntry {
    pub sorted_entry_buffer: Buffer,
    pub count: usize,
}

[MiniMinerX]

@mosure How could I go about working on this issue? Is this related to issue #12. I see your note about entry_buffer_a but don't know what that means. I'd love to help!

[mosure]

@MiniMinerX I appreciate the help!

the gaussian splat forward rasterization requires per-view (camera) depth sort.

in the current implementation, the sort buffers are instantiated once per gaussian cloud entity. this is a problem for the current sort strategies (radix sort is a render-graph level compute node, CPU sort is done in the main world), since there is only one buffer, fix-sized to 1x number of gaussians, see CPU sort where multiple cameras write to the same memory:

bevy_gaussian_splatting/src/sort/std.rs

Line 61 in 08f6e4a

for (

there are two strategies:

create a new sorted entry buffer for each camera

• the entry buffer (and sorted_bind_group) would need to be created per View:
  

  bevy_gaussian_splatting/src/render/mod.rs

  Line 812 in 08f6e4a

  let sorted_bind_group = render_device.create_bind_group(

• edit I am not sure this works because the ECS structure becomes a bit messy, e.g. attaching sort buffers to cameras doesn't work well with multiple gaussian entities. attaching multiple sort buffers to a single gaussian entity is also a bit odd. I think the dynamic index solution makes the most sense.

use a single storage buffer with dynamic index

• the single sorted_entry_buffer would need to be sized based on the number of cameras in the scene:
  

  bevy_gaussian_splatting/src/render/mod.rs

  Line 821 in 08f6e4a

  size: BufferSize::new((cloud.count * std::mem::size_of::<(u32, u32)>()) as u64),

• the sorted_bind_group would need to be set with dynamic index, example of setting dynamic index here:
  

  bevy_gaussian_splatting/src/render/mod.rs

  Line 987 in 08f6e4a

  pass.set_bind_group(3, &bind_groups.sorted_bind_group, &[]);

in either case, writing to SortedEntries (on CPU) or (on GPU) would need to be aware of the buffer layout/indexing strategy.

---

the error in #111 would also need to be addressed. I think it refers to a bug in queue_gaussian_view_bind_groups, for some reason, the view_binding in Res<ViewUniforms> is not correlated correctly with the ViewUniformOffset. It might be good to log which entity has the invalid offset for further debugging.