bevy_shader_plugin/
gpu_buffer.rs

1use std::{
2    any::{TypeId, type_name},
3    collections::HashMap,
4    marker::PhantomData,
5    num::NonZeroU64,
6};
7
8use bevy::{
9    asset::RenderAssetUsages,
10    prelude::*,
11    render::{
12        Render, RenderApp,
13        render_asset::RenderAssets,
14        render_resource::{BufferUsages, ShaderSize, encase::internal::WriteInto},
15        storage::{GpuShaderBuffer, ShaderBuffer},
16    },
17};
18
19use crate::{
20    ShaderSystems,
21    auto_array::{array_stride, 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    // One element: what a runtime-sized `array<T>` requires. Not the full length, since
99    // the asset behind `GpuBuffer::handle` can be replaced with a different-sized one.
100    let min_size = (len > 0)
101        .then(|| NonZeroU64::new(array_stride::<T>() as u64))
102        .flatten();
103    register_binding(
104        render_app,
105        group_index,
106        binding_index,
107        AutoBufferKind::Storage { read_only: false },
108        min_size,
109    );
110
111    render_app.add_systems(
112        Render,
113        (move |buffers: Res<RenderAssets<GpuShaderBuffer>>, mut table: ResMut<BindingTable>| {
114            if let Some(gpu_buffer) = buffers.get(asset_id) {
115                table.set_buffer(group_index, binding_index, &gpu_buffer.buffer);
116            }
117        })
118        .in_set(ShaderSystems::PrepareBindings),
119    );
120}
bevy_shader_plugin/gpu_buffer.rs

bevy_shader_plugin/
gpu_buffer.rs
