18604768490

Committed 17 Oct 2025 09:02PM UTC coverage: 66.442% (-0.1%) from 66.545%

Build # 18604768490

Build Type

Pull #508

github

Committed by

web-flow

Commit Message

Merge 78fa9c93e into ce96d01b7

Pull Request Pull Request #508: Pack multiple effects per particle slab

Run Details

25 of 54 new or added lines in 4 files covered. (46.3%)

7 existing lines in 1 file now uncovered.

5128 of 7718 relevant lines covered (66.44%)

214.74 hits per line

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

42.0

/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::{
            BindGroup, BindGroupLayout, Buffer, BufferId, CachedComputePipelineId,
            CachedPipelineState, ComputePipelineDescriptor, PipelineCache,
        },
        renderer::RenderDevice,
    },
    shader::Shader,
    utils::default,
};
use wgpu::{
    BindGroupEntry, BindGroupLayoutEntry, BindingResource, BindingType, BufferBinding,
    BufferBindingType, BufferDescriptor, BufferUsages, CommandEncoder, ShaderStages,
};

use super::{gpu_buffer::GpuBuffer, GpuDispatchIndirectArgs, GpuEffectMetadata, StorageType};
use crate::{render::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,
}

#[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_layouts:
        HashMap<SortFillBindGroupLayoutKey, (BindGroupLayout, CachedComputePipelineId)>,
    /// Bind groups for group #0 of the sort-fill compute pass.
    sort_fill_bind_groups: HashMap<SortFillBindGroupKey, BindGroup>,
    /// Bind group for group #0 of the sort compute pass.
    sort_bind_group: BindGroup,
    sort_copy_bind_group_layout: BindGroupLayout,
    /// 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 = render_device.create_bind_group_layout(
            "hanabi:bind_group_layout:sort",
            &[BindGroupLayoutEntry {
                binding: 0,
                visibility: ShaderStages::COMPUTE,
                ty: BindingType::Buffer {
                    ty: BufferBindingType::Storage { read_only: false },
                    has_dynamic_offset: false,
                    min_binding_size: Some(NonZeroU64::new(16).unwrap()), // count + dual kv pair
                },
                count: None,
            }],
        );

        let sort_bind_group = render_device.create_bind_group(
            "hanabi:bind_group:sort",
            &sort_bind_group_layout,
            &[
                // @group(0) @binding(0) var<storage, read_write> pairs : array<KeyValuePair>;
                BindGroupEntry {
                    binding: 0,
                    resource: BindingResource::Buffer(BufferBinding {
                        buffer: &sort_buffer,
                        offset: 0,
                        size: None,
                    }),
                },
            ],
        );

        let sort_pipeline_id = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
            label: Some("hanabi:pipeline:sort".into()),
            layout: vec![sort_bind_group_layout],
            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 = render_device.create_bind_group_layout(
            "hanabi:bind_group_layout:sort_copy",
            &[
                // @group(0) @binding(0) var<storage, read_write> indirect_index_buffer :
                // IndirectIndexBuffer;
                BindGroupLayoutEntry {
                    binding: 0,
                    visibility: ShaderStages::COMPUTE,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: false,
                        min_binding_size: Some(NonZeroU64::new(12).unwrap()), // ping/pong+dead
                    },
                    count: None,
                },
                // @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
                BindGroupLayoutEntry {
                    binding: 1,
                    visibility: ShaderStages::COMPUTE,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: Some(NonZeroU64::new(16).unwrap()), /* count + dual kv
                                                                               * pair */
                    },
                    count: None,
                },
                // @group(0) @binding(2) var<storage, read_write> effect_metadata : EffectMetadata;
                BindGroupLayoutEntry {
                    binding: 2,
                    visibility: ShaderStages::COMPUTE,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: true,
                        min_binding_size: Some(GpuEffectMetadata::aligned_size(alignment)),
                    },
                    count: None,
                },
                // @group(0) @binding(3) var<storage, read> spawner : Spawner;
                BindGroupLayoutEntry {
                    binding: 3,
                    visibility: ShaderStages::COMPUTE,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: true,
                        min_binding_size: Some(GpuSpawnerParams::aligned_size(alignment)),
                    },
                    count: None,
                },
            ],
        );

        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.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_layouts: default(),
            sort_fill_bind_groups: default(),
            sort_bind_group,
            sort_copy_bind_group_layout,
            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_bind_group(&self) -> &BindGroup {
        &self.sort_bind_group
    }

    #[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(
        &mut self,
        pipeline_cache: &PipelineCache,
        particle_layout: &ParticleLayout,
    ) -> Result<&BindGroupLayout, ()> {
        let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;
        let (layout, _) = self
            .sort_fill_bind_group_layouts
            .entry(key)
            .or_insert_with(|| {
                let alignment = self
                    .render_device
                    .limits()
                    .min_storage_buffer_offset_alignment;
                let bind_group_layout = self.render_device.create_bind_group_layout(
                    "hanabi:bind_group_layout:sort_fill",
                    &[
                        // @group(0) @binding(0) var<storage, read_write> sort_buffer : SortBuffer;
                        BindGroupLayoutEntry {
                            binding: 0,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: false },
                                has_dynamic_offset: false,
                                min_binding_size: Some(NonZeroU64::new(16).unwrap()), // count + dual kv pair
                            },
                            count: None,
                        },
                        // @group(0) @binding(1) var<storage, read> particle_buffer : RawParticleBuffer;
                        BindGroupLayoutEntry {
                            binding: 1,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: Some(key.particle_min_binding_size.into()),
                            },
                            count: None,
                        },
                        // @group(0) @binding(2) var<storage, read> indirect_index_buffer : array<u32>;
                        BindGroupLayoutEntry {
                            binding: 2,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: false,
                                min_binding_size: Some(NonZeroU64::new(12).unwrap()), // ping/pong+dead
                            },
                            count: None,
                        },
                        // @group(0) @binding(3) var<storage, read_write> effect_metadata : EffectMetadata;
                        BindGroupLayoutEntry {
                            binding: 3,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: false },
                                has_dynamic_offset: true,
                                min_binding_size: Some(GpuEffectMetadata::aligned_size(alignment)),
                            },
                            count: None,
                        },
                        // @group(0) @binding(4) var<storage, read> spawner : Spawner;
                        BindGroupLayoutEntry {
                            binding: 4,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: true,
                                min_binding_size: Some(GpuSpawnerParams::aligned_size(alignment)),
                            },
                            count: None,
                        },
                    ],
                );
                let pipeline_id =
                    pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
                        label: Some("hanabi:pipeline:sort_fill".into()),
                        layout: vec![bind_group_layout.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, 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(
        &self,
        particle_layout: &ParticleLayout,
    ) -> Option<&BindGroupLayout> {
        let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;
        self.sort_fill_bind_group_layouts
            .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_layouts
            .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,
    ) -> 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 = &self.sort_fill_bind_group_layouts.get(&key).ok_or(())?.0;
                entry.insert(
                    self.render_device.create_bind_group(
                        "hanabi:bind_group:sort_fill",
                        layout,
                        &[
                            // @group(0) @binding(0) var<storage, read_write> pairs:
                            // array<KeyValuePair>;
                            BindGroupEntry {
                                binding: 0,
                                resource: self.sort_buffer.as_entire_binding(),
                            },
                            // @group(0) @binding(1) var<storage, read> particle_buffer:
                            // ParticleBuffer;
                            BindGroupEntry {
                                binding: 1,
                                resource: particle.as_entire_binding(),
                            },
                            // @group(0) @binding(2) var<storage, read> indirect_index_buffer :
                            // array<u32>;
                            BindGroupEntry {
                                binding: 2,
                                resource: indirect_index.as_entire_binding(),
                            },
                            // @group(0) @binding(3) var<storage, read> effect_metadata :
                            // EffectMetadata;
                            BindGroupEntry {
                                binding: 3,
                                resource: BindingResource::Buffer(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;
                            BindGroupEntry {
                                binding: 4,
                                resource: BindingResource::Buffer(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)
    }

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

djeedai / bevy_hanabi / 18604768490

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