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::prelude::ReflectDefault;
use bevy_reflect::Reflect; use bevy_reflect::Reflect;
use bevy_utils::default; use bevy_utils::default;
use encase::internal::WriteInto;
use encase::ShaderType;
use wgpu::util::BufferInitDescriptor; use wgpu::util::BufferInitDescriptor;
/// Adds [`ShaderStorageBuffer`] as an asset that is extracted and uploaded to the GPU. /// 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.asset_usage = asset_usage;
storage 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. /// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.

View file

@ -4,7 +4,6 @@ use bevy::{
reflect::TypePath, reflect::TypePath,
render::render_resource::{AsBindGroup, ShaderRef}, render::render_resource::{AsBindGroup, ShaderRef},
}; };
use bevy_render::render_asset::RenderAssetUsages;
use bevy_render::storage::ShaderStorageBuffer; use bevy_render::storage::ShaderStorageBuffer;
const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wgsl"; const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wgsl";
@ -33,10 +32,7 @@ fn setup(
[0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0],
]; ];
let colors = buffers.add(ShaderStorageBuffer::new( let colors = buffers.add(ShaderStorageBuffer::from(color_data));
bytemuck::cast_slice(color_data.as_slice()),
RenderAssetUsages::default(),
));
// Create the custom material with the storage buffer // Create the custom material with the storage buffer
let custom_material = CustomMaterial { colors }; let custom_material = CustomMaterial { colors };
@ -72,22 +68,19 @@ fn update(
) { ) {
let material = materials.get_mut(&material_handle.0).unwrap(); let material = materials.get_mut(&material_handle.0).unwrap();
let buffer = buffers.get_mut(&material.colors).unwrap(); let buffer = buffers.get_mut(&material.colors).unwrap();
buffer.data = Some( buffer.set_data(
bytemuck::cast_slice( (0..5)
(0..5) .map(|i| {
.map(|i| { let t = time.elapsed_seconds() * 5.0;
let t = time.elapsed_seconds() * 5.0; [
[ (t + i as f32).sin() / 2.0 + 0.5,
(t + i as f32).sin() / 2.0 + 0.5, (t + i as f32 + 2.0).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,
(t + i as f32 + 4.0).sin() / 2.0 + 0.5, 1.0,
1.0, ]
] })
}) .collect::<Vec<[f32; 4]>>()
.collect::<Vec<[f32; 4]>>() .as_slice(),
.as_slice(),
)
.to_vec(),
); );
} }