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(&wgpu::InstanceDescriptor::default());
246/// #
247/// # let adapter = instance
248/// # .request_adapter(&wgpu::RequestAdapterOptions::default())
249/// # .await
250/// # .expect("adapter");
251/// #
252/// # let (device, _queue) = adapter
253/// # .request_device(&wgpu::DeviceDescriptor {
254/// # label: Some("Device"),
255/// # required_limits: adapter.limits(),
256/// # ..Default::default()
257/// # })
258/// # .await
259/// # .expect("device");
260/// #
261/// # const MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor =
262/// # wgpu::BindGroupLayoutDescriptor {
263/// # label: Some("My Custom Bind Group Layout"),
264/// # entries: &[wgpu::BindGroupLayoutEntry {
265/// # binding: 0,
266/// # visibility: wgpu::ShaderStages::COMPUTE,
267/// # ty: wgpu::BindingType::Buffer {
268/// # ty: wgpu::BufferBindingType::Uniform,
269/// # has_dynamic_offset: false,
270/// # min_binding_size: None,
271/// # },
272/// # count: None,
273/// # }],
274/// # };
275/// #
276/// # let my_buffer = device.create_buffer(&wgpu::BufferDescriptor {
277/// # label: Some("My Buffer"),
278/// # size: 4,
279/// # usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
280/// # mapped_at_creation: false,
281/// # });
282/// #
283/// # let my_existing_selection_buffer = SelectionBuffer::new(&device, 1024);
284/// #
285/// // Create an editor that holds the buffers for the Gaussians
286/// let editor = Editor::new(
287/// &device,
288/// &vec![core::Gaussian {
289/// rot: Quat::IDENTITY,
290/// pos: Vec3::ZERO,
291/// color: U8Vec4::ZERO,
292/// sh: [Vec3::ZERO; 15],
293/// scale: Vec3::ONE,
294/// }],
295/// );
296///
297/// // Create the selection custom operation compute bundle
298/// let my_selection_custom_op_bundle = core::ComputeBundleBuilder::new()
299/// .label("My Selection")
300/// .bind_group_layouts([
301/// &SelectionBundle::<GaussianPod>::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
302/// &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR,
303/// ])
304/// .resolver({
305/// let mut resolver =
306/// wesl::StandardResolver::new("path/to/my/folder/containing/wesl");
307/// // Required for using core buffer structs.
308/// resolver.add_package(&core::shader::PACKAGE);
309/// // Optionally add this for some utility functions.
310/// resolver.add_package(&shader::PACKAGE);
311/// resolver
312/// })
313/// .main_shader("package::my_wesl_filename".parse().unwrap())
314/// .entry_point("main")
315/// .wesl_compile_options(wesl::CompileOptions {
316/// // Required for enabling the correct features for core struct.
317/// features: GaussianPod::wesl_features(),
318/// ..Default::default()
319/// })
320/// .build_without_bind_groups(&device)
321/// .map_err(|e| log::error!("{e}"))
322/// .expect("my selection custom op bundle");
323///
324/// // Create the selection bundle
325/// let selection_bundle =
326/// SelectionBundle::<GaussianPod>::new(&device, vec![my_selection_custom_op_bundle]);
327///
328/// // Create the bind group for your custom operation
329/// let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
330/// .create_bind_group(
331/// &device,
332/// 1, // Index 0 is always the Gaussians buffer
333/// [my_buffer.buffer().as_entire_binding()],
334/// )
335/// .unwrap();
336///
337/// // Create the selection expression
338/// let selection_expr = SelectionExpr::selection(
339/// 0, // The bundle index for your custom operation in the selection bundle
340/// vec![my_selection_custom_op_bind_group],
341/// )
342/// .union(
343/// // Combine with other selection expressions using different functions
344/// // Here is an existing selection buffer for example
345/// SelectionExpr::Buffer(my_existing_selection_buffer),
346/// );
347///
348/// // Create a selection buffer for the result
349/// let dest_selection_buffer =
350/// SelectionBuffer::new(&device, editor.gaussians_buffer.len() as u32);
351///
352/// # let mut encoder =
353/// # device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
354///
355/// // Evaluate the selection expression
356/// selection_bundle.evaluate(
357/// &device,
358/// &mut encoder,
359/// &selection_expr,
360/// &dest_selection_buffer,
361/// &editor.model_transform_buffer,
362/// &editor.gaussian_transform_buffer,
363/// &editor.gaussians_buffer,
364/// );
365/// # });
366/// ```
367///
368/// ## Shader Format
369///
370/// You may copy and paste the following shader bindings for
371/// [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`] into your custom selection operation
372/// shader to ensure that the bindings are correct, then add your own bindings after that.
373///
374/// ```wgsl
375/// import wgpu_3dgs_core::{
376/// gaussian::Gaussian,
377/// gaussian_transform::GaussianTransform,
378/// model_transform::{model_to_world, ModelTransform},
379/// };
380///
381/// @group(0) @binding(0)
382/// var<uniform> op: u32;
383///
384/// @group(0) @binding(1)
385/// var<storage, read> source: array<u32>;
386///
387/// @group(0) @binding(2)
388/// var<storage, read_write> dest: array<atomic<u32>>;
389///
390/// @group(0) @binding(3)
391/// var<uniform> model_transform: ModelTransform;
392///
393/// @group(0) @binding(4)
394/// var<uniform> gaussian_transform: GaussianTransform;
395///
396/// @group(0) @binding(5)
397/// var<storage, read> gaussians: array<Gaussian>;
398///
399/// // Your custom bindings here...
400///
401/// override workgroup_size: u32;
402///
403/// @compute @workgroup_size(workgroup_size)
404/// fn main(@builtin(global_invocation_id) id: vec3<u32>) {
405/// let index = id.x;
406///
407/// if index >= arrayLength(&gaussians) {
408/// return;
409/// }
410///
411/// let gaussian = gaussians[index];
412///
413/// let world_pos = model_to_world(model_transform, gaussian.position);
414///
415/// // Your custom selection operation code here...
416///
417/// let word_index = index / 32u;
418/// let bit_index = index % 32u;
419/// let bit_mask = 1u << bit_index;
420/// if /* Condition for selecting the Gaussian */ {
421/// atomicOr(&dest[word_index], bit_mask);
422/// } else {
423/// atomicAnd(&dest[word_index], ~bit_mask);
424/// }
425/// }
426/// ```
427#[derive(Debug)]
428pub struct SelectionBundle<G: GaussianPod> {
429 /// The compute bundle for primitive selection operations.
430 primitive_bundle: ComputeBundle<()>,
431 /// The compute bundles for selection custom operations.
432 pub bundles: Vec<ComputeBundle<()>>,
433 /// The Gaussian pod marker.
434 gaussian_pod_marker: std::marker::PhantomData<G>,
435}
436
437impl<G: GaussianPod> SelectionBundle<G> {
438 /// The Gaussians bind group layout descriptors.
439 ///
440 /// This bind group layout takes the following buffers:
441 /// - [`SelectionOpBuffer`]
442 /// - Source [`SelectionBuffer`]
443 /// - Destination [`SelectionBuffer`]
444 /// - [`ModelTransformBuffer`]
445 /// - [`GaussianTransformBuffer`]
446 /// - [`GaussiansBuffer`]
447 pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
448 wgpu::BindGroupLayoutDescriptor {
449 label: Some("Selection Gaussians Bind Group Layout"),
450 entries: &[
451 // Selection operation buffer
452 wgpu::BindGroupLayoutEntry {
453 binding: 0,
454 visibility: wgpu::ShaderStages::COMPUTE,
455 ty: wgpu::BindingType::Buffer {
456 ty: wgpu::BufferBindingType::Uniform,
457 has_dynamic_offset: false,
458 min_binding_size: None,
459 },
460 count: None,
461 },
462 // Source selection buffer
463 wgpu::BindGroupLayoutEntry {
464 binding: 1,
465 visibility: wgpu::ShaderStages::COMPUTE,
466 ty: wgpu::BindingType::Buffer {
467 ty: wgpu::BufferBindingType::Storage { read_only: true },
468 has_dynamic_offset: false,
469 min_binding_size: None,
470 },
471 count: None,
472 },
473 // Destination selection buffer
474 wgpu::BindGroupLayoutEntry {
475 binding: 2,
476 visibility: wgpu::ShaderStages::COMPUTE,
477 ty: wgpu::BindingType::Buffer {
478 ty: wgpu::BufferBindingType::Storage { read_only: false },
479 has_dynamic_offset: false,
480 min_binding_size: None,
481 },
482 count: None,
483 },
484 // Model transform buffer
485 wgpu::BindGroupLayoutEntry {
486 binding: 3,
487 visibility: wgpu::ShaderStages::COMPUTE,
488 ty: wgpu::BindingType::Buffer {
489 ty: wgpu::BufferBindingType::Uniform,
490 has_dynamic_offset: false,
491 min_binding_size: None,
492 },
493 count: None,
494 },
495 // Gaussian transform buffer
496 wgpu::BindGroupLayoutEntry {
497 binding: 4,
498 visibility: wgpu::ShaderStages::COMPUTE,
499 ty: wgpu::BindingType::Buffer {
500 ty: wgpu::BufferBindingType::Uniform,
501 has_dynamic_offset: false,
502 min_binding_size: None,
503 },
504 count: None,
505 },
506 // Gaussians buffer
507 wgpu::BindGroupLayoutEntry {
508 binding: 5,
509 visibility: wgpu::ShaderStages::COMPUTE,
510 ty: wgpu::BindingType::Buffer {
511 ty: wgpu::BufferBindingType::Storage { read_only: true },
512 has_dynamic_offset: false,
513 min_binding_size: None,
514 },
515 count: None,
516 },
517 ],
518 };
519
520 /// Create a new selection bundle.
521 ///
522 /// `bundles` are used for [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], or
523 /// [`SelectionExpr::Selection`] as custom operations, they must have the same bind group 0 as
524 /// the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`], see documentation of
525 /// [`SelectionBundle`] for more details.
526 pub fn new(device: &wgpu::Device, bundles: Vec<ComputeBundle<()>>) -> Self {
527 let primitive_bundle = Self::create_primitive_bundle(device);
528
529 Self {
530 primitive_bundle,
531 bundles,
532 gaussian_pod_marker: std::marker::PhantomData,
533 }
534 }
535
536 /// Get the Gaussians bind group layout.
537 pub fn gaussians_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
538 &self.primitive_bundle.bind_group_layouts()[0]
539 }
540
541 /// Evaluate and apply the selection expression.
542 #[allow(clippy::too_many_arguments)]
543 pub fn evaluate(
544 &self,
545 device: &wgpu::Device,
546 encoder: &mut wgpu::CommandEncoder,
547 expr: &SelectionExpr,
548 dest: &SelectionBuffer,
549 model_transform: &ModelTransformBuffer,
550 gaussian_transform: &GaussianTransformBuffer,
551 gaussians: &GaussiansBuffer<G>,
552 ) {
553 if let SelectionExpr::Identity = expr {
554 return;
555 } else if let SelectionExpr::Buffer(buffer) = expr {
556 encoder.copy_buffer_to_buffer(
557 buffer.buffer(),
558 0,
559 dest.buffer(),
560 0,
561 dest.buffer().size(),
562 );
563 return;
564 }
565
566 let d = dest;
567 let m = model_transform;
568 let g = gaussian_transform;
569 let gs = gaussians;
570
571 let op = SelectionOpBuffer::new(device, expr.as_u32().expect("operation expression"));
572 let source = SelectionBuffer::new(device, gaussians.len() as u32);
573
574 match expr {
575 SelectionExpr::Union(l, r) => {
576 self.evaluate(device, encoder, l, &source, m, g, gs);
577 self.evaluate(device, encoder, r, d, m, g, gs);
578 }
579 SelectionExpr::Intersection(l, r) => {
580 self.evaluate(device, encoder, l, &source, m, g, gs);
581 self.evaluate(device, encoder, r, d, m, g, gs);
582 }
583 SelectionExpr::Difference(l, r) => {
584 self.evaluate(device, encoder, l, &source, m, g, gs);
585 self.evaluate(device, encoder, r, d, m, g, gs);
586 }
587 SelectionExpr::SymmetricDifference(l, r) => {
588 self.evaluate(device, encoder, l, &source, m, g, gs);
589 self.evaluate(device, encoder, r, d, m, g, gs);
590 }
591 SelectionExpr::Complement(e) => {
592 self.evaluate(device, encoder, e, d, m, g, gs);
593 }
594 SelectionExpr::Unary(_, e, _) => {
595 self.evaluate(device, encoder, e, d, m, g, gs);
596 }
597 SelectionExpr::Binary(l, _, r, _) => {
598 self.evaluate(device, encoder, l, &source, m, g, gs);
599 self.evaluate(device, encoder, r, d, m, g, gs);
600 }
601 SelectionExpr::Selection(_, _) => {}
602 SelectionExpr::Identity | SelectionExpr::Buffer(_) => {
603 unreachable!();
604 }
605 }
606
607 let gaussians_bind_group = self
608 .primitive_bundle
609 .create_bind_group(
610 device,
611 0,
612 [
613 op.buffer().as_entire_binding(),
614 source.buffer().as_entire_binding(),
615 d.buffer().as_entire_binding(),
616 m.buffer().as_entire_binding(),
617 g.buffer().as_entire_binding(),
618 gs.buffer().as_entire_binding(),
619 ],
620 )
621 .expect("gaussians bind group");
622
623 match expr.custom_op_index_and_bind_groups() {
624 None => self.primitive_bundle.dispatch(
625 encoder,
626 (gaussians.len() as u32).div_ceil(32),
627 [&gaussians_bind_group],
628 ),
629 Some((i, bind_groups)) => {
630 let bind_groups = std::iter::once(&gaussians_bind_group)
631 .chain(bind_groups)
632 .collect::<Vec<_>>();
633
634 let bundle = &self.bundles[i];
635
636 bundle.dispatch(encoder, gaussians.len() as u32, bind_groups);
637 }
638 }
639 }
640
641 /// Create the selection primitive operation [`ComputeBundle`].
642 ///
643 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
644 ///
645 /// You usually do not need to use this method, it is used internally for creating the
646 /// primitive operation bundle for evaluation.
647 pub fn create_primitive_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
648 ComputeBundleBuilder::new()
649 .label("Selection Primitive Operations")
650 .bind_group_layout(&Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR)
651 .resolver({
652 let mut resolver = wesl::PkgResolver::new();
653 resolver.add_package(&core::shader::PACKAGE);
654 resolver.add_package(&shader::PACKAGE);
655 resolver
656 })
657 .main_shader(
658 "wgpu_3dgs_editor::selection::primitive"
659 .parse()
660 .expect("selection::primitive module path"),
661 )
662 .entry_point("main")
663 .wesl_compile_options(wesl::CompileOptions {
664 features: G::wesl_features(),
665 ..Default::default()
666 })
667 .build_without_bind_groups(device)
668 .map_err(|e| log::error!("{e}"))
669 .expect("primitive bundle")
670 }
671
672 /// The sphere selection bind group layout descriptor.
673 ///
674 /// This bind group layout takes the following buffers:
675 /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
676 pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
677 wgpu::BindGroupLayoutDescriptor {
678 label: Some("Sphere Selection Bind Group Layout"),
679 entries: &[
680 // Inverse transform uniform buffer
681 wgpu::BindGroupLayoutEntry {
682 binding: 0,
683 visibility: wgpu::ShaderStages::COMPUTE,
684 ty: wgpu::BindingType::Buffer {
685 ty: wgpu::BufferBindingType::Uniform,
686 has_dynamic_offset: false,
687 min_binding_size: None,
688 },
689 count: None,
690 },
691 ],
692 };
693
694 /// Create a sphere selection custom operation.
695 ///
696 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
697 /// - Bind group 1 is [`SelectionBundle::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR`].
698 pub fn create_sphere_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
699 let mut resolver = wesl::PkgResolver::new();
700 resolver.add_package(&core::shader::PACKAGE);
701 resolver.add_package(&shader::PACKAGE);
702
703 ComputeBundleBuilder::new()
704 .label("Sphere Selection")
705 .bind_group_layouts([
706 &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
707 &Self::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR,
708 ])
709 .main_shader(
710 "wgpu_3dgs_editor::selection::sphere"
711 .parse()
712 .expect("selection::sphere module path"),
713 )
714 .entry_point("main")
715 .wesl_compile_options(wesl::CompileOptions {
716 features: G::wesl_features(),
717 ..Default::default()
718 })
719 .resolver(resolver)
720 .build_without_bind_groups(device)
721 .map_err(|e| log::error!("{e}"))
722 .expect("sphere selection compute bundle")
723 }
724
725 /// The box selection bind group layout descriptor.
726 ///
727 /// This bind group layout takes the following buffers:
728 /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
729 pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
730 wgpu::BindGroupLayoutDescriptor {
731 label: Some("Box Selection Bind Group Layout"),
732 entries: &[
733 // Inverse transform uniform buffer
734 wgpu::BindGroupLayoutEntry {
735 binding: 0,
736 visibility: wgpu::ShaderStages::COMPUTE,
737 ty: wgpu::BindingType::Buffer {
738 ty: wgpu::BufferBindingType::Uniform,
739 has_dynamic_offset: false,
740 min_binding_size: None,
741 },
742 count: None,
743 },
744 ],
745 };
746
747 /// Create a box selection custom operation.
748 ///
749 /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
750 /// - Bind group 1 is [`SelectionBundle::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR`].
751 pub fn create_box_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
752 let mut resolver = wesl::PkgResolver::new();
753 resolver.add_package(&core::shader::PACKAGE);
754 resolver.add_package(&shader::PACKAGE);
755
756 ComputeBundleBuilder::new()
757 .label("Box Selection")
758 .bind_group_layouts([
759 &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
760 &Self::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR,
761 ])
762 .main_shader(
763 "wgpu_3dgs_editor::selection::box"
764 .parse()
765 .expect("selection::box module path"),
766 )
767 .entry_point("main")
768 .wesl_compile_options(wesl::CompileOptions {
769 features: G::wesl_features(),
770 ..Default::default()
771 })
772 .resolver(resolver)
773 .build_without_bind_groups(device)
774 .map_err(|e| log::error!("{e}"))
775 .expect("box selection compute bundle")
776 }
777}