SelectionBundle

Struct SelectionBundle 

Source
pub struct SelectionBundle<G: GaussianPod> {
    pub bundles: Vec<ComputeBundle<()>>,
    /* private fields */
}
Expand description

A collection of specialized ComputeBundle for selection operations.

§Custom Operations

All ComputeBundles supplied to this bundle as a SelectionExpr::Unary, SelectionExpr::Binary, or SelectionExpr::Selection custom operation must have the same bind group 0 as the SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR. They must also not have the bind group itself, as it will be supplied automatically during evaluation.

It is recommended to use ComputeBundleBuilder to create the custom operation bundles, and build them using ComputeBundleBuilder::build_without_bind_groups.

// Create the selection custom operation compute bundle
let my_selection_custom_op_bundle = ComputeBundleBuilder::new()
    .label("My Selection")
    .bind_group_layouts([
        &SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
        &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR, // Put your custom bind group layout here.
    ])
    .resolver({
        let mut resolver = wesl::PkgResolver::new();
        resolver.add_package(&core::shader::PACKAGE); // Required for using core buffer structs.
        resolver.add_package(&shader::PACKAGE); // Optionally add this for some utility functions.
        resolver
    })
    .main_shader("path::to::my::wesl::module".parse().unwrap())
    .entry_point("main")
    .wesl_compile_options(wesl::CompileOptions {
        features: G::wesl_features(), // Required for enabling the correct features for core structs.
        ..Default::default()
    })
    .build_without_bind_groups(&device)
    .map_err(|e| log::error!("{e}"))
    .expect("my selection custom op bundle");

// Create the selection bundle
let selection_bundle = SelectionBundle::<GaussianPod>::new(
    &device,
    vec![my_selection_custom_op_bundle],
);

// Create the bind group for your custom operation
let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
    .create_bind_group(
        &device,
        1, // Index 0 is the Gaussians buffer, so remember to start from 1 for your bind groups
        [my_selection_custom_op_buffer.buffer().as_entire_binding()], // Your custom bind group resources
    )
    .unwrap()

// Create the selection expression
let selection_expr = gs::SelectionExpr::selection(
    0, // The bundle index for your custom operation in the selection bundle
    vec![my_selection_custom_op_bind_group],
)
.union( // Combine with other selection expressions using different functions
    gs::SelectionExpr::Buffer(my_existing_selection_buffer) // An existing selection buffer for example
);

// Create a selection buffer for the result
let dest_selection_buffer = SelectionBuffer::new(&device, gaussians_buffer.len() as u32);

// Evaluate the selection expression
selection_bundle.evaluate(
    &device,
    &mut encoder,
    &selection_expr,
    &dest_selection_buffer,
    &model_transform,
    &gaussian_transform,
    &gaussians_buffer,
);

§Shader Format

You may copy and paste the following shader bindings for SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR into your custom selection operation shader to ensure that the bindings are correct, then add your own bindings after that.

import wgpu_3dgs_core::{
    gaussian::Gaussian,
    gaussian_transform::GaussianTransform,
    model_transform::{model_to_world, ModelTransform},
};

@group(0) @binding(0)
var<uniform> op: u32;

@group(0) @binding(1)
var<storage, read> source: array<u32>;

@group(0) @binding(2)
var<storage, read_write> dest: array<atomic<u32>>;

@group(0) @binding(3)
var<uniform> model_transform: ModelTransform;

@group(0) @binding(4)
var<uniform> gaussian_transform: GaussianTransform;

@group(0) @binding(5)
var<storage, read> gaussians: array<Gaussian>;

// Your custom bindings here...

override workgroup_size: u32;

@compute @workgroup_size(workgroup_size)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let index = id.x;

    if index >= arrayLength(&gaussians) {
        return;
    }

    let gaussian = gaussians[index];

    let world_pos = model_to_world(model_transform, gaussian.position);

    // Your custom selection operation code here...

    let word_index = index / 32u;
    let bit_index = index % 32u;
    let bit_mask = 1u << bit_index;
    if /* Condition for selecting the Gaussian */ {
        atomicOr(&dest[word_index], bit_mask);
    } else {
        atomicAnd(&dest[word_index], ~bit_mask);
    }
}

Fields§

§bundles: Vec<ComputeBundle<()>>

The compute bundles for selection custom operations.

Implementations§

Source§

impl<G: GaussianPod> SelectionBundle<G>

Source

pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>

The Gaussians bind group layout descriptors.

This bind group layout takes the following buffers:

Source

pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>

The sphere selection bind group layout descriptor.

This bind group layout takes the following buffers:

Source

pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: BindGroupLayoutDescriptor<'static>

The box selection bind group layout descriptor.

This bind group layout takes the following buffers:

Source

pub fn new<'a>(device: &Device, bundles: Vec<ComputeBundle<()>>) -> Self

Create a new selection bundle.

bundles are used for SelectionExpr::Unary, SelectionExpr::Binary, or SelectionExpr::Selection as custom operations, they must have the same bind group 0 as the SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR, see documentation of SelectionBundle for more details.

Examples found in repository?
examples/filter_selection.rs (line 143)
88async fn main() {
89    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
90
91    let args = Args::parse();
92    let model_path = &args.model;
93    let pos = Vec3::from_slice(&args.pos);
94    let rot = Quat::from_slice(&args.rot);
95    let scale = Vec3::from_slice(&args.scale);
96    let shape = match args.shape {
97        Shape::Sphere => gs::SelectionBundle::<GaussianPod>::create_sphere_bundle,
98        Shape::Box => gs::SelectionBundle::<GaussianPod>::create_box_bundle,
99    };
100    let repeat = args.repeat;
101    let offset = Vec3::from_slice(&args.offset);
102
103    log::debug!("Creating wgpu instance");
104    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
105
106    log::debug!("Requesting adapter");
107    let adapter = instance
108        .request_adapter(&wgpu::RequestAdapterOptions::default())
109        .await
110        .expect("adapter");
111
112    log::debug!("Requesting device");
113    let (device, queue) = adapter
114        .request_device(&wgpu::DeviceDescriptor {
115            label: Some("Device"),
116            required_features: wgpu::Features::empty(),
117            required_limits: adapter.limits(),
118            memory_hints: wgpu::MemoryHints::default(),
119            trace: wgpu::Trace::Off,
120        })
121        .await
122        .expect("device");
123
124    log::debug!("Creating gaussians");
125    let f = std::fs::File::open(model_path).expect("ply file");
126    let mut reader = std::io::BufReader::new(f);
127    let gaussians = gs::core::Gaussians::read_ply(&mut reader).expect("gaussians");
128
129    log::debug!("Creating gaussians buffer");
130    let gaussians_buffer =
131        gs::core::GaussiansBuffer::<GaussianPod>::new(&device, &gaussians.gaussians);
132
133    log::debug!("Creating model transform buffer");
134    let model_transform = gs::core::ModelTransformBuffer::new(&device);
135
136    log::debug!("Creating Gaussian transform buffer");
137    let gaussian_transform = gs::core::GaussianTransformBuffer::new(&device);
138
139    log::debug!("Creating shape selection compute bundle");
140    let shape_selection = shape(&device);
141
142    log::debug!("Creating selection bundle");
143    let selection_bundle = gs::SelectionBundle::<GaussianPod>::new(&device, vec![shape_selection]);
144
145    log::debug!("Creating shape selection buffers");
146    let shape_selection_buffers = (0..repeat)
147        .map(|i| {
148            let offset_pos = pos + offset * i as f32;
149            let buffer = gs::InvTransformBuffer::new(&device);
150            buffer.update_with_scale_rot_pos(&queue, scale, rot, offset_pos);
151            buffer
152        })
153        .collect::<Vec<_>>();
154
155    log::debug!("Creating shape selection bind groups");
156    let shape_selection_bind_groups = shape_selection_buffers
157        .iter()
158        .map(|buffer| {
159            selection_bundle.bundles[0]
160                .create_bind_group(
161                    &device,
162                    // index 0 is the Gaussians buffer, so we use 1,
163                    // see docs of create_sphere_bundle or create_box_bundle
164                    1,
165                    [buffer.buffer().as_entire_binding()],
166                )
167                .expect("bind group")
168        })
169        .collect::<Vec<_>>();
170
171    log::debug!("Creating selection expression");
172    let selection_expr = shape_selection_bind_groups.into_iter().fold(
173        gs::SelectionExpr::Identity,
174        |acc, bind_group| {
175            acc.union(gs::SelectionExpr::selection(
176                0, // the 0 here is the bundle index in the selection bundle
177                vec![bind_group],
178            ))
179        },
180    );
181
182    log::debug!("Creating destination buffer");
183    let dest = gs::SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
184
185    log::info!("Starting selection process");
186    let time = std::time::Instant::now();
187
188    log::debug!("Selecting Gaussians");
189    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
190        label: Some("Selection Encoder"),
191    });
192
193    selection_bundle.evaluate(
194        &device,
195        &mut encoder,
196        &selection_expr,
197        &dest,
198        &model_transform,
199        &gaussian_transform,
200        &gaussians_buffer,
201    );
202
203    queue.submit(Some(encoder.finish()));
204
205    #[allow(unused_must_use)]
206    device.poll(wgpu::PollType::Wait);
207
208    log::info!("Editing process completed in {:?}", time.elapsed());
209
210    log::debug!("Filtering Gaussians");
211    let selected_gaussians = gs::core::Gaussians {
212        gaussians: dest
213            .download::<u32>(&device, &queue)
214            .await
215            .expect("selected download")
216            .iter()
217            .flat_map(|group| {
218                std::iter::repeat_n(group, 32)
219                    .enumerate()
220                    .map(|(i, g)| g & (1 << i) != 0)
221            })
222            .zip(gaussians.gaussians.iter())
223            .filter(|(selected, _)| *selected)
224            .map(|(_, g)| g.clone())
225            .collect::<Vec<_>>(),
226    };
227
228    log::debug!("Writing modified Gaussians to output file");
229    let output_file = std::fs::File::create(&args.output).expect("output file");
230    let mut writer = std::io::BufWriter::new(output_file);
231    selected_gaussians
232        .write_ply(&mut writer)
233        .expect("write modified Gaussians to output file");
234
235    log::info!("Filtered Gaussians written to {}", args.output);
236}
Source

pub fn gaussians_bind_group_layout(&self) -> &BindGroupLayout

Get the Gaussians bind group layout.

Source

pub fn evaluate( &self, device: &Device, encoder: &mut CommandEncoder, expr: &SelectionExpr, dest: &SelectionBuffer, model_transform: &ModelTransformBuffer, gaussian_transform: &GaussianTransformBuffer, gaussians: &GaussiansBuffer<G>, )

Evaluate and apply the selection expression.

Examples found in repository?
examples/filter_selection.rs (lines 193-201)
88async fn main() {
89    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
90
91    let args = Args::parse();
92    let model_path = &args.model;
93    let pos = Vec3::from_slice(&args.pos);
94    let rot = Quat::from_slice(&args.rot);
95    let scale = Vec3::from_slice(&args.scale);
96    let shape = match args.shape {
97        Shape::Sphere => gs::SelectionBundle::<GaussianPod>::create_sphere_bundle,
98        Shape::Box => gs::SelectionBundle::<GaussianPod>::create_box_bundle,
99    };
100    let repeat = args.repeat;
101    let offset = Vec3::from_slice(&args.offset);
102
103    log::debug!("Creating wgpu instance");
104    let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
105
106    log::debug!("Requesting adapter");
107    let adapter = instance
108        .request_adapter(&wgpu::RequestAdapterOptions::default())
109        .await
110        .expect("adapter");
111
112    log::debug!("Requesting device");
113    let (device, queue) = adapter
114        .request_device(&wgpu::DeviceDescriptor {
115            label: Some("Device"),
116            required_features: wgpu::Features::empty(),
117            required_limits: adapter.limits(),
118            memory_hints: wgpu::MemoryHints::default(),
119            trace: wgpu::Trace::Off,
120        })
121        .await
122        .expect("device");
123
124    log::debug!("Creating gaussians");
125    let f = std::fs::File::open(model_path).expect("ply file");
126    let mut reader = std::io::BufReader::new(f);
127    let gaussians = gs::core::Gaussians::read_ply(&mut reader).expect("gaussians");
128
129    log::debug!("Creating gaussians buffer");
130    let gaussians_buffer =
131        gs::core::GaussiansBuffer::<GaussianPod>::new(&device, &gaussians.gaussians);
132
133    log::debug!("Creating model transform buffer");
134    let model_transform = gs::core::ModelTransformBuffer::new(&device);
135
136    log::debug!("Creating Gaussian transform buffer");
137    let gaussian_transform = gs::core::GaussianTransformBuffer::new(&device);
138
139    log::debug!("Creating shape selection compute bundle");
140    let shape_selection = shape(&device);
141
142    log::debug!("Creating selection bundle");
143    let selection_bundle = gs::SelectionBundle::<GaussianPod>::new(&device, vec![shape_selection]);
144
145    log::debug!("Creating shape selection buffers");
146    let shape_selection_buffers = (0..repeat)
147        .map(|i| {
148            let offset_pos = pos + offset * i as f32;
149            let buffer = gs::InvTransformBuffer::new(&device);
150            buffer.update_with_scale_rot_pos(&queue, scale, rot, offset_pos);
151            buffer
152        })
153        .collect::<Vec<_>>();
154
155    log::debug!("Creating shape selection bind groups");
156    let shape_selection_bind_groups = shape_selection_buffers
157        .iter()
158        .map(|buffer| {
159            selection_bundle.bundles[0]
160                .create_bind_group(
161                    &device,
162                    // index 0 is the Gaussians buffer, so we use 1,
163                    // see docs of create_sphere_bundle or create_box_bundle
164                    1,
165                    [buffer.buffer().as_entire_binding()],
166                )
167                .expect("bind group")
168        })
169        .collect::<Vec<_>>();
170
171    log::debug!("Creating selection expression");
172    let selection_expr = shape_selection_bind_groups.into_iter().fold(
173        gs::SelectionExpr::Identity,
174        |acc, bind_group| {
175            acc.union(gs::SelectionExpr::selection(
176                0, // the 0 here is the bundle index in the selection bundle
177                vec![bind_group],
178            ))
179        },
180    );
181
182    log::debug!("Creating destination buffer");
183    let dest = gs::SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
184
185    log::info!("Starting selection process");
186    let time = std::time::Instant::now();
187
188    log::debug!("Selecting Gaussians");
189    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
190        label: Some("Selection Encoder"),
191    });
192
193    selection_bundle.evaluate(
194        &device,
195        &mut encoder,
196        &selection_expr,
197        &dest,
198        &model_transform,
199        &gaussian_transform,
200        &gaussians_buffer,
201    );
202
203    queue.submit(Some(encoder.finish()));
204
205    #[allow(unused_must_use)]
206    device.poll(wgpu::PollType::Wait);
207
208    log::info!("Editing process completed in {:?}", time.elapsed());
209
210    log::debug!("Filtering Gaussians");
211    let selected_gaussians = gs::core::Gaussians {
212        gaussians: dest
213            .download::<u32>(&device, &queue)
214            .await
215            .expect("selected download")
216            .iter()
217            .flat_map(|group| {
218                std::iter::repeat_n(group, 32)
219                    .enumerate()
220                    .map(|(i, g)| g & (1 << i) != 0)
221            })
222            .zip(gaussians.gaussians.iter())
223            .filter(|(selected, _)| *selected)
224            .map(|(_, g)| g.clone())
225            .collect::<Vec<_>>(),
226    };
227
228    log::debug!("Writing modified Gaussians to output file");
229    let output_file = std::fs::File::create(&args.output).expect("output file");
230    let mut writer = std::io::BufWriter::new(output_file);
231    selected_gaussians
232        .write_ply(&mut writer)
233        .expect("write modified Gaussians to output file");
234
235    log::info!("Filtered Gaussians written to {}", args.output);
236}
Source

pub fn create_primitive_bundle(device: &Device) -> ComputeBundle<()>

Create the selection primitive operation ComputeBundle.

You usually do not need to use this method, it is used internally for creating the primitive operation bundle for evaluation.

Source

pub fn create_sphere_bundle(device: &Device) -> ComputeBundle<()>

Create a sphere selection custom operation.

Source

pub fn create_box_bundle(device: &Device) -> ComputeBundle<()>

Create a box selection custom operation.

Trait Implementations§

Source§

impl<G: Debug + GaussianPod> Debug for SelectionBundle<G>

Source§

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more

Auto Trait Implementations§

§

impl<G> Freeze for SelectionBundle<G>

§

impl<G> !RefUnwindSafe for SelectionBundle<G>

§

impl<G> Send for SelectionBundle<G>

§

impl<G> Sync for SelectionBundle<G>

§

impl<G> Unpin for SelectionBundle<G>
where G: Unpin,

§

impl<G> !UnwindSafe for SelectionBundle<G>

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> Downcast<T> for T

Source§

fn downcast(&self) -> &T

Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> IntoEither for T

Source§

fn into_either(self, into_left: bool) -> Either<Self, Self>

Converts self into a Left variant of Either<Self, Self> if into_left is true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
where F: FnOnce(&Self) -> bool,

Converts self into a Left variant of Either<Self, Self> if into_left(&self) returns true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> Upcast<T> for T

Source§

fn upcast(&self) -> Option<&T>

Source§

impl<T> WasmNotSend for T
where T: Send,

Source§

impl<T> WasmNotSendSync for T

Source§

impl<T> WasmNotSync for T
where T: Sync,