22800469877

Committed 07 Mar 2026 02:04PM UTC coverage: 57.546% (-0.2%) from 57.75%

Build # 22800469877

Build Type

push

github

Committed by

web-flow

Commit Message

Batch spawners and properties (#525)

Bind the entire spawner and property arrays in all passes, instead of a
single entry. This removes the need to pad those structures. Access them
via an offset in the effect's own metadata.

Add a new `BatchInfo` struct holding the per-batch data. This clarifies
the responsibilites between this and `EffectMetadata`, the latter
holding per-effect (not per-batch) data.

Remove the `DispatchBufferIndices` component, which was used to track
the allocated entry for indirect compute dispatch. Instead, align those
allocations 1:1 with the GPU batch info allocations, which are
re-computed each frame.

Add a prefix sum pass before the update pass, which computes the prefix
sum of alive particles after the init pass. This is used to enable
batched update compute dispatch, where the number of compute threads
maps to the total number of alive particles in the batch. In that case,
we need to find which thread updates which particle of which batch,
using that prefix sum. Note that in this change, due to other
limitations still present, each effect instance is still in its own
batch (there's effectively no batching). Enabling full batching requires
more work, notably on the sort pass for ribbons, and the GPU-based init
pass with GPU events.

Change the allocation of spawners to occur after sorting. This ensures
all effects in a same batch have sequential allocations, which enables
accessing those spawners with a simple {offset + index} strategy.

Change the render pass to use the same bind group "spawner@2" than other
passes. This binds the property buffer, although in this change the
metadata buffer is still not available, so properties can't be used yet
in the render pass. The bind groups should be reviewed anyway because,
with batching approaching, and with the current change, assumptions
about frequency of changes is now wrong, and individual bindings should
be re-grouped in more suitable frequency-based groups.

Finally, the... (continued)

Coverage Stats

193 of 404 new or added lines in 7 files covered. (47.77%)

26 existing lines in 3 files now uncovered.

4793 of 8329 relevant lines covered (57.55%)

198.51 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

28.18

/src/render/sort.rs

use std::num::{NonZeroU32, NonZeroU64};

use bevy::{
    asset::Handle,
    ecs::{resource::Resource, world::World},
    platform::collections::{hash_map::Entry, HashMap},
    render::{
        render_resource::{
            binding_types::{
                storage_buffer, storage_buffer_read_only, storage_buffer_read_only_sized,
                storage_buffer_sized,
            },
            BindGroup, BindGroupEntries, BindGroupLayoutDescriptor, BindGroupLayoutEntries, Buffer,
            BufferId, CachedComputePipelineId, CachedPipelineState, ComputePipelineDescriptor,
            PipelineCache, ShaderType,
        },
        renderer::RenderDevice,
    },
    shader::Shader,
    utils::default,
};
use bytemuck::{Pod, Zeroable};
use wgpu::{BufferBinding, BufferDescriptor, BufferUsages, CommandEncoder, ShaderStages};

use super::{gpu_buffer::GpuBuffer, GpuDispatchIndirectArgs, GpuEffectMetadata, StorageType};
use crate::{
    render::{GpuIndirectIndex, GpuSpawnerParams},
    Attribute, ParticleLayout,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct SortFillBindGroupLayoutKey {
    particle_min_binding_size: NonZeroU32,
    particle_ribbon_id_offset: u32,
    particle_age_offset: u32,
}

impl SortFillBindGroupLayoutKey {
    pub fn from_particle_layout(particle_layout: &ParticleLayout) -> Result<Self, ()> {
        let particle_ribbon_id_offset = particle_layout
            .byte_offset(Attribute::RIBBON_ID)
            .ok_or(())?;
        let particle_age_offset = particle_layout.byte_offset(Attribute::AGE).ok_or(())?;
        let key = SortFillBindGroupLayoutKey {
            particle_min_binding_size: particle_layout.min_binding_size32(),
            particle_ribbon_id_offset,
            particle_age_offset,
        };
        Ok(key)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct SortFillBindGroupKey {
    particle: BufferId,
    indirect_index: BufferId,
    effect_metadata: BufferId,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct SortCopyBindGroupKey {
    indirect_index: BufferId,
    sort: BufferId,
    effect_metadata: BufferId,
}

/// GPU representation of a single dual-key value pair, with the added buffer
/// count as prefix. This is mainly used for shorcuts in bindings, not directly
/// as a type.
#[repr(C)]
#[derive(Debug, Copy, Clone, Pod, Zeroable, ShaderType)]
struct GpuSortBufferSingleEntry {
    /// Number of key-value pairs to sort. This is the first element of the
    /// entire buffer.
    pub count: u32,
    /// Key for the first entry.
    pub key: u32,
    /// Secondary key for the first entry.
    pub key2: u32,
    /// Value for the first entry.
    pub value: u32,
}

#[derive(Resource)]
pub struct SortBindGroups {
    /// Render device.
    render_device: RenderDevice,
    /// Sort-fill pass compute shader.
    sort_fill_shader: Handle<Shader>,
    /// GPU buffer of key-value pairs to sort.
    sort_buffer: Buffer,
    /// GPU buffer containing the [`GpuDispatchIndirect`] structs for the
    /// sort-fill and sort passes.
    indirect_buffer: GpuBuffer<GpuDispatchIndirectArgs>,
    /// Bind group layouts for group #0 of the sort-fill compute pass.
    sort_fill_bind_group_layout_descs:
        HashMap<SortFillBindGroupLayoutKey, (BindGroupLayoutDescriptor, CachedComputePipelineId)>,
    /// Bind groups for group #0 of the sort-fill compute pass.
    sort_fill_bind_groups: HashMap<SortFillBindGroupKey, BindGroup>,
    /// Bind group layout descriptor for group #0 of the sort compute pass.
    sort_bind_group_layout_desc: BindGroupLayoutDescriptor,
    /// Bind group for group #0 of the sort compute pass.
    sort_bind_group: Option<BindGroup>,
    sort_copy_bind_group_layout_desc: BindGroupLayoutDescriptor,
    /// Pipeline for sort pass.
    sort_pipeline_id: CachedComputePipelineId,
    /// Pipeline for sort-copy pass.
    sort_copy_pipeline_id: CachedComputePipelineId,
    /// Bind groups for group #0 of the sort-copy compute pass.
    sort_copy_bind_groups: HashMap<SortCopyBindGroupKey, BindGroup>,
}

impl SortBindGroups {
    pub fn new(
        world: &mut World,
        sort_fill_shader: Handle<Shader>,
        sort_shader: Handle<Shader>,
        sort_copy_shader: Handle<Shader>,
    ) -> Self {
        let render_device = world.resource::<RenderDevice>();
        let pipeline_cache = world.resource::<PipelineCache>();

        let sort_buffer = render_device.create_buffer(&BufferDescriptor {
            label: Some("hanabi:buffer:sort:pairs"),
            size: 3 * 1024 * 1024,
            usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
            mapped_at_creation: false,
        });

        let indirect_buffer_size = 3 * 1024;
        let indirect_buffer = render_device.create_buffer(&BufferDescriptor {
            label: Some("hanabi:buffer:sort:indirect"),
            size: indirect_buffer_size,
            usage: BufferUsages::COPY_SRC
                | BufferUsages::COPY_DST
                | BufferUsages::STORAGE
                | BufferUsages::INDIRECT,
            mapped_at_creation: false,
        });
        let indirect_buffer = GpuBuffer::new_allocated(
            indirect_buffer,
            indirect_buffer_size as u32,
            Some("hanabi:buffer:sort:indirect".to_string()),
        );

        let sort_bind_group_layout_desc = BindGroupLayoutDescriptor::new(
            "hanabi:bgl:sort",
            &BindGroupLayoutEntries::single(
                ShaderStages::COMPUTE,
                storage_buffer_sized(false, Some(NonZeroU64::new(16).unwrap())), /* count + dual
                                                                                  * kv pair */
            ),
        );

        let sort_pipeline_id = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
            label: Some("hanabi:pipeline:sort".into()),
            layout: vec![sort_bind_group_layout_desc.clone()],
            shader: sort_shader,
            shader_defs: vec!["HAS_DUAL_KEY".into()],
            entry_point: Some("main".into()),
            push_constant_ranges: vec![],
            zero_initialize_workgroup_memory: false,
        });

        let alignment = render_device.limits().min_storage_buffer_offset_alignment;
        let sort_copy_bind_group_layout_desc = BindGroupLayoutDescriptor::new(
            "hanabi:bgl:sort_copy",
            &BindGroupLayoutEntries::sequential(
                ShaderStages::COMPUTE,
                (
                    // @group(0) @binding(0) var<storage, read_write> indirect_index_buffer :
                    // IndirectIndexBuffer;
                    storage_buffer::<GpuIndirectIndex>(false),
                    // @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
                    storage_buffer_read_only::<GpuSortBufferSingleEntry>(false),
                    // @group(0) @binding(2) var<storage, read_write> effect_metadata :
                    // EffectMetadata;
                    storage_buffer_sized(true, Some(GpuEffectMetadata::aligned_size(alignment))),
                    // @group(0) @binding(3) var<storage, read> spawner : Spawner;
                    storage_buffer_read_only_sized(
                        true,
                        Some(GpuSpawnerParams::aligned_size(alignment)),
                    ),
                ),
            ),
        );

        let sort_copy_pipeline_id =
            pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
                label: Some("hanabi:pipeline:sort_copy".into()),
                layout: vec![sort_copy_bind_group_layout_desc.clone()],
                shader: sort_copy_shader,
                shader_defs: vec![],
                entry_point: Some("main".into()),
                push_constant_ranges: vec![],
                zero_initialize_workgroup_memory: false,
            });

        Self {
            render_device: render_device.clone(),
            sort_fill_shader,
            sort_buffer,
            indirect_buffer,
            sort_fill_bind_group_layout_descs: default(),
            sort_fill_bind_groups: default(),
            sort_bind_group_layout_desc,
            // This bind group is created later, once the pipeline and its bind group layouts are
            // created by the pipeline cache. Technically we could create the bind group layout
            // immediately because the PipelineCache pretends to but actually creates them
            // on-the-fly in get_bind_group_layout(), but this is brittle as any behavior change
            // would break Hanabi. Instead we create this bind group alongside all others, which is
            // more consistent too.
            sort_bind_group: None,
            sort_copy_bind_group_layout_desc,
            sort_pipeline_id,
            sort_copy_pipeline_id,
            sort_copy_bind_groups: default(),
        }
    }

    #[inline]
    pub fn clear_indirect_dispatch_buffer(&mut self) {
        self.indirect_buffer.clear();
    }

    #[inline]
    pub fn allocate_indirect_dispatch(&mut self) -> u32 {
        self.indirect_buffer.allocate()
    }

    #[inline]
    pub fn get_indirect_dispatch_byte_offset(&self, index: u32) -> u32 {
        self.indirect_buffer.item_size() as u32 * index
    }

    #[inline]
    #[allow(dead_code)]
    pub fn sort_buffer(&self) -> &Buffer {
        &self.sort_buffer
    }

    #[inline]
    pub fn indirect_buffer(&self) -> Option<&Buffer> {
        self.indirect_buffer.buffer()
    }

    #[inline]
    pub fn sort_pipeline_id(&self) -> CachedComputePipelineId {
        self.sort_pipeline_id
    }

    /// Check if the sort pipeline is ready to run for the given effect
    /// instance.
    ///
    /// This ensures all compute pipelines are compiled and ready to be used
    /// this frame.
    pub fn is_pipeline_ready(
        &self,
        particle_layout: &ParticleLayout,
        pipeline_cache: &PipelineCache,
    ) -> bool {
        // Validate the sort-fill pipeline. It was created and queued for compile by
        // ensure_sort_fill_bind_group_layout(), which normally is called just before
        // is_pipeline_ready().
        let Some(pipeline_id) = self.get_sort_fill_pipeline_id(particle_layout) else {
            return false;
        };
        if !matches!(
            pipeline_cache.get_compute_pipeline_state(pipeline_id),
            CachedPipelineState::Ok(_)
        ) {
            return false;
        }

        // The 2 pipelines below are created and queued for compile in new(), so are
        // almost always ready.
        // FIXME - they could be checked once a frame only, not once per effect...

        // Validate the sort pipeline
        if !matches!(
            pipeline_cache.get_compute_pipeline_state(self.sort_pipeline_id()),
            CachedPipelineState::Ok(_)
        ) {
            return false;
        }

        // Validate the sort-copy pipeline
        if !matches!(
            pipeline_cache.get_compute_pipeline_state(self.get_sort_copy_pipeline_id()),
            CachedPipelineState::Ok(_)
        ) {
            return false;
        }

        true
    }

    #[inline]
    pub fn prepare_buffers(&mut self, render_device: &RenderDevice) {
        self.indirect_buffer.prepare_buffers(render_device);
    }

    #[inline]
    pub fn write_buffers(&self, command_encoder: &mut CommandEncoder) {
        self.indirect_buffer.write_buffers(command_encoder);
    }

    #[inline]
    pub fn clear_previous_frame_resizes(&mut self) {
        self.indirect_buffer.clear_previous_frame_resizes();
    }

    pub fn ensure_sort_fill_bind_group_layout_desc(
        &mut self,
        pipeline_cache: &PipelineCache,
        particle_layout: &ParticleLayout,
    ) -> Result<&BindGroupLayoutDescriptor, ()> {
        let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;
        let (layout, _) = self
            .sort_fill_bind_group_layout_descs
            .entry(key)
            .or_insert_with(|| {
                let alignment = self
                    .render_device
                    .limits()
                    .min_storage_buffer_offset_alignment;
                let bind_group_layout_desc = BindGroupLayoutDescriptor::new(
                    "hanabi:bgl:sort_fill",
                    &BindGroupLayoutEntries::sequential(
                        ShaderStages::COMPUTE,
                        (
                            // @group(0) @binding(0) var<storage, read_write> sort_buffer :
                            // SortBuffer;
                            storage_buffer::<GpuSortBufferSingleEntry>(false),
                            // @group(0) @binding(1) var<storage, read> particle_buffer :
                            // RawParticleBuffer;
                            storage_buffer_read_only_sized(
                                false,
                                Some(key.particle_min_binding_size.into()),
                            ),
                            // @group(0) @binding(2) var<storage, read> indirect_index_buffer :
                            // array<u32>;
                            storage_buffer_read_only::<GpuIndirectIndex>(false),
                            // @group(0) @binding(3) var<storage, read_write> effect_metadata :
                            // EffectMetadata;
                            storage_buffer_sized(
                                true,
                                Some(GpuEffectMetadata::aligned_size(alignment)),
                            ),
                            // @group(0) @binding(4) var<storage, read> spawner : Spawner;
                            storage_buffer_read_only_sized(
                                true,
                                Some(GpuSpawnerParams::aligned_size(alignment)),
                            ),
                        ),
                    ),
                );
                let pipeline_id =
                    pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
                        label: Some("hanabi:pipeline:sort_fill".into()),
                        layout: vec![bind_group_layout_desc.clone()],
                        shader: self.sort_fill_shader.clone(),
                        shader_defs: vec!["HAS_DUAL_KEY".into()],
                        entry_point: Some("main".into()),
                        push_constant_ranges: vec![],
                        zero_initialize_workgroup_memory: false,
                    });
                (bind_group_layout_desc, pipeline_id)
            });
        Ok(layout)
    }

    // We currently only use the bind group layout internally in
    // ensure_sort_fill_bind_group()
    #[allow(dead_code)]
    pub fn get_sort_fill_bind_group_layout_desc(
        &self,
        particle_layout: &ParticleLayout,
    ) -> Option<&BindGroupLayoutDescriptor> {
        let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;
        self.sort_fill_bind_group_layout_descs
            .get(&key)
            .map(|(layout, _)| layout)
    }

    pub fn get_sort_fill_pipeline_id(
        &self,
        particle_layout: &ParticleLayout,
    ) -> Option<CachedComputePipelineId> {
        let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;
        self.sort_fill_bind_group_layout_descs
            .get(&key)
            .map(|(_, pipeline_id)| *pipeline_id)
    }

    pub fn get_sort_copy_pipeline_id(&self) -> CachedComputePipelineId {
        self.sort_copy_pipeline_id
    }

    pub fn ensure_sort_fill_bind_group(
        &mut self,
        particle_layout: &ParticleLayout,
        particle: &Buffer,
        indirect_index: &Buffer,
        effect_metadata: &Buffer,
        spawner_buffer: &Buffer,
        pipeline_cache: &PipelineCache,
    ) -> Result<&BindGroup, ()> {
        let key = SortFillBindGroupKey {
            particle: particle.id(),
            indirect_index: indirect_index.id(),
            effect_metadata: effect_metadata.id(),
        };
        let entry = self.sort_fill_bind_groups.entry(key);
        let bind_group = match entry {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => {
                // Note: can't use get_bind_group_layout() because the function call mixes the
                // lifetimes of the two hash maps and complains the bind group one is already
                // borrowed. Doing a manual access to the layout one instead makes the compiler
                // happy.
                let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;
                let layout_desc = &self
                    .sort_fill_bind_group_layout_descs
                    .get(&key)
                    .ok_or(())?
                    .0;
                entry.insert(
                    self.render_device.create_bind_group(
                        "hanabi:bg:sort_fill",
                        &pipeline_cache.get_bind_group_layout(layout_desc),
                        &BindGroupEntries::sequential((
                            // @group(0) @binding(0) var<storage, read_write> pairs:
                            // array<KeyValuePair>;
                            self.sort_buffer.as_entire_binding(),
                            // @group(0) @binding(1) var<storage, read> particle_buffer:
                            // ParticleBuffer;
                            particle.as_entire_binding(),
                            // @group(0) @binding(2) var<storage, read> indirect_index_buffer :
                            // array<u32>;
                            indirect_index.as_entire_binding(),
                            // @group(0) @binding(3) var<storage, read> effect_metadata :
                            // EffectMetadata;
                            BufferBinding {
                                buffer: effect_metadata,
                                offset: 0,
                                size: Some(GpuEffectMetadata::aligned_size(
                                    self.render_device
                                        .limits()
                                        .min_storage_buffer_offset_alignment,
                                )),
                            },
                            // @group(0) @binding(4) var<storage, read> spawner : Spawner;
                            BufferBinding {
                                buffer: spawner_buffer,
                                offset: 0,
                                size: Some(GpuSpawnerParams::aligned_size(
                                    self.render_device
                                        .limits()
                                        .min_storage_buffer_offset_alignment,
                                )),
                            },
                        )),
                    ),
                )
            }
        };
        Ok(bind_group)
    }

    pub fn sort_fill_bind_group(
        &self,
        particle: BufferId,
        indirect_index: BufferId,
        effect_metadata: BufferId,
    ) -> Option<&BindGroup> {
        let key = SortFillBindGroupKey {
            particle,
            indirect_index,
            effect_metadata,
        };
        self.sort_fill_bind_groups.get(&key)
    }

    /// Ensure the bind group for the sort pass is created.
    pub fn ensure_sort_bind_group(
        &mut self,
        pipeline_cache: &PipelineCache,
    ) -> Result<&BindGroup, ()> {
        if self.sort_bind_group.is_none() {
            let sort_bind_group = self.render_device.create_bind_group(
                "hanabi:bg:sort",
                &pipeline_cache.get_bind_group_layout(&self.sort_bind_group_layout_desc),
                // @group(0) @binding(0) var<storage, read_write> pairs : array<KeyValuePair>;
                &BindGroupEntries::single(self.sort_buffer.as_entire_binding()),
            );
            self.sort_bind_group = Some(sort_bind_group);
        }
        Ok(self.sort_bind_group.as_ref().unwrap())
    }

    #[inline]
    pub fn sort_bind_group(&self) -> Option<&BindGroup> {
        self.sort_bind_group.as_ref()
    }

    pub fn ensure_sort_copy_bind_group(
        &mut self,
        indirect_index_buffer: &Buffer,
        effect_metadata_buffer: &Buffer,
        spawner_buffer: &Buffer,
        pipeline_cache: &PipelineCache,
    ) -> Result<&BindGroup, ()> {
        let key = SortCopyBindGroupKey {
            indirect_index: indirect_index_buffer.id(),
            sort: self.sort_buffer.id(),
            effect_metadata: effect_metadata_buffer.id(),
        };
        let entry = self.sort_copy_bind_groups.entry(key);
        let bind_group = match entry {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => {
                entry.insert(
                    self.render_device.create_bind_group(
                        "hanabi:bg:sort_copy",
                        &pipeline_cache
                            .get_bind_group_layout(&self.sort_copy_bind_group_layout_desc),
                        &BindGroupEntries::sequential((
                            // @group(0) @binding(0) var<storage, read_write> indirect_index_buffer
                            // : IndirectIndexBuffer;
                            indirect_index_buffer.as_entire_binding(),
                            // @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
                            self.sort_buffer.as_entire_binding(),
                            // @group(0) @binding(2) var<storage, read> effect_metadata :
                            // EffectMetadata;
                            BufferBinding {
                                buffer: effect_metadata_buffer,
                                offset: 0,
                                size: Some(GpuEffectMetadata::aligned_size(
                                    self.render_device
                                        .limits()
                                        .min_storage_buffer_offset_alignment,
                                )),
                            },
                            // @group(0) @binding(3) var<storage, read> spawner : Spawner;
                            BufferBinding {
                                buffer: spawner_buffer,
                                offset: 0,
                                size: Some(GpuSpawnerParams::aligned_size(
                                    self.render_device
                                        .limits()
                                        .min_storage_buffer_offset_alignment,
                                )),
                            },
                        )),
                    ),
                )
            }
        };
        Ok(bind_group)
    }

    pub fn sort_copy_bind_group(
        &self,
        indirect_index: BufferId,
        effect_metadata: BufferId,
    ) -> Option<&BindGroup> {
        let key = SortCopyBindGroupKey {
            indirect_index,
            sort: self.sort_buffer.id(),
            effect_metadata,
        };
        self.sort_copy_bind_groups.get(&key)
    }
}

1	use std::num::{NonZeroU32, NonZeroU64};
2
3	use bevy::{
4	asset::Handle,
5	ecs::{resource::Resource, world::World},
6	platform::collections::{hash_map::Entry, HashMap},
7	render::{
8	render_resource::{
9	binding_types::{
10	storage_buffer, storage_buffer_read_only, storage_buffer_read_only_sized,
11	storage_buffer_sized,
12	},
13	BindGroup, BindGroupEntries, BindGroupLayoutDescriptor, BindGroupLayoutEntries, Buffer,
14	BufferId, CachedComputePipelineId, CachedPipelineState, ComputePipelineDescriptor,
15	PipelineCache, ShaderType,
16	},
17	renderer::RenderDevice,
18	},
19	shader::Shader,
20	utils::default,
21	};
22	use bytemuck::{Pod, Zeroable};
23	use wgpu::{BufferBinding, BufferDescriptor, BufferUsages, CommandEncoder, ShaderStages};
24
25	use super::{gpu_buffer::GpuBuffer, GpuDispatchIndirectArgs, GpuEffectMetadata, StorageType};
26	use crate::{
27	render::{GpuIndirectIndex, GpuSpawnerParams},
28	Attribute, ParticleLayout,
29	};
30
31	#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32	struct SortFillBindGroupLayoutKey {
33	particle_min_binding_size: NonZeroU32,
34	particle_ribbon_id_offset: u32,
35	particle_age_offset: u32,
36	}
37
38	impl SortFillBindGroupLayoutKey {
39	pub fn from_particle_layout(particle_layout: &ParticleLayout) -> Result<Self, ()> {	×
40	let particle_ribbon_id_offset = particle_layout	×
41	.byte_offset(Attribute::RIBBON_ID)	×
42	.ok_or(())?;	×
43	let particle_age_offset = particle_layout.byte_offset(Attribute::AGE).ok_or(())?;	×
44	let key = SortFillBindGroupLayoutKey {
45	particle_min_binding_size: particle_layout.min_binding_size32(),	×
46	particle_ribbon_id_offset,
47	particle_age_offset,
48	};
49	Ok(key)	×
50	}
51	}
52
53	#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
54	struct SortFillBindGroupKey {
55	particle: BufferId,
56	indirect_index: BufferId,
57	effect_metadata: BufferId,
58	}
59
60	#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
61	struct SortCopyBindGroupKey {
62	indirect_index: BufferId,
63	sort: BufferId,
64	effect_metadata: BufferId,
65	}
66
67	/// GPU representation of a single dual-key value pair, with the added buffer
68	/// count as prefix. This is mainly used for shorcuts in bindings, not directly
69	/// as a type.
70	#[repr(C)]
71	#[derive(Debug, Copy, Clone, Pod, Zeroable, ShaderType)]
72	struct GpuSortBufferSingleEntry {
73	/// Number of key-value pairs to sort. This is the first element of the
74	/// entire buffer.
75	pub count: u32,
76	/// Key for the first entry.
77	pub key: u32,
78	/// Secondary key for the first entry.
79	pub key2: u32,
80	/// Value for the first entry.
81	pub value: u32,
82	}
83
84	#[derive(Resource)]
85	pub struct SortBindGroups {
86	/// Render device.
87	render_device: RenderDevice,
88	/// Sort-fill pass compute shader.
89	sort_fill_shader: Handle<Shader>,
90	/// GPU buffer of key-value pairs to sort.
91	sort_buffer: Buffer,
92	/// GPU buffer containing the [`GpuDispatchIndirect`] structs for the
93	/// sort-fill and sort passes.
94	indirect_buffer: GpuBuffer<GpuDispatchIndirectArgs>,
95	/// Bind group layouts for group #0 of the sort-fill compute pass.
96	sort_fill_bind_group_layout_descs:
97	HashMap<SortFillBindGroupLayoutKey, (BindGroupLayoutDescriptor, CachedComputePipelineId)>,
98	/// Bind groups for group #0 of the sort-fill compute pass.
99	sort_fill_bind_groups: HashMap<SortFillBindGroupKey, BindGroup>,
100	/// Bind group layout descriptor for group #0 of the sort compute pass.
101	sort_bind_group_layout_desc: BindGroupLayoutDescriptor,
102	/// Bind group for group #0 of the sort compute pass.
103	sort_bind_group: Option<BindGroup>,
104	sort_copy_bind_group_layout_desc: BindGroupLayoutDescriptor,
105	/// Pipeline for sort pass.
106	sort_pipeline_id: CachedComputePipelineId,
107	/// Pipeline for sort-copy pass.
108	sort_copy_pipeline_id: CachedComputePipelineId,
109	/// Bind groups for group #0 of the sort-copy compute pass.
110	sort_copy_bind_groups: HashMap<SortCopyBindGroupKey, BindGroup>,
111	}
112
113	impl SortBindGroups {
114	pub fn new(	3✔
115	world: &mut World,
116	sort_fill_shader: Handle<Shader>,
117	sort_shader: Handle<Shader>,
118	sort_copy_shader: Handle<Shader>,
119	) -> Self {
120	let render_device = world.resource::<RenderDevice>();	9✔
121	let pipeline_cache = world.resource::<PipelineCache>();	9✔
122
123	let sort_buffer = render_device.create_buffer(&BufferDescriptor {	12✔
124	label: Some("hanabi:buffer:sort:pairs"),	6✔
125	size: 3 * 1024 * 1024,	3✔
126	usage: BufferUsages::COPY_DST \| BufferUsages::STORAGE,	3✔
127	mapped_at_creation: false,	3✔
128	});
129
130	let indirect_buffer_size = 3 * 1024;	6✔
131	let indirect_buffer = render_device.create_buffer(&BufferDescriptor {	12✔
132	label: Some("hanabi:buffer:sort:indirect"),	6✔
133	size: indirect_buffer_size,	3✔
134	usage: BufferUsages::COPY_SRC	3✔
135	\| BufferUsages::COPY_DST	3✔
136	\| BufferUsages::STORAGE	3✔
137	\| BufferUsages::INDIRECT,	3✔
138	mapped_at_creation: false,	3✔
139	});
140	let indirect_buffer = GpuBuffer::new_allocated(
141	indirect_buffer,	3✔
142	indirect_buffer_size as u32,	3✔
143	Some("hanabi:buffer:sort:indirect".to_string()),	3✔
144	);
145
146	let sort_bind_group_layout_desc = BindGroupLayoutDescriptor::new(
147	"hanabi:bgl:sort",
148	&BindGroupLayoutEntries::single(	6✔
149	ShaderStages::COMPUTE,	3✔
150	storage_buffer_sized(false, Some(NonZeroU64::new(16).unwrap())), /* count + dual	6✔
151	* kv pair */	3✔
152	),
153	);
154
155	let sort_pipeline_id = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {	12✔
156	label: Some("hanabi:pipeline:sort".into()),	6✔
157	layout: vec![sort_bind_group_layout_desc.clone()],	12✔
158	shader: sort_shader,	6✔
159	shader_defs: vec!["HAS_DUAL_KEY".into()],	12✔
160	entry_point: Some("main".into()),	3✔
161	push_constant_ranges: vec![],	3✔
162	zero_initialize_workgroup_memory: false,	3✔
163	});
164
165	let alignment = render_device.limits().min_storage_buffer_offset_alignment;	6✔
166	let sort_copy_bind_group_layout_desc = BindGroupLayoutDescriptor::new(
167	"hanabi:bgl:sort_copy",
168	&BindGroupLayoutEntries::sequential(	6✔
169	ShaderStages::COMPUTE,	3✔
170	(
171	// @group(0) @binding(0) var<storage, read_write> indirect_index_buffer :
172	// IndirectIndexBuffer;
173	storage_buffer::<GpuIndirectIndex>(false),	6✔
174	// @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
175	storage_buffer_read_only::<GpuSortBufferSingleEntry>(false),	6✔
176	// @group(0) @binding(2) var<storage, read_write> effect_metadata :
177	// EffectMetadata;
178	storage_buffer_sized(true, Some(GpuEffectMetadata::aligned_size(alignment))),	9✔
179	// @group(0) @binding(3) var<storage, read> spawner : Spawner;
180	storage_buffer_read_only_sized(	6✔
181	true,	3✔
182	Some(GpuSpawnerParams::aligned_size(alignment)),	3✔
183	),
184	),
185	),
186	);
187
188	let sort_copy_pipeline_id =	3✔
189	pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {	9✔
190	label: Some("hanabi:pipeline:sort_copy".into()),	6✔
191	layout: vec![sort_copy_bind_group_layout_desc.clone()],	12✔
192	shader: sort_copy_shader,	6✔
193	shader_defs: vec![],	6✔
194	entry_point: Some("main".into()),	3✔
195	push_constant_ranges: vec![],	3✔
196	zero_initialize_workgroup_memory: false,	3✔
197	});
198
199	Self {
200	render_device: render_device.clone(),	9✔
201	sort_fill_shader,
202	sort_buffer,
203	indirect_buffer,
204	sort_fill_bind_group_layout_descs: default(),	6✔
205	sort_fill_bind_groups: default(),	6✔
206	sort_bind_group_layout_desc,
207	// This bind group is created later, once the pipeline and its bind group layouts are
208	// created by the pipeline cache. Technically we could create the bind group layout
209	// immediately because the PipelineCache pretends to but actually creates them
210	// on-the-fly in get_bind_group_layout(), but this is brittle as any behavior change
211	// would break Hanabi. Instead we create this bind group alongside all others, which is
212	// more consistent too.
213	sort_bind_group: None,
214	sort_copy_bind_group_layout_desc,
215	sort_pipeline_id,
216	sort_copy_pipeline_id,
217	sort_copy_bind_groups: default(),	3✔
218	}
219	}
220
221	#[inline]
222	pub fn clear_indirect_dispatch_buffer(&mut self) {	330✔
223	self.indirect_buffer.clear();	660✔
224	}
225
226	#[inline]
227	pub fn allocate_indirect_dispatch(&mut self) -> u32 {	×
228	self.indirect_buffer.allocate()	×
229	}
230
231	#[inline]
232	pub fn get_indirect_dispatch_byte_offset(&self, index: u32) -> u32 {	×
233	self.indirect_buffer.item_size() as u32 * index	×
234	}
235
236	#[inline]
237	#[allow(dead_code)]
238	pub fn sort_buffer(&self) -> &Buffer {	×
239	&self.sort_buffer	×
240	}
241
242	#[inline]
UNCOV 243	pub fn indirect_buffer(&self) -> Option<&Buffer> {	×
UNCOV 244	self.indirect_buffer.buffer()	×
245	}
246
247	#[inline]
248	pub fn sort_pipeline_id(&self) -> CachedComputePipelineId {	×
249	self.sort_pipeline_id	×
250	}
251
252	/// Check if the sort pipeline is ready to run for the given effect
253	/// instance.
254	///
255	/// This ensures all compute pipelines are compiled and ready to be used
256	/// this frame.
257	pub fn is_pipeline_ready(	×
258	&self,
259	particle_layout: &ParticleLayout,
260	pipeline_cache: &PipelineCache,
261	) -> bool {
262	// Validate the sort-fill pipeline. It was created and queued for compile by
263	// ensure_sort_fill_bind_group_layout(), which normally is called just before
264	// is_pipeline_ready().
265	let Some(pipeline_id) = self.get_sort_fill_pipeline_id(particle_layout) else {	×
266	return false;	×
267	};
268	if !matches!(	×
269	pipeline_cache.get_compute_pipeline_state(pipeline_id),	×
270	CachedPipelineState::Ok(_)
271	) {
272	return false;	×
273	}
274
275	// The 2 pipelines below are created and queued for compile in new(), so are
276	// almost always ready.
277	// FIXME - they could be checked once a frame only, not once per effect...
278
279	// Validate the sort pipeline
280	if !matches!(	×
281	pipeline_cache.get_compute_pipeline_state(self.sort_pipeline_id()),	×
282	CachedPipelineState::Ok(_)
283	) {
284	return false;	×
285	}
286
287	// Validate the sort-copy pipeline
288	if !matches!(	×
289	pipeline_cache.get_compute_pipeline_state(self.get_sort_copy_pipeline_id()),	×
290	CachedPipelineState::Ok(_)
291	) {
292	return false;	×
293	}
294
295	true	×
296	}
297
298	#[inline]
299	pub fn prepare_buffers(&mut self, render_device: &RenderDevice) {	330✔
300	self.indirect_buffer.prepare_buffers(render_device);	990✔
301	}
302
303	#[inline]
304	pub fn write_buffers(&self, command_encoder: &mut CommandEncoder) {	330✔
305	self.indirect_buffer.write_buffers(command_encoder);	990✔
306	}
307
308	#[inline]
309	pub fn clear_previous_frame_resizes(&mut self) {	330✔
310	self.indirect_buffer.clear_previous_frame_resizes();	660✔
311	}
312
313	pub fn ensure_sort_fill_bind_group_layout_desc(	×
314	&mut self,
315	pipeline_cache: &PipelineCache,
316	particle_layout: &ParticleLayout,
317	) -> Result<&BindGroupLayoutDescriptor, ()> {
318	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;	×
319	let (layout, _) = self	×
320	.sort_fill_bind_group_layout_descs	×
321	.entry(key)	×
322	.or_insert_with(\|\| {	×
323	let alignment = self	×
324	.render_device	×
325	.limits()	×
326	.min_storage_buffer_offset_alignment;	×
327	let bind_group_layout_desc = BindGroupLayoutDescriptor::new(	×
328	"hanabi:bgl:sort_fill",
329	&BindGroupLayoutEntries::sequential(	×
330	ShaderStages::COMPUTE,	×
331	(
332	// @group(0) @binding(0) var<storage, read_write> sort_buffer :
333	// SortBuffer;
334	storage_buffer::<GpuSortBufferSingleEntry>(false),	×
335	// @group(0) @binding(1) var<storage, read> particle_buffer :
336	// RawParticleBuffer;
337	storage_buffer_read_only_sized(	×
338	false,	×
339	Some(key.particle_min_binding_size.into()),	×
340	),
341	// @group(0) @binding(2) var<storage, read> indirect_index_buffer :
342	// array<u32>;
343	storage_buffer_read_only::<GpuIndirectIndex>(false),	×
344	// @group(0) @binding(3) var<storage, read_write> effect_metadata :
345	// EffectMetadata;
346	storage_buffer_sized(	×
347	true,	×
348	Some(GpuEffectMetadata::aligned_size(alignment)),	×
349	),
350	// @group(0) @binding(4) var<storage, read> spawner : Spawner;
351	storage_buffer_read_only_sized(	×
352	true,	×
353	Some(GpuSpawnerParams::aligned_size(alignment)),	×
354	),
355	),
356	),
357	);
358	let pipeline_id =	×
359	pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {	×
360	label: Some("hanabi:pipeline:sort_fill".into()),	×
361	layout: vec![bind_group_layout_desc.clone()],	×
362	shader: self.sort_fill_shader.clone(),	×
363	shader_defs: vec!["HAS_DUAL_KEY".into()],	×
364	entry_point: Some("main".into()),	×
365	push_constant_ranges: vec![],	×
366	zero_initialize_workgroup_memory: false,	×
367	});
368	(bind_group_layout_desc, pipeline_id)	×
369	});
370	Ok(layout)	×
371	}
372
373	// We currently only use the bind group layout internally in
374	// ensure_sort_fill_bind_group()
375	#[allow(dead_code)]
376	pub fn get_sort_fill_bind_group_layout_desc(	×
377	&self,
378	particle_layout: &ParticleLayout,
379	) -> Option<&BindGroupLayoutDescriptor> {
380	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;	×
381	self.sort_fill_bind_group_layout_descs	×
382	.get(&key)	×
383	.map(\|(layout, _)\| layout)	×
384	}
385
386	pub fn get_sort_fill_pipeline_id(	×
387	&self,
388	particle_layout: &ParticleLayout,
389	) -> Option<CachedComputePipelineId> {
390	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;	×
391	self.sort_fill_bind_group_layout_descs	×
392	.get(&key)	×
393	.map(\|(_, pipeline_id)\| *pipeline_id)	×
394	}
395
396	pub fn get_sort_copy_pipeline_id(&self) -> CachedComputePipelineId {	×
397	self.sort_copy_pipeline_id	×
398	}
399
400	pub fn ensure_sort_fill_bind_group(	×
401	&mut self,
402	particle_layout: &ParticleLayout,
403	particle: &Buffer,
404	indirect_index: &Buffer,
405	effect_metadata: &Buffer,
406	spawner_buffer: &Buffer,
407	pipeline_cache: &PipelineCache,
408	) -> Result<&BindGroup, ()> {
409	let key = SortFillBindGroupKey {
410	particle: particle.id(),	×
411	indirect_index: indirect_index.id(),	×
412	effect_metadata: effect_metadata.id(),	×
413	};
414	let entry = self.sort_fill_bind_groups.entry(key);	×
415	let bind_group = match entry {	×
416	Entry::Occupied(entry) => entry.into_mut(),	×
417	Entry::Vacant(entry) => {	×
418	// Note: can't use get_bind_group_layout() because the function call mixes the
419	// lifetimes of the two hash maps and complains the bind group one is already
420	// borrowed. Doing a manual access to the layout one instead makes the compiler
421	// happy.
422	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;	×
423	let layout_desc = &self	×
424	.sort_fill_bind_group_layout_descs	×
425	.get(&key)	×
426	.ok_or(())?	×
427	.0;
428	entry.insert(	×
429	self.render_device.create_bind_group(	×
430	"hanabi:bg:sort_fill",	×
431	&pipeline_cache.get_bind_group_layout(layout_desc),	×
432	&BindGroupEntries::sequential((	×
433	// @group(0) @binding(0) var<storage, read_write> pairs:
434	// array<KeyValuePair>;
435	self.sort_buffer.as_entire_binding(),	×
436	// @group(0) @binding(1) var<storage, read> particle_buffer:
437	// ParticleBuffer;
438	particle.as_entire_binding(),	×
439	// @group(0) @binding(2) var<storage, read> indirect_index_buffer :
440	// array<u32>;
441	indirect_index.as_entire_binding(),	×
442	// @group(0) @binding(3) var<storage, read> effect_metadata :
443	// EffectMetadata;
444	BufferBinding {	×
445	buffer: effect_metadata,	×
446	offset: 0,	×
447	size: Some(GpuEffectMetadata::aligned_size(	×
448	self.render_device	×
449	.limits()	×
450	.min_storage_buffer_offset_alignment,	×
451	)),
452	},
453	// @group(0) @binding(4) var<storage, read> spawner : Spawner;
454	BufferBinding {	×
455	buffer: spawner_buffer,	×
456	offset: 0,	×
457	size: Some(GpuSpawnerParams::aligned_size(	×
458	self.render_device	×
459	.limits()	×
460	.min_storage_buffer_offset_alignment,	×
461	)),
462	},
463	)),
464	),
465	)
466	}
467	};
468	Ok(bind_group)	×
469	}
470
471	pub fn sort_fill_bind_group(	×
472	&self,
473	particle: BufferId,
474	indirect_index: BufferId,
475	effect_metadata: BufferId,
476	) -> Option<&BindGroup> {
477	let key = SortFillBindGroupKey {
478	particle,
479	indirect_index,
480	effect_metadata,
481	};
482	self.sort_fill_bind_groups.get(&key)	×
483	}
484
485	/// Ensure the bind group for the sort pass is created.
486	pub fn ensure_sort_bind_group(	×
487	&mut self,
488	pipeline_cache: &PipelineCache,
489	) -> Result<&BindGroup, ()> {
490	if self.sort_bind_group.is_none() {	×
491	let sort_bind_group = self.render_device.create_bind_group(	×
492	"hanabi:bg:sort",
493	&pipeline_cache.get_bind_group_layout(&self.sort_bind_group_layout_desc),	×
494	// @group(0) @binding(0) var<storage, read_write> pairs : array<KeyValuePair>;
495	&BindGroupEntries::single(self.sort_buffer.as_entire_binding()),	×
496	);
497	self.sort_bind_group = Some(sort_bind_group);	×
498	}
499	Ok(self.sort_bind_group.as_ref().unwrap())	×
500	}
501
502	#[inline]
503	pub fn sort_bind_group(&self) -> Option<&BindGroup> {	×
504	self.sort_bind_group.as_ref()	×
505	}
506
507	pub fn ensure_sort_copy_bind_group(	×
508	&mut self,
509	indirect_index_buffer: &Buffer,
510	effect_metadata_buffer: &Buffer,
511	spawner_buffer: &Buffer,
512	pipeline_cache: &PipelineCache,
513	) -> Result<&BindGroup, ()> {
514	let key = SortCopyBindGroupKey {
515	indirect_index: indirect_index_buffer.id(),	×
516	sort: self.sort_buffer.id(),	×
517	effect_metadata: effect_metadata_buffer.id(),	×
518	};
519	let entry = self.sort_copy_bind_groups.entry(key);	×
520	let bind_group = match entry {	×
521	Entry::Occupied(entry) => entry.into_mut(),	×
522	Entry::Vacant(entry) => {	×
523	entry.insert(	×
524	self.render_device.create_bind_group(	×
525	"hanabi:bg:sort_copy",	×
526	&pipeline_cache	×
527	.get_bind_group_layout(&self.sort_copy_bind_group_layout_desc),	×
528	&BindGroupEntries::sequential((	×
529	// @group(0) @binding(0) var<storage, read_write> indirect_index_buffer
530	// : IndirectIndexBuffer;
531	indirect_index_buffer.as_entire_binding(),	×
532	// @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
533	self.sort_buffer.as_entire_binding(),	×
534	// @group(0) @binding(2) var<storage, read> effect_metadata :
535	// EffectMetadata;
536	BufferBinding {	×
537	buffer: effect_metadata_buffer,	×
538	offset: 0,	×
539	size: Some(GpuEffectMetadata::aligned_size(	×
540	self.render_device	×
541	.limits()	×
542	.min_storage_buffer_offset_alignment,	×
543	)),
544	},
545	// @group(0) @binding(3) var<storage, read> spawner : Spawner;
546	BufferBinding {	×
547	buffer: spawner_buffer,	×
548	offset: 0,	×
549	size: Some(GpuSpawnerParams::aligned_size(	×
550	self.render_device	×
551	.limits()	×
552	.min_storage_buffer_offset_alignment,	×
553	)),
554	},
555	)),
556	),
557	)
558	}
559	};
560	Ok(bind_group)	×
561	}
562
563	pub fn sort_copy_bind_group(	×
564	&self,
565	indirect_index: BufferId,
566	effect_metadata: BufferId,
567	) -> Option<&BindGroup> {
568	let key = SortCopyBindGroupKey {
569	indirect_index,
570	sort: self.sort_buffer.id(),	×
571	effect_metadata,
572	};
573	self.sort_copy_bind_groups.get(&key)	×
574	}
575	}

djeedai / bevy_hanabi / 22800469877

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous