Skip to content

Commit

Permalink
Add convenience methods for constructing and setting storage buffer d…
Browse files Browse the repository at this point in the history
…ata (#15044)

Adds some methods to assist in building `ShaderStorageBuffer` without
using `bytemuck`. We keep the `&[u8]` constructors since this is still
modeled as a thin wrapper around the buffer descriptor, but should make
it easier to interact with at the cost of an extra allocation in the
`ShaderType` path for the buffer writer.

Follow up from #14663
  • Loading branch information
tychedelia committed Sep 9, 2024
1 parent e939d6c commit 5eca832
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 21 deletions.
25 changes: 25 additions & 0 deletions crates/bevy_render/src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ use bevy_ecs::system::SystemParamItem;
use bevy_reflect::prelude::ReflectDefault;
use bevy_reflect::Reflect;
use bevy_utils::default;
use encase::internal::WriteInto;
use encase::ShaderType;
use wgpu::util::BufferInitDescriptor;

/// Adds [`ShaderStorageBuffer`] as an asset that is extracted and uploaded to the GPU.
Expand Down Expand Up @@ -72,6 +74,29 @@ impl ShaderStorageBuffer {
storage.asset_usage = asset_usage;
storage
}

/// Sets the data of the storage buffer to the given [`ShaderType`].
pub fn set_data<T>(&mut self, value: T)
where
T: ShaderType + WriteInto,
{
let size = value.size().get() as usize;
let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
wrapper.write(&value).unwrap();
self.data = Some(wrapper.into_inner());
}
}

impl<T> From<T> for ShaderStorageBuffer
where
T: ShaderType + WriteInto,
{
fn from(value: T) -> Self {
let size = value.size().get() as usize;
let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
wrapper.write(&value).unwrap();
Self::new(wrapper.as_ref(), RenderAssetUsages::default())
}
}

/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
Expand Down
35 changes: 14 additions & 21 deletions examples/shader/storage_buffer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ use bevy::{
reflect::TypePath,
render::render_resource::{AsBindGroup, ShaderRef},
};
use bevy_render::render_asset::RenderAssetUsages;
use bevy_render::storage::ShaderStorageBuffer;

const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wgsl";
Expand Down Expand Up @@ -33,10 +32,7 @@ fn setup(
[0.0, 1.0, 1.0, 1.0],
];

let colors = buffers.add(ShaderStorageBuffer::new(
bytemuck::cast_slice(color_data.as_slice()),
RenderAssetUsages::default(),
));
let colors = buffers.add(ShaderStorageBuffer::from(color_data));

// Create the custom material with the storage buffer
let custom_material = CustomMaterial { colors };
Expand Down Expand Up @@ -72,22 +68,19 @@ fn update(
) {
let material = materials.get_mut(&material_handle.0).unwrap();
let buffer = buffers.get_mut(&material.colors).unwrap();
buffer.data = Some(
bytemuck::cast_slice(
(0..5)
.map(|i| {
let t = time.elapsed_seconds() * 5.0;
[
(t + i as f32).sin() / 2.0 + 0.5,
(t + i as f32 + 2.0).sin() / 2.0 + 0.5,
(t + i as f32 + 4.0).sin() / 2.0 + 0.5,
1.0,
]
})
.collect::<Vec<[f32; 4]>>()
.as_slice(),
)
.to_vec(),
buffer.set_data(
(0..5)
.map(|i| {
let t = time.elapsed_seconds() * 5.0;
[
(t + i as f32).sin() / 2.0 + 0.5,
(t + i as f32 + 2.0).sin() / 2.0 + 0.5,
(t + i as f32 + 4.0).sin() / 2.0 + 0.5,
1.0,
]
})
.collect::<Vec<[f32; 4]>>()
.as_slice(),
);
}

Expand Down

0 comments on commit 5eca832

Please sign in to comment.