Trait Modifier

Source

pub trait Modifier<G: GaussianPod> {
    // Required method
    fn apply(
        &self,
        device: &Device,
        encoder: &mut CommandEncoder,
        gaussians: &GaussiansBuffer<G>,
        model_transform: &ModelTransformBuffer,
        gaussian_transform: &GaussianTransformBuffer,
    );
}

Expand description

A trait to apply modifier to Gaussians.

§Overview

This trait simply defines a method to apply modifications to a set of Gaussians stored in a GaussiansBuffer. It makes it convenient for users to apply a sequence of modifications.

The trait is also blanket implemented for closures with the same signature, allowing users to easily create modifier closures instead of having to define a modifier struct.

Editor also provides an apply method which takes a slice of Modifiers to apply them in sequence to the stored Gaussians.

§Usage

There are many ways to use this but the recommended way is to implement this trait for a closure which dispatch a ComputeBundle.

// Create an editor that holds the buffers for the Gaussians and will apply the modifier
let editor = Editor::new(
    &device,
    &vec![core::Gaussian {
        rot: Quat::IDENTITY,
        pos: Vec3::ZERO,
        color: U8Vec4::ZERO,
        sh: [Vec3::ZERO; 15],
        scale: Vec3::ONE,
    }],
);

// Create the modifier compute bundle
let my_modifier_bundle = core::ComputeBundleBuilder::new()
    .label("My Modifier")
    .bind_group_layouts([
        // For accessing Gaussians and transforms
        &MODIFIER_GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
        // Your custom bind group layout here
        &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR,
    ])
    .resolver({
        let mut resolver =
            wesl::StandardResolver::new("path/to/my/folder/containing/wesl");
        // Required for using core buffer structs.
        resolver.add_package(&core::shader::PACKAGE);
        // Optionally add this for some utility functions.
        resolver.add_package(&shader::PACKAGE);
        resolver
    })
    .main_shader("package::my_wesl_filename".parse().unwrap())
    .entry_point("main")
    .wesl_compile_options(wesl::CompileOptions {
        // Required for enabling the correct features for core struct.
        features: GaussianPod::wesl_features(),
        ..Default::default()
    })
    .build(
        &device,
        [
            vec![
                editor.gaussians_buffer.buffer().as_entire_binding(),
                editor.model_transform_buffer.buffer().as_entire_binding(),
                editor.gaussian_transform_buffer.buffer().as_entire_binding(),
            ],
            vec![my_buffer.buffer().as_entire_binding()],
        ],
    )
    .map_err(|e| log::error!("{e}"))
    .expect("my modifier bundle");

// Create the modifier closure
// This function signature implements Modifier by default
let my_modifier =
    |_device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        gaussians: &GaussiansBuffer<GaussianPod>,
        _model_transform: &ModelTransformBuffer,
        _gaussian_transform: &GaussianTransformBuffer| {
        my_modifier_bundle.dispatch(encoder, gaussians.len() as u32);
    };


// Apply the modifier using the editor
editor.apply(
    &device,
    &mut encoder,
    [&my_modifier as &dyn Modifier<GaussianPod>],
);

§Shader Format

You may copy and paste the following shader bindings for MODIFIER_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::ModelTransform,
};

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

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

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

// Your custom bindings here...
//
// You may also apply modifier to selected gaussians only by adding:
// @group(1) @binding(N)
// var<storage, read> selection: array<u32>;

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;
    }

    @if(/* using selection buffer */) {
        let word_index = index / 32u;
        let bit_index = index % 32u;
        let bit_mask = 1u << bit_index;
        if (selection[word_index] & bit_mask) == 0 {
            return;
        }
    }

    var gaussian = gaussians[index];

    // Your custom modifier operation code here...

    gaussians[index] = gaussian;
}