28128423775

Committed 24 Jun 2026 08:44PM UTC coverage: 58.625% (-0.2%) from 58.797%

Build # 28128423775

Build Type

push

github

Committed by

web-flow

Commit Message

Upgrade to Bevy 0.19 (#545)

Coverage Stats

107 of 365 new or added lines in 6 files covered. (29.32%)

17 existing lines in 2 files now uncovered.

5244 of 8945 relevant lines covered (58.62%)

187.27 hits per line

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

27.6

/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, ShaderSize, 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,
    // Bound at binding 4 with a per-effect dynamic offset; reallocates as
    // effects grow, so it must be keyed or a stale buffer overruns the offset.
    spawner: BufferId,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct SortCopyBindGroupKey {
    indirect_index: BufferId,
    sort: BufferId,
    effect_metadata: BufferId,
    // Bound at binding 3 with a per-effect dynamic offset; reallocates as
    // effects grow, so it must be keyed or a stale buffer overruns the offset.
    spawner: 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,
        });

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

djeedai / bevy_hanabi / 28128423775

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