18260856969

Committed 05 Oct 2025 03:39PM UTC coverage: 66.606% (+0.03%) from 66.58%

Build # 18260856969

Build Type

push

github

Committed by

web-flow

Commit Message

Upgrade to Bevy v0.17 (#502)

Run Details

31 of 38 new or added lines in 11 files covered. (81.58%)

2 existing lines in 1 file now uncovered.

5120 of 7687 relevant lines covered (66.61%)

148.45 hits per line

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

41.99

/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::{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 effect_metadata_min_binding_size = GpuEffectMetadata::aligned_size(
            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: true,
                        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(effect_metadata_min_binding_size),
                    },
                    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> pairs: array<KeyValuePair>;
                        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: ParticleBuffer;
                        BindGroupLayoutEntry {
                            binding: 1,
                            visibility: ShaderStages::COMPUTE,
                            ty: BindingType::Buffer {
                                ty: BufferBindingType::Storage { read_only: true },
                                has_dynamic_offset: true,
                                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: true,
                                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,
                        },
                    ],
                );
                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,
    ) -> 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: BindingResource::Buffer(BufferBinding {
                                    buffer: &self.sort_buffer,
                                    offset: 0,
                                    size: None,
                                }),
                            },
                            // @group(0) @binding(1) var<storage, read> particle_buffer:
                            // ParticleBuffer;
                            BindGroupEntry {
                                binding: 1,
                                resource: BindingResource::Buffer(BufferBinding {
                                    buffer: particle,
                                    offset: 0,
                                    size: None,
                                }),
                            },
                            // @group(0) @binding(2) var<storage, read> indirect_index_buffer :
                            // array<u32>;
                            BindGroupEntry {
                                binding: 2,
                                resource: BindingResource::Buffer(BufferBinding {
                                    buffer: indirect_index,
                                    offset: 0,
                                    size: None,
                                }),
                            },
                            // @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,
                                    )),
                                }),
                            },
                        ],
                    ),
                )
            }
        };
        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,
    ) -> 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: BindingResource::Buffer(BufferBinding {
                                    buffer: indirect_index_buffer,
                                    offset: 0,
                                    size: None,
                                }),
                            },
                            // @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
                            BindGroupEntry {
                                binding: 1,
                                resource: BindingResource::Buffer(BufferBinding {
                                    buffer: &self.sort_buffer,
                                    offset: 0,
                                    size: None,
                                }),
                            },
                            // @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,
                                    )),
                                }),
                            },
                        ],
                    ),
                )
            }
        };
        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::{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 effect_metadata_min_binding_size = GpuEffectMetadata::aligned_size(
162	render_device.limits().min_storage_buffer_offset_alignment,	3✔
163	);
164	let sort_copy_bind_group_layout = render_device.create_bind_group_layout(	9✔
165	"hanabi:bind_group_layout:sort_copy",
166	&[	3✔
167	// @group(0) @binding(0) var<storage, read_write> indirect_index_buffer :
168	// IndirectIndexBuffer;
169	BindGroupLayoutEntry {	6✔
170	binding: 0,	6✔
171	visibility: ShaderStages::COMPUTE,	6✔
172	ty: BindingType::Buffer {	6✔
173	ty: BufferBindingType::Storage { read_only: false },	9✔
174	has_dynamic_offset: true,	6✔
175	min_binding_size: Some(NonZeroU64::new(12).unwrap()), // ping/pong+dead	6✔
176	},
177	count: None,	6✔
178	},
179	// @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
180	BindGroupLayoutEntry {	6✔
181	binding: 1,	6✔
182	visibility: ShaderStages::COMPUTE,	6✔
183	ty: BindingType::Buffer {	6✔
184	ty: BufferBindingType::Storage { read_only: true },	9✔
185	has_dynamic_offset: false,	6✔
186	min_binding_size: Some(NonZeroU64::new(16).unwrap()), /* count + dual kv	6✔
187	* pair */	6✔
188	},
189	count: None,	6✔
190	},
191	// @group(0) @binding(2) var<storage, read_write> effect_metadata : EffectMetadata;
192	BindGroupLayoutEntry {	3✔
193	binding: 2,	3✔
194	visibility: ShaderStages::COMPUTE,	3✔
195	ty: BindingType::Buffer {	3✔
196	ty: BufferBindingType::Storage { read_only: false },	3✔
197	has_dynamic_offset: true,	3✔
198	min_binding_size: Some(effect_metadata_min_binding_size),	3✔
199	},
200	count: None,	3✔
201	},
202	],
203	);
204
205	let sort_copy_pipeline_id =	3✔
206	pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {	9✔
207	label: Some("hanabi:pipeline:sort_copy".into()),	6✔
208	layout: vec![sort_copy_bind_group_layout.clone()],	12✔
209	shader: sort_copy_shader,	6✔
210	shader_defs: vec![],	6✔
211	entry_point: Some("main".into()),	3✔
212	push_constant_ranges: vec![],	3✔
213	zero_initialize_workgroup_memory: false,	3✔
214	});
215
216	Self {
217	render_device: render_device.clone(),	9✔
218	sort_fill_shader,
219	sort_buffer,
220	indirect_buffer,
221	sort_fill_bind_group_layouts: default(),	6✔
222	sort_fill_bind_groups: default(),	6✔
223	sort_bind_group,
224	sort_copy_bind_group_layout,
225	sort_pipeline_id,
226	sort_copy_pipeline_id,
227	sort_copy_bind_groups: default(),	3✔
228	}
229	}
230
231	#[inline]
232	pub fn clear_indirect_dispatch_buffer(&mut self) {	330✔
233	self.indirect_buffer.clear();	660✔
234	}
235
236	#[inline]
237	pub fn allocate_indirect_dispatch(&mut self) -> u32 {	×
238	self.indirect_buffer.allocate()	×
239	}
240
241	#[inline]
242	pub fn get_indirect_dispatch_byte_offset(&self, index: u32) -> u32 {	×
243	self.indirect_buffer.item_size() as u32 * index	×
244	}
245
246	#[inline]
247	#[allow(dead_code)]
248	pub fn sort_buffer(&self) -> &Buffer {	×
249	&self.sort_buffer	×
250	}
251
252	#[inline]
253	pub fn indirect_buffer(&self) -> Option<&Buffer> {	312✔
254	self.indirect_buffer.buffer()	624✔
255	}
256
257	#[inline]
258	pub fn sort_bind_group(&self) -> &BindGroup {	×
259	&self.sort_bind_group	×
260	}
261
262	#[inline]
263	pub fn sort_pipeline_id(&self) -> CachedComputePipelineId {	×
264	self.sort_pipeline_id	×
265	}
266
267	/// Check if the sort pipeline is ready to run for the given effect
268	/// instance.
269	///
270	/// This ensures all compute pipelines are compiled and ready to be used
271	/// this frame.
272	pub fn is_pipeline_ready(	×
273	&self,
274	particle_layout: &ParticleLayout,
275	pipeline_cache: &PipelineCache,
276	) -> bool {
277	// Validate the sort-fill pipeline. It was created and queued for compile by
278	// ensure_sort_fill_bind_group_layout(), which normally is called just before
279	// is_pipeline_ready().
280	let Some(pipeline_id) = self.get_sort_fill_pipeline_id(particle_layout) else {	×
281	return false;	×
282	};
283	if !matches!(	×
284	pipeline_cache.get_compute_pipeline_state(pipeline_id),
285	CachedPipelineState::Ok(_)
286	) {
287	return false;	×
288	}
289
290	// The 2 pipelines below are created and queued for compile in new(), so are
291	// almost always ready.
292	// FIXME - they could be checked once a frame only, not once per effect...
293
294	// Validate the sort pipeline
295	if !matches!(	×
296	pipeline_cache.get_compute_pipeline_state(self.sort_pipeline_id()),	×
297	CachedPipelineState::Ok(_)
298	) {
299	return false;	×
300	}
301
302	// Validate the sort-copy pipeline
303	if !matches!(	×
304	pipeline_cache.get_compute_pipeline_state(self.get_sort_copy_pipeline_id()),	×
305	CachedPipelineState::Ok(_)
306	) {
307	return false;	×
308	}
309
310	true	×
311	}
312
313	#[inline]
314	pub fn prepare_buffers(&mut self, render_device: &RenderDevice) {	330✔
315	self.indirect_buffer.prepare_buffers(render_device);	990✔
316	}
317
318	#[inline]
319	pub fn write_buffers(&self, command_encoder: &mut CommandEncoder) {	330✔
320	self.indirect_buffer.write_buffers(command_encoder);	990✔
321	}
322
323	#[inline]
324	pub fn clear_previous_frame_resizes(&mut self) {	330✔
325	self.indirect_buffer.clear_previous_frame_resizes();	660✔
326	}
327
328	pub fn ensure_sort_fill_bind_group_layout(	×
329	&mut self,
330	pipeline_cache: &PipelineCache,
331	particle_layout: &ParticleLayout,
332	) -> Result<&BindGroupLayout, ()> {
333	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;	×
334	let (layout, _) = self
335	.sort_fill_bind_group_layouts
336	.entry(key)
337	.or_insert_with(\|\| {	×
338	let alignment = self	×
339	.render_device	×
340	.limits()	×
341	.min_storage_buffer_offset_alignment;	×
342	let bind_group_layout = self.render_device.create_bind_group_layout(	×
343	"hanabi:bind_group_layout:sort_fill",
344	&[	×
345	// @group(0) @binding(0) var<storage, read_write> pairs: array<KeyValuePair>;
346	BindGroupLayoutEntry {	×
347	binding: 0,	×
348	visibility: ShaderStages::COMPUTE,	×
349	ty: BindingType::Buffer {	×
350	ty: BufferBindingType::Storage { read_only: false },	×
351	has_dynamic_offset: false,	×
352	min_binding_size: Some(NonZeroU64::new(16).unwrap()), // count + dual kv pair	×
353	},
354	count: None,	×
355	},
356	// @group(0) @binding(1) var<storage, read> particle_buffer: ParticleBuffer;
357	BindGroupLayoutEntry {	×
358	binding: 1,	×
359	visibility: ShaderStages::COMPUTE,	×
360	ty: BindingType::Buffer {	×
361	ty: BufferBindingType::Storage { read_only: true },	×
362	has_dynamic_offset: true,	×
363	min_binding_size: Some(key.particle_min_binding_size.into()),	×
364	},
365	count: None,	×
366	},
367	// @group(0) @binding(2) var<storage, read> indirect_index_buffer : array<u32>;
368	BindGroupLayoutEntry {	×
369	binding: 2,	×
370	visibility: ShaderStages::COMPUTE,	×
371	ty: BindingType::Buffer {	×
372	ty: BufferBindingType::Storage { read_only: true },	×
373	has_dynamic_offset: true,	×
374	min_binding_size: Some(NonZeroU64::new(12).unwrap()), // ping/pong+dead	×
375	},
376	count: None,	×
377	},
378	// @group(0) @binding(3) var<storage, read_write> effect_metadata : EffectMetadata;
379	BindGroupLayoutEntry {	×
380	binding: 3,	×
381	visibility: ShaderStages::COMPUTE,	×
382	ty: BindingType::Buffer {	×
383	ty: BufferBindingType::Storage { read_only: false },	×
384	has_dynamic_offset: true,	×
385	min_binding_size: Some(GpuEffectMetadata::aligned_size(alignment)),	×
386	},
387	count: None,	×
388	},
389	],
390	);
391	let pipeline_id =	×
392	pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {	×
393	label: Some("hanabi:pipeline:sort_fill".into()),	×
394	layout: vec![bind_group_layout.clone()],	×
395	shader: self.sort_fill_shader.clone(),	×
396	shader_defs: vec!["HAS_DUAL_KEY".into()],	×
NEW 397	entry_point: Some("main".into()),	×
398	push_constant_ranges: vec![],	×
399	zero_initialize_workgroup_memory: false,	×
400	});
401	(bind_group_layout, pipeline_id)	×
402	});
403	Ok(layout)
404	}
405
406	// We currently only use the bind group layout internally in
407	// ensure_sort_fill_bind_group()
408	#[allow(dead_code)]
409	pub fn get_sort_fill_bind_group_layout(	×
410	&self,
411	particle_layout: &ParticleLayout,
412	) -> Option<&BindGroupLayout> {
413	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;	×
414	self.sort_fill_bind_group_layouts
415	.get(&key)
416	.map(\|(layout, _)\| layout)
417	}
418
419	pub fn get_sort_fill_pipeline_id(	×
420	&self,
421	particle_layout: &ParticleLayout,
422	) -> Option<CachedComputePipelineId> {
423	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout).ok()?;	×
424	self.sort_fill_bind_group_layouts
425	.get(&key)
426	.map(\|(_, pipeline_id)\| *pipeline_id)
427	}
428
429	pub fn get_sort_copy_pipeline_id(&self) -> CachedComputePipelineId {	×
430	self.sort_copy_pipeline_id	×
431	}
432
433	pub fn ensure_sort_fill_bind_group(	×
434	&mut self,
435	particle_layout: &ParticleLayout,
436	particle: &Buffer,
437	indirect_index: &Buffer,
438	effect_metadata: &Buffer,
439	) -> Result<&BindGroup, ()> {
440	let key = SortFillBindGroupKey {
441	particle: particle.id(),	×
442	indirect_index: indirect_index.id(),	×
443	effect_metadata: effect_metadata.id(),	×
444	};
445	let entry = self.sort_fill_bind_groups.entry(key);	×
446	let bind_group = match entry {	×
447	Entry::Occupied(entry) => entry.into_mut(),	×
448	Entry::Vacant(entry) => {	×
449	// Note: can't use get_bind_group_layout() because the function call mixes the
450	// lifetimes of the two hash maps and complains the bind group one is already
451	// borrowed. Doing a manual access to the layout one instead makes the compiler
452	// happy.
453	let key = SortFillBindGroupLayoutKey::from_particle_layout(particle_layout)?;	×
454	let layout = &self.sort_fill_bind_group_layouts.get(&key).ok_or(())?.0;	×
455	entry.insert(
456	self.render_device.create_bind_group(
457	"hanabi:bind_group:sort_fill",
458	layout,
459	&[
460	// @group(0) @binding(0) var<storage, read_write> pairs:
461	// array<KeyValuePair>;
462	BindGroupEntry {
463	binding: 0,
464	resource: BindingResource::Buffer(BufferBinding {
465	buffer: &self.sort_buffer,
466	offset: 0,
467	size: None,
468	}),
469	},
470	// @group(0) @binding(1) var<storage, read> particle_buffer:
471	// ParticleBuffer;
472	BindGroupEntry {
473	binding: 1,
474	resource: BindingResource::Buffer(BufferBinding {
475	buffer: particle,
476	offset: 0,
477	size: None,
478	}),
479	},
480	// @group(0) @binding(2) var<storage, read> indirect_index_buffer :
481	// array<u32>;
482	BindGroupEntry {
483	binding: 2,
484	resource: BindingResource::Buffer(BufferBinding {
485	buffer: indirect_index,
486	offset: 0,
487	size: None,
488	}),
489	},
490	// @group(0) @binding(3) var<storage, read> effect_metadata :
491	// EffectMetadata;
492	BindGroupEntry {
493	binding: 3,
494	resource: BindingResource::Buffer(BufferBinding {
495	buffer: effect_metadata,
496	offset: 0,
497	size: Some(GpuEffectMetadata::aligned_size(
498	self.render_device
499	.limits()
500	.min_storage_buffer_offset_alignment,
501	)),
502	}),
503	},
504	],
505	),
506	)
507	}
508	};
509	Ok(bind_group)
510	}
511
512	pub fn sort_fill_bind_group(	×
513	&self,
514	particle: BufferId,
515	indirect_index: BufferId,
516	effect_metadata: BufferId,
517	) -> Option<&BindGroup> {
518	let key = SortFillBindGroupKey {
519	particle,
520	indirect_index,
521	effect_metadata,
522	};
523	self.sort_fill_bind_groups.get(&key)	×
524	}
525
526	pub fn ensure_sort_copy_bind_group(	×
527	&mut self,
528	indirect_index_buffer: &Buffer,
529	effect_metadata_buffer: &Buffer,
530	) -> Result<&BindGroup, ()> {
531	let key = SortCopyBindGroupKey {
532	indirect_index: indirect_index_buffer.id(),	×
533	sort: self.sort_buffer.id(),	×
534	effect_metadata: effect_metadata_buffer.id(),	×
535	};
536	let entry = self.sort_copy_bind_groups.entry(key);	×
537	let bind_group = match entry {	×
538	Entry::Occupied(entry) => entry.into_mut(),	×
539	Entry::Vacant(entry) => {	×
540	entry.insert(	×
541	self.render_device.create_bind_group(	×
542	"hanabi:bind_group:sort_copy",	×
543	&self.sort_copy_bind_group_layout,	×
544	&[	×
545	// @group(0) @binding(0) var<storage, read_write> indirect_index_buffer
546	// : IndirectIndexBuffer;
547	BindGroupEntry {	×
548	binding: 0,	×
549	resource: BindingResource::Buffer(BufferBinding {	×
550	buffer: indirect_index_buffer,	×
551	offset: 0,	×
552	size: None,	×
553	}),
554	},
555	// @group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
556	BindGroupEntry {	×
557	binding: 1,	×
558	resource: BindingResource::Buffer(BufferBinding {	×
559	buffer: &self.sort_buffer,	×
560	offset: 0,	×
561	size: None,	×
562	}),
563	},
564	// @group(0) @binding(2) var<storage, read> effect_metadata :
565	// EffectMetadata;
566	BindGroupEntry {	×
567	binding: 2,	×
568	resource: BindingResource::Buffer(BufferBinding {	×
569	buffer: effect_metadata_buffer,	×
570	offset: 0,	×
571	size: Some(GpuEffectMetadata::aligned_size(	×
572	self.render_device	×
573	.limits()	×
574	.min_storage_buffer_offset_alignment,	×
575	)),
576	}),
577	},
578	],
579	),
580	)
581	}
582	};
583	Ok(bind_group)	×
584	}
585
586	pub fn sort_copy_bind_group(	×
587	&self,
588	indirect_index: BufferId,
589	effect_metadata: BufferId,
590	) -> Option<&BindGroup> {
591	let key = SortCopyBindGroupKey {
592	indirect_index,
593	sort: self.sort_buffer.id(),	×
594	effect_metadata,
595	};
596	self.sort_copy_bind_groups.get(&key)	×
597	}
598	}

djeedai / bevy_hanabi / 18260856969

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