bevy_shader_plugin/
gpu_buffer.rs

1use std::{
2    any::{TypeId, type_name},
3    collections::HashMap,
4    marker::PhantomData,
5};
6
7use bevy::{
8    asset::RenderAssetUsages,
9    prelude::*,
10    render::{
11        Render, RenderApp,
12        render_asset::RenderAssets,
13        render_resource::BufferUsages,
14        storage::{GpuShaderBuffer, ShaderBuffer},
15    },
16};
17use encase::{ShaderSize, internal::WriteInto};
18
19use crate::{
20    ShaderSystems,
21    auto_array::default_filled_bytes,
22    bindings::{AutoBufferKind, BindingTable, register_binding},
23};
24
25/// Main-world handle to a GPU-owned buffer registered with
26/// [`register_gpu_buffer`](crate::ShaderAppExt::register_gpu_buffer).
27///
28/// The buffer lives entirely on the GPU: it starts filled with `T::default()` and is
29/// never re-uploaded from the CPU, so whatever shaders write to it persists across frames.
30#[derive(Resource)]
31pub struct GpuBuffer<Tag> {
32    handle: Handle<ShaderBuffer>,
33    len: u32,
34    _marker: PhantomData<Tag>,
35}
36
37impl<Tag> GpuBuffer<Tag> {
38    /// The storage buffer asset backing this buffer.
39    pub fn handle(&self) -> &Handle<ShaderBuffer> {
40        &self.handle
41    }
42
43    /// Number of elements in the buffer.
44    pub fn len(&self) -> u32 {
45        self.len
46    }
47
48    /// `true` if the buffer has no elements.
49    pub fn is_empty(&self) -> bool {
50        self.len == 0
51    }
52}
53
54/// Element count of every registered GPU buffer, keyed by its tag type.
55/// Used to size [`Workgroups::over`](crate::Workgroups::over) dispatches.
56#[derive(Resource, Default)]
57pub(crate) struct GpuBufferLengths(pub HashMap<TypeId, u32>);
58
59/// Register a GPU-owned buffer. See
60/// [`ShaderAppExt::register_gpu_buffer`](crate::ShaderAppExt::register_gpu_buffer).
61pub(crate) fn register<Tag, T>(app: &mut App, group_index: u32, binding_index: u32, len: u32)
62where
63    Tag: Send + Sync + 'static,
64    T: ShaderSize + WriteInto + Default + Send + Sync + 'static,
65{
66    let mut storage = ShaderBuffer::new(
67        &default_filled_bytes::<T>(len as usize),
68        RenderAssetUsages::default(),
69    );
70    storage.buffer_description.label = Some("gpu_buffer");
71    // COPY_SRC lets readback (ours, or Bevy's `Readback::buffer`) copy out of this buffer.
72    storage.buffer_description.usage |= BufferUsages::COPY_DST | BufferUsages::COPY_SRC;
73
74    let handle = app
75        .world_mut()
76        .resource_mut::<Assets<ShaderBuffer>>()
77        .add(storage);
78    let asset_id = handle.id();
79
80    let previous = app
81        .world_mut()
82        .get_resource_or_insert_with(GpuBufferLengths::default)
83        .0
84        .insert(TypeId::of::<Tag>(), len);
85    assert!(
86        previous.is_none(),
87        "register_gpu_buffer: tag `{}` is already registered; use a distinct tag per buffer",
88        type_name::<Tag>()
89    );
90
91    app.insert_resource(GpuBuffer::<Tag> {
92        handle,
93        len,
94        _marker: PhantomData,
95    });
96
97    let render_app = app.sub_app_mut(RenderApp);
98    register_binding(
99        render_app,
100        group_index,
101        binding_index,
102        AutoBufferKind::Storage { read_only: false },
103    );
104
105    render_app.add_systems(
106        Render,
107        (move |buffers: Res<RenderAssets<GpuShaderBuffer>>, mut table: ResMut<BindingTable>| {
108            if let Some(gpu_buffer) = buffers.get(asset_id) {
109                table.set_buffer(group_index, binding_index, &gpu_buffer.buffer);
110            }
111        })
112        .in_set(ShaderSystems::PrepareBindings),
113    );
114}
bevy_shader_plugin/gpu_buffer.rs

bevy_shader_plugin/
gpu_buffer.rs
