wgpu_3dgs_editor/selection.rs
1use glam::*;
2
3use crate::{
4 SelectionBuffer, SelectionOpBuffer,
5 core::{
6 self, BufferWrapper, ComputeBundle, ComputeBundleBuilder, GaussianPod,
7 GaussianTransformBuffer, GaussiansBuffer, ModelTransformBuffer,
8 },
9 shader,
10};
11
12/// A selection expression tree.
13///
14/// ## Overview
15///
16/// This can be used to carry out operations on selection buffers, these operations are evaluated
17/// by [`SelectionBundle::evaluate`] in a recursive manner (depth-first).
18///
19/// ## Custom Operations
20///
21/// [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], and [`SelectionExpr::Selection`] are
22/// custom operations that can be defined with additional [`ComputeBundle`]s, so they also
23/// carry a vector of bind groups that are used in the operation when dispatched/evaluated.
24/// These vectors should correspond to the selection bundle's bind groups starting at index 1,
25/// because index 0 must be defined by [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
26#[derive(Debug, Default)]
27pub enum SelectionExpr {
28 /// Apply an identity operation.
29 #[default]
30 Identity,
31 /// Union of the two selections.
32 Union(Box<SelectionExpr>, Box<SelectionExpr>),
33 /// Interaction of the two selections.
34 Intersection(Box<SelectionExpr>, Box<SelectionExpr>),
35 /// Difference of the two selections.
36 Difference(Box<SelectionExpr>, Box<SelectionExpr>),
37 /// Symmetric difference of the two selections.
38 SymmetricDifference(Box<SelectionExpr>, Box<SelectionExpr>),
39 /// Complement of the selection.
40 Complement(Box<SelectionExpr>),
41 /// Apply a custom unary operation.
42 Unary(usize, Box<SelectionExpr>, Vec<wgpu::BindGroup>),
43 /// Apply a custom binary operation.
44 Binary(
45 Box<SelectionExpr>,
46 usize,
47 Box<SelectionExpr>,
48 Vec<wgpu::BindGroup>,
49 ),
50 /// Create a selection.
51 Selection(usize, Vec<wgpu::BindGroup>),
52 /// Directly use a selection buffer.
53 Buffer(SelectionBuffer),
54}
55
56impl SelectionExpr {
57 /// The first u32 value for a custom operation.
58 pub const CUSTOM_OP_START: u32 = 5;
59
60 /// Create a new [`SelectionExpr::Identity`].
61 pub fn identity() -> Self {
62 Self::Identity
63 }
64
65 /// Create a new [`SelectionExpr::Union`].
66 pub fn union(self, other: Self) -> Self {
67 Self::Union(Box::new(self), Box::new(other))
68 }
69
70 /// Create a new [`SelectionExpr::Intersection`].
71 pub fn intersection(self, other: Self) -> Self {
72 Self::Intersection(Box::new(self), Box::new(other))
73 }
74
75 /// Create a new [`SelectionExpr::Difference`].
76 pub fn difference(self, other: Self) -> Self {
77 Self::Difference(Box::new(self), Box::new(other))
78 }
79
80 /// Create a new [`SelectionExpr::SymmetricDifference`].
81 pub fn symmetric_difference(self, other: Self) -> Self {
82 Self::SymmetricDifference(Box::new(self), Box::new(other))
83 }
84
85 /// Create a new [`SelectionExpr::Complement`].
86 pub fn complement(self) -> Self {
87 Self::Complement(Box::new(self))
88 }
89
90 /// Create a new [`SelectionExpr::Unary`].
91 pub fn unary(self, op: usize, bind_groups: Vec<wgpu::BindGroup>) -> Self {
92 Self::Unary(op, Box::new(self), bind_groups)
93 }
94
95 /// Create a new [`SelectionExpr::Binary`].
96 pub fn binary(self, op: usize, other: Self, bind_groups: Vec<wgpu::BindGroup>) -> Self {
97 Self::Binary(Box::new(self), op, Box::new(other), bind_groups)
98 }
99
100 /// Create a new [`SelectionExpr::Selection`].
101 pub fn selection(op: usize, bind_groups: Vec<wgpu::BindGroup>) -> Self {
102 Self::Selection(op, bind_groups)
103 }
104
105 /// Create a new [`SelectionExpr::Buffer`].
106 pub fn buffer(buffer: SelectionBuffer) -> Self {
107 Self::Buffer(buffer)
108 }
109
110 /// Update the expression in place.
111 pub fn update_with(&mut self, f: impl FnOnce(Self) -> Self) {
112 *self = f(std::mem::take(self));
113 }
114
115 /// Get the u32 associated with this expression's operation.
116 ///
117 /// The value returned is not the same as that returned by [`SelectionExpr::custom_op_index`],
118 /// but rather a value that can be used to identify the operation by the compute shader, custom
119 /// operation's index are offset by [`SelectionExpr::CUSTOM_OP_START`].
120 ///
121 /// You usually do not need to use this method, it is used internally for evaluation of the
122 /// compute shader.
123 pub fn as_u32(&self) -> Option<u32> {
124 match self {
125 SelectionExpr::Union(_, _) => Some(0),
126 SelectionExpr::Intersection(_, _) => Some(1),
127 SelectionExpr::SymmetricDifference(_, _) => Some(2),
128 SelectionExpr::Difference(_, _) => Some(3),
129 SelectionExpr::Complement(_) => Some(4),
130 SelectionExpr::Unary(op, _, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
131 SelectionExpr::Binary(_, op, _, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
132 SelectionExpr::Selection(op, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
133 SelectionExpr::Buffer(_) => None,
134 SelectionExpr::Identity => None,
135 }
136 }
137
138 /// Whether this expression is an identity operation.
139 pub fn is_identity(&self) -> bool {
140 matches!(self, SelectionExpr::Identity)
141 }
142
143 /// Whether this expression is a primitive operation.
144 pub fn is_primitive(&self) -> bool {
145 matches!(
146 self,
147 SelectionExpr::Union(..)
148 | SelectionExpr::Intersection(..)
149 | SelectionExpr::Difference(..)
150 | SelectionExpr::SymmetricDifference(..)
151 | SelectionExpr::Complement(..)
152 )
153 }
154
155 /// Whether this expression is a custom operation.
156 pub fn is_custom(&self) -> bool {
157 matches!(
158 self,
159 SelectionExpr::Unary(..) | SelectionExpr::Binary(..) | SelectionExpr::Selection(..)
160 )
161 }
162
163 /// Whether this expression is a selection operation.
164 pub fn is_operation(&self) -> bool {
165 matches!(
166 self,
167 SelectionExpr::Union(..)
168 | SelectionExpr::Intersection(..)
169 | SelectionExpr::Difference(..)
170 | SelectionExpr::SymmetricDifference(..)
171 | SelectionExpr::Complement(..)
172 | SelectionExpr::Unary(..)
173 | SelectionExpr::Binary(..)
174 | SelectionExpr::Selection(..)
175 )
176 }
177
178 /// Whether this expression is a selection buffer.
179 pub fn is_buffer(&self) -> bool {
180 matches!(self, SelectionExpr::Buffer(_))
181 }
182
183 /// Get the custom operation index.
184 ///
185 /// This is the index of the custom operation in [`SelectionBundle::bundles`] vector.
186 pub fn custom_op_index(&self) -> Option<usize> {
187 match self {
188 SelectionExpr::Unary(op, _, _)
189 | SelectionExpr::Binary(_, op, _, _)
190 | SelectionExpr::Selection(op, _) => Some(*op),
191 _ => None,
192 }
193 }
194
195 /// Get the custom operation bind groups for this expression.
196 pub fn custom_bind_groups(&self) -> Option<&Vec<wgpu::BindGroup>> {
197 match self {
198 SelectionExpr::Unary(_, _, bind_groups) => Some(bind_groups),
199 SelectionExpr::Binary(_, _, _, bind_groups) => Some(bind_groups),
200 SelectionExpr::Selection(_, bind_groups) => Some(bind_groups),
201 _ => None,
202 }
203 }
204
205 /// Get the custom operation index and bind groups for this expression.
206 pub fn custom_op_index_and_bind_groups(&self) -> Option<(usize, &Vec<wgpu::BindGroup>)> {
207 match self {
208 SelectionExpr::Unary(op, _, bind_groups)
209 | SelectionExpr::Binary(_, op, _, bind_groups)
210 | SelectionExpr::Selection(op, bind_groups) => Some((*op, bind_groups)),
211 _ => None,
212 }
213 }
214}
215
216/// A collection of specialized [`ComputeBundle`] for selection operations.
217///
218/// ## Custom Operations
219///
220/// All [`ComputeBundle`]s supplied to this bundle as a [`SelectionExpr::Unary`],
221/// [`SelectionExpr::Binary`], or [`SelectionExpr::Selection`] custom operation must have the same
222/// bind group 0 as the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`]. They must also
223/// not have the bind group itself, as it will be supplied automatically during evaluation.
224///
225/// Note that [`SelectionExpr::Unary`] will also get the source selection buffer, but it will be
226/// empty (all zeros), you should operate on the destination selection buffer only.
227///
228/// It is recommended to use [`ComputeBundleBuilder`] to create the custom operation bundles,
229/// and build them using [`ComputeBundleBuilder::build_without_bind_groups`].
230///
231/// ```rust no_run
232/// # pollster::block_on(async {
233/// # use wgpu_3dgs_editor::{
234/// # Editor, MODIFIER_GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR, Modifier, SelectionBuffer,
235/// # SelectionBundle, SelectionExpr,
236/// # core::{
237/// # self, BufferWrapper, GaussianPod as _, GaussianTransformBuffer,
238/// # GaussiansBuffer, ModelTransformBuffer, glam::*,
239/// # },
240/// # shader,
241/// # };
242/// #
243/// # type GaussianPod = core::GaussianPodWithShSingleCov3dSingleConfigs;
244/// #
245/// # let instance = wgpu::Instance::new(
246/// # wgpu::InstanceDescriptor::new_without_display_handle_from_env()
247/// # );
248/// #
249/// # let adapter = instance
250/// # .request_adapter(&wgpu::RequestAdapterOptions::default())
251/// # .await
252/// # .expect("adapter");
253/// #
254/// # let (device, _queue) = adapter
255/// # .request_device(&wgpu::DeviceDescriptor {
256/// # label: Some("Device"),
257/// # required_limits: adapter.limits(),
258/// # ..Default::default()
259/// # })
260/// # .await
261/// # .expect("device");
262/// #
263/// # const MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor =
264/// # wgpu::BindGroupLayoutDescriptor {
265/// # label: Some("My Custom Bind Group Layout"),
266/// # entries: &[wgpu::BindGroupLayoutEntry {
267/// # binding: 0,
268/// # visibility: wgpu::ShaderStages::COMPUTE,
269/// # ty: wgpu::BindingType::Buffer {
270/// # ty: wgpu::BufferBindingType::Uniform,
271/// # has_dynamic_offset: false,
272/// # min_binding_size: None,
273/// # },
274/// # count: None,
275/// # }],
276/// # };
277/// #
278/// # let my_buffer = device.create_buffer(&wgpu::BufferDescriptor {
279/// # label: Some("My Buffer"),
280/// # size: 4,
281/// # usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
282/// # mapped_at_creation: false,
283/// # });
284/// #
285/// # let my_existing_selection_buffer = SelectionBuffer::new(&device, 1024);
286/// #
287/// // Create an editor that holds the buffers for the Gaussians
288/// let editor = Editor::new(
289/// &device,
290/// &vec![core::Gaussian {
291/// rot: Quat::IDENTITY,
292/// pos: Vec3::ZERO,
293/// color: U8Vec4::ZERO,
294/// sh: [Vec3::ZERO; 15],
295/// scale: Vec3::ONE,
296/// }],
297/// );
298///
299/// // Create the selection custom operation compute bundle
300/// let my_selection_custom_op_bundle = core::ComputeBundleBuilder::new()
301/// .label("My Selection")
302/// .bind_group_layouts([
303/// &SelectionBundle::<GaussianPod>::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
304/// &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR,
305/// ])
306/// .resolver({
307/// let mut resolver =
308/// wesl::StandardResolver::new("path/to/my/folder/containing/wesl");
309/// // Required for using core buffer structs.
310/// resolver.add_package(&core::shader::PACKAGE);
311/// // Optionally add this for some utility functions.
312/// resolver.add_package(&shader::PACKAGE);
313/// resolver
314/// })
315/// .main_shader("package::my_wesl_filename".parse().unwrap())
316/// .entry_point("main")
317/// .wesl_compile_options(wesl::CompileOptions {
318/// // Required for enabling the correct features for core struct.
319/// features: GaussianPod::wesl_features(),
320/// ..Default::default()
321/// })
322/// .build_without_bind_groups(&device)
323/// .map_err(|e| log::error!("{e}"))
324/// .expect("my selection custom op bundle");
325///
326/// // Create the selection bundle
327/// let selection_bundle =
328/// SelectionBundle::<GaussianPod>::new(&device, vec![my_selection_custom_op_bundle]);
329///
330/// // Create the bind group for your custom operation
331/// let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
332/// .create_bind_group(
333/// &device,
334/// 1, // Index 0 is always the Gaussians buffer
335/// [my_buffer.buffer().as_entire_binding()],
336/// )
337/// .unwrap();
338///
339/// // Create the selection expression
340/// let selection_expr = SelectionExpr::selection(
341/// 0, // The bundle index for your custom operation in the selection bundle
342/// vec![my_selection_custom_op_bind_group],
343/// )
344/// .union(
345/// // Combine with other selection expressions using different functions
346/// // Here is an existing selection buffer for example
347/// SelectionExpr::Buffer(my_existing_selection_buffer),
348/// );
349///
350/// // Create a selection buffer for the result
351/// let dest_selection_buffer =
352/// SelectionBuffer::new(&device, editor.gaussians_buffer.len() as u32);
353///
354/// # let mut encoder =
355/// # device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
356///
357/// // Evaluate the selection expression
358/// selection_bundle.evaluate(
359/// &device,
360/// &mut encoder,
361/// &selection_expr,
362/// &dest_selection_buffer,
363/// &editor.model_transform_buffer,
364/// &editor.gaussian_transform_buffer,
365/// &editor.gaussians_buffer,
366/// );
367/// # });
368/// ```
369///
370/// ## Shader Format
371///
372/// You may copy and paste the following shader bindings for
373/// [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`] into your custom selection operation
374/// shader to ensure that the bindings are correct, then add your own bindings after that.
375///
376/// ```wgsl
377/// import wgpu_3dgs_core::{
378/// gaussian::Gaussian,
379/// gaussian_transform::GaussianTransform,
380/// model_transform::{model_to_world, ModelTransform},
381/// };
382///
383/// @group(0) @binding(0)
384/// var<uniform> op: u32;
385///
386/// @group(0) @binding(1)
387/// var<storage, read> source: array<u32>;
388///
389/// @group(0) @binding(2)
390/// var<storage, read_write> dest: array<atomic<u32>>;
391///
392/// @group(0) @binding(3)
393/// var<uniform> model_transform: ModelTransform;
394///
395/// @group(0) @binding(4)
396/// var<uniform> gaussian_transform: GaussianTransform;
397///
398/// @group(0) @binding(5)
399/// var<storage, read> gaussians: array<Gaussian>;
400///
401/// // Your custom bindings here...
402///
403/// override workgroup_size: u32;
404///
405/// @compute @workgroup_size(workgroup_size)
406/// fn main(@builtin(global_invocation_id) id: vec3<u32>) {
407/// let index = id.x;
408///
409/// if index >= arrayLength(&gaussians) {
410/// return;
411/// }
412///
413/// let gaussian = gaussians[index];
414///
415/// let world_pos = model_to_world(model_transform, gaussian.position);
416///
417/// // Your custom selection operation code here...
418///
419/// let word_index = index / 32u;
420/// let bit_index = index % 32u;
421/// let bit_mask = 1u << bit_index;
422/// if /* Condition for selecting the Gaussian */ {
423/// atomicOr(&dest[word_index], bit_mask);
424/// } else {
425/// atomicAnd(&dest[word_index], ~bit_mask);
426/// }
427/// }
428/// ```
429#[derive(Debug)]
430pub struct SelectionBundle<G: GaussianPod> {
431 /// The compute bundle for primitive selection operations.
432 primitive_bundle: ComputeBundle<()>,
433 /// The compute bundles for selection custom operations.
434 pub bundles: Vec<ComputeBundle<()>>,
435 /// The Gaussian pod marker.
436 gaussian_pod_marker: std::marker::PhantomData<G>,
437}
438
439impl<G: GaussianPod> SelectionBundle<G> {
440 /// The Gaussians bind group layout descriptors.
441 ///
442 /// This bind group layout takes the following buffers:
443 /// - [`SelectionOpBuffer`]
444 /// - Source [`SelectionBuffer`]
445 /// - Destination [`SelectionBuffer`]
446 /// - [`ModelTransformBuffer`]
447 /// - [`GaussianTransformBuffer`]
448 /// - [`GaussiansBuffer`]
449 pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
450 wgpu::BindGroupLayoutDescriptor {
451 label: Some("Selection Gaussians Bind Group Layout"),
452 entries: &[
453 // Selection operation buffer
454 wgpu::BindGroupLayoutEntry {
455 binding: 0,
456 visibility: wgpu::ShaderStages::COMPUTE,
457 ty: wgpu::BindingType::Buffer {
458 ty: wgpu::BufferBindingType::Uniform,
459 has_dynamic_offset: false,
460 min_binding_size: None,
461 },
462 count: None,
463 },
464 // Source selection buffer
465 wgpu::BindGroupLayoutEntry {
466 binding: 1,
467 visibility: wgpu::ShaderStages::COMPUTE,
468 ty: wgpu::BindingType::Buffer {
469 ty: wgpu::BufferBindingType::Storage { read_only: true },
470 has_dynamic_offset: false,
471 min_binding_size: None,
472 },
473 count: None,
474 },
475 // Destination selection buffer
476 wgpu::BindGroupLayoutEntry {
477 binding: 2,
478 visibility: wgpu::ShaderStages::COMPUTE,
479 ty: wgpu::BindingType::Buffer {
480 ty: wgpu::BufferBindingType::Storage { read_only: false },
481 has_dynamic_offset: false,
482 min_binding_size: None,
483 },
484 count: None,
485 },
486 // Model transform buffer
487 wgpu::BindGroupLayoutEntry {
488 binding: 3,
489 visibility: wgpu::ShaderStages::COMPUTE,
490 ty: wgpu::BindingType::Buffer {
491 ty: wgpu::BufferBindingType::Uniform,
492 has_dynamic_offset: false,
493 min_binding_size: None,
494 },
495 count: None,
496 },
497 // Gaussian transform buffer
498 wgpu::BindGroupLayoutEntry {
499 binding: 4,
500 visibility: wgpu::ShaderStages::COMPUTE,
501 ty: wgpu::BindingType::Buffer {
502 ty: wgpu::BufferBindingType::Uniform,
503 has_dynamic_offset: false,
504 min_binding_size: None,
505 },
506 count: None,
507 },
508 // Gaussians buffer
509 wgpu::BindGroupLayoutEntry {
510 binding: 5,
511 visibility: wgpu::ShaderStages::COMPUTE,
512 ty: wgpu::BindingType::Buffer {
513 ty: wgpu::BufferBindingType::Storage { read_only: true },
514 has_dynamic_offset: false,
515 min_binding_size: None,
516 },
517 count: None,
518 },
519 ],
520 };
521
522 /// Create a new selection bundle.
523 ///
524 /// `bundles` are used for [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], or
525 /// [`SelectionExpr::Selection`] as custom operations, they must have the same bind group 0 as
526 /// the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`], see documentation of
527 /// [`SelectionBundle`] for more details.
528 pub fn new(device: &wgpu::Device, bundles: Vec<ComputeBundle<()>>) -> Self {
529 let primitive_bundle = Self::create_primitive_bundle(device);
530
531 Self {
532 primitive_bundle,
533 bundles,
534 gaussian_pod_marker: std::marker::PhantomData,
535 }
536 }
537
538 /// Get the Gaussians bind group layout.
539 pub fn gaussians_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
540 &self.primitive_bundle.bind_group_layouts()[0]
541 }
542
543 /// Evaluate and apply the selection expression.
544 #[allow(clippy::too_many_arguments)]
545 pub fn evaluate(
546 &self,
547 device: &wgpu::Device,
548 encoder: &mut wgpu::CommandEncoder,
549 expr: &SelectionExpr,
550 dest: &SelectionBuffer,
551 model_transform: &ModelTransformBuffer,
552 gaussian_transform: &GaussianTransformBuffer,
553 gaussians: &GaussiansBuffer<G>,
554 ) {
555 if let SelectionExpr::Identity = expr {
556 return;
557 } else if let SelectionExpr::Buffer(buffer) = expr {
558 encoder.copy_buffer_to_buffer(
559 buffer.buffer(),
560 0,
561 dest.buffer(),
562 0,
563 dest.buffer().size(),
564 );
565 return;
566 }
567
568 let d = dest;
569 let m = model_transform;
570 let g = gaussian_transform;
571 let gs = gaussians;
572
573 let op = SelectionOpBuffer::new(device, expr.as_u32().expect("operation expression"));
574 let source = SelectionBuffer::new(device, gaussians.len() as u32);
575
576 match expr {
577 SelectionExpr::Union(l, r) => {
578 self.evaluate(device, encoder, l, &source, m, g, gs);
579 self.evaluate(device, encoder, r, d, m, g, gs);
580 }
581 SelectionExpr::Intersection(l, r) => {
582 self.evaluate(device, encoder, l, &source, m, g, gs);
583 self.evaluate(device, encoder, r, d, m, g, gs);
584 }
585 SelectionExpr::Difference(l, r) => {
586 self.evaluate(device, encoder, l, &source, m, g, gs);
587 self.evaluate(device, encoder, r, d, m, g, gs);
588 }
589 SelectionExpr::SymmetricDifference(l, r) => {
590 self.evaluate(device, encoder, l, &source, m, g, gs);
591 self.evaluate(device, encoder, r, d, m, g, gs);
592 }
593 SelectionExpr::Complement(e) => {
594 self.evaluate(device, encoder, e, d, m, g, gs);
595 }
596 SelectionExpr::Unary(_, e, _) => {
597 self.evaluate(device, encoder, e, d, m, g, gs);
598 }
599 SelectionExpr::Binary(l, _, r, _) => {
600 self.evaluate(device, encoder, l, &source, m, g, gs);
601 self.evaluate(device, encoder, r, d, m, g, gs);
602 }
603 SelectionExpr::Selection(_, _) => {}
604 SelectionExpr::Identity | SelectionExpr::Buffer(_) => {
605 unreachable!();
606 }
607 }
608
609 let gaussians_bind_group = self
610 .primitive_bundle
611 .create_bind_group(
612 device,
613 0,
614 [
615 op.buffer().as_entire_binding(),
616 source.buffer().as_entire_binding(),
617 d.buffer().as_entire_binding(),
618 m.buffer().as_entire_binding(),
619 g.buffer().as_entire_binding(),
620 gs.buffer().as_entire_binding(),
621 ],
622 )
623 .expect("gaussians bind group");
624
625 match expr.custom_op_index_and_bind_groups() {
626 None => self.primitive_bundle.dispatch(
627 encoder,
628 (gaussians.len() as u32).div_ceil(32),
629 [&gaussians_bind_group],
630 ),
631 Some((i, bind_groups)) => {
632 let bind_groups = std::iter::once(&gaussians_bind_group)
633 .chain(bind_groups)
634 .collect::<Vec<_>>();
635
636 let bundle = &self.bundles[i];
637
638 bundle.dispatch(encoder, gaussians.len() as u32, bind_groups);
639 }
640 }
641 }
642
643 /// Create the selection primitive operation [`ComputeBundle`].
644 ///
645 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
646 ///
647 /// You usually do not need to use this method, it is used internally for creating the
648 /// primitive operation bundle for evaluation.
649 pub fn create_primitive_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
650 ComputeBundleBuilder::new()
651 .label("Selection Primitive Operations")
652 .bind_group_layout(&Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR)
653 .resolver({
654 let mut resolver = wesl::PkgResolver::new();
655 resolver.add_package(&core::shader::PACKAGE);
656 resolver.add_package(&shader::PACKAGE);
657 resolver
658 })
659 .main_shader(
660 "wgpu_3dgs_editor::selection::primitive"
661 .parse()
662 .expect("selection::primitive module path"),
663 )
664 .entry_point("main")
665 .wesl_compile_options(wesl::CompileOptions {
666 features: G::wesl_features(),
667 ..Default::default()
668 })
669 .build_without_bind_groups(device)
670 .map_err(|e| log::error!("{e}"))
671 .expect("primitive bundle")
672 }
673
674 /// The sphere selection bind group layout descriptor.
675 ///
676 /// This bind group layout takes the following buffers:
677 /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
678 pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
679 wgpu::BindGroupLayoutDescriptor {
680 label: Some("Sphere Selection Bind Group Layout"),
681 entries: &[
682 // Inverse transform uniform buffer
683 wgpu::BindGroupLayoutEntry {
684 binding: 0,
685 visibility: wgpu::ShaderStages::COMPUTE,
686 ty: wgpu::BindingType::Buffer {
687 ty: wgpu::BufferBindingType::Uniform,
688 has_dynamic_offset: false,
689 min_binding_size: None,
690 },
691 count: None,
692 },
693 ],
694 };
695
696 /// Create a sphere selection custom operation.
697 ///
698 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
699 /// - Bind group 1 is [`SelectionBundle::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR`].
700 pub fn create_sphere_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
701 let mut resolver = wesl::PkgResolver::new();
702 resolver.add_package(&core::shader::PACKAGE);
703 resolver.add_package(&shader::PACKAGE);
704
705 ComputeBundleBuilder::new()
706 .label("Sphere Selection")
707 .bind_group_layouts([
708 &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
709 &Self::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR,
710 ])
711 .main_shader(
712 "wgpu_3dgs_editor::selection::sphere"
713 .parse()
714 .expect("selection::sphere module path"),
715 )
716 .entry_point("main")
717 .wesl_compile_options(wesl::CompileOptions {
718 features: G::wesl_features(),
719 ..Default::default()
720 })
721 .resolver(resolver)
722 .build_without_bind_groups(device)
723 .map_err(|e| log::error!("{e}"))
724 .expect("sphere selection compute bundle")
725 }
726
727 /// The box selection bind group layout descriptor.
728 ///
729 /// This bind group layout takes the following buffers:
730 /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
731 pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
732 wgpu::BindGroupLayoutDescriptor {
733 label: Some("Box Selection Bind Group Layout"),
734 entries: &[
735 // Inverse transform uniform buffer
736 wgpu::BindGroupLayoutEntry {
737 binding: 0,
738 visibility: wgpu::ShaderStages::COMPUTE,
739 ty: wgpu::BindingType::Buffer {
740 ty: wgpu::BufferBindingType::Uniform,
741 has_dynamic_offset: false,
742 min_binding_size: None,
743 },
744 count: None,
745 },
746 ],
747 };
748
749 /// Create a box selection custom operation.
750 ///
751 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
752 /// - Bind group 1 is [`SelectionBundle::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR`].
753 pub fn create_box_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
754 let mut resolver = wesl::PkgResolver::new();
755 resolver.add_package(&core::shader::PACKAGE);
756 resolver.add_package(&shader::PACKAGE);
757
758 ComputeBundleBuilder::new()
759 .label("Box Selection")
760 .bind_group_layouts([
761 &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
762 &Self::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR,
763 ])
764 .main_shader(
765 "wgpu_3dgs_editor::selection::box"
766 .parse()
767 .expect("selection::box module path"),
768 )
769 .entry_point("main")
770 .wesl_compile_options(wesl::CompileOptions {
771 features: G::wesl_features(),
772 ..Default::default()
773 })
774 .resolver(resolver)
775 .build_without_bind_groups(device)
776 .map_err(|e| log::error!("{e}"))
777 .expect("box selection compute bundle")
778 }
779}