Add convenience methods for constructing and setting storage buffer data (#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
This commit is contained in:
charlotte 2024-09-09 08:28:31 -07:00 committed by GitHub
parent e939d6c33f
commit 5eca832cee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 39 additions and 21 deletions

View file

@ -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.
@ -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.

View file

@ -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";
@ -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 };
@ -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(),
);
}