From a4640046fcf1ce8e3e2bb94838aa520f0f99cfef Mon Sep 17 00:00:00 2001 From: charlotte Date: Mon, 2 Sep 2024 09:46:34 -0700 Subject: [PATCH] Adds `ShaderStorageBuffer` asset (#14663) Adds a new `Handle` asset type that can be used as a render asset, particularly for use with `AsBindGroup`. Closes: #13658 # Objective Allow users to create storage buffers in the main world without having to access the `RenderDevice`. While this resource is technically available, it's bad form to use in the main world and requires mixing rendering details with main world code. Additionally, this makes storage buffers easier to use with `AsBindGroup`, particularly in the following scenarios: - Sharing the same buffers between a compute stage and material shader. We already have examples of this for storage textures (see game of life example) and these changes allow a similar pattern to be used with storage buffers. - Preventing repeated gpu upload (see the previous easier to use `Vec` `AsBindGroup` option). - Allow initializing custom materials using `Default`. Previously, the lack of a `Default` implement for the raw `wgpu::Buffer` type made implementing a `AsBindGroup + Default` bound difficult in the presence of buffers. ## Solution Adds a new `Handle` asset type that is prepared into a `GpuStorageBuffer` render asset. This asset can either be initialized with a `Vec` of properly aligned data or with a size hint. Users can modify the underlying `wgpu::BufferDescriptor` to provide additional usage flags. ## Migration Guide The `AsBindGroup` `storage` attribute has been modified to reference the new `Handle` asset instead. Usages of Vec` should be converted into assets instead. --------- Co-authored-by: IceSentry --- Cargo.toml | 11 ++ assets/shaders/storage_buffer.wgsl | 38 ++++++ .../bevy_render/macros/src/as_bind_group.rs | 36 ++---- crates/bevy_render/src/lib.rs | 3 + .../src/render_resource/bind_group.rs | 15 ++- crates/bevy_render/src/storage.rs | 109 +++++++++++++++++ examples/README.md | 1 + examples/shader/storage_buffer.rs | 113 ++++++++++++++++++ 8 files changed, 297 insertions(+), 29 deletions(-) create mode 100644 assets/shaders/storage_buffer.wgsl create mode 100644 crates/bevy_render/src/storage.rs create mode 100644 examples/shader/storage_buffer.rs diff --git a/Cargo.toml b/Cargo.toml index 6e9daa74a4..e8101a191c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2417,6 +2417,17 @@ description = "A shader that shows how to bind and sample multiple textures as a category = "Shaders" wasm = false +[[example]] +name = "storage_buffer" +path = "examples/shader/storage_buffer.rs" +doc-scrape-examples = true + +[package.metadata.example.storage_buffer] +name = "Storage Buffer" +description = "A shader that shows how to bind a storage buffer using a custom material." +category = "Shaders" +wasm = true + [[example]] name = "specialized_mesh_pipeline" path = "examples/shader/specialized_mesh_pipeline.rs" diff --git a/assets/shaders/storage_buffer.wgsl b/assets/shaders/storage_buffer.wgsl new file mode 100644 index 0000000000..c052411e3f --- /dev/null +++ b/assets/shaders/storage_buffer.wgsl @@ -0,0 +1,38 @@ +#import bevy_pbr::{ + mesh_functions, + view_transformations::position_world_to_clip +} + +@group(2) @binding(0) var colors: array, 5>; + +struct Vertex { + @builtin(instance_index) instance_index: u32, + @location(0) position: vec3, +}; + +struct VertexOutput { + @builtin(position) clip_position: vec4, + @location(0) world_position: vec4, + @location(1) color: vec4, +}; + +@vertex +fn vertex(vertex: Vertex) -> VertexOutput { + var out: VertexOutput; + var world_from_local = mesh_functions::get_world_from_local(vertex.instance_index); + out.world_position = mesh_functions::mesh_position_local_to_world(world_from_local, vec4(vertex.position, 1.0)); + out.clip_position = position_world_to_clip(out.world_position.xyz); + + // We have 5 colors in the storage buffer, but potentially many instances of the mesh, so + // we use the instance index to select a color from the storage buffer. + out.color = colors[vertex.instance_index % 5]; + + return out; +} + +@fragment +fn fragment( + mesh: VertexOutput, +) -> @location(0) vec4 { + return mesh.color; +} \ No newline at end of file diff --git a/crates/bevy_render/macros/src/as_bind_group.rs b/crates/bevy_render/macros/src/as_bind_group.rs index 81fb920f3c..9e4a4ab9fd 100644 --- a/crates/bevy_render/macros/src/as_bind_group.rs +++ b/crates/bevy_render/macros/src/as_bind_group.rs @@ -212,13 +212,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { visibility.hygienic_quote("e! { #render_path::render_resource }); let field_name = field.ident.as_ref().unwrap(); - let field_ty = &field.ty; - - let min_binding_size = if buffer { - quote! {None} - } else { - quote! {Some(<#field_ty as #render_path::render_resource::ShaderType>::min_size())} - }; if buffer { binding_impls.push(quote! { @@ -230,21 +223,15 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { ) }); } else { - binding_impls.push(quote! {{ - use #render_path::render_resource::AsBindGroupShaderType; - let mut buffer = #render_path::render_resource::encase::StorageBuffer::new(Vec::new()); - buffer.write(&self.#field_name).unwrap(); - ( - #binding_index, - #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data( - &#render_path::render_resource::BufferInitDescriptor { - label: None, - usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::STORAGE, - contents: buffer.as_ref(), - }, - )) - ) - }}); + binding_impls.push(quote! { + ( + #binding_index, + #render_path::render_resource::OwnedBindingResource::Buffer({ + let handle: &#asset_path::Handle<#render_path::storage::ShaderStorageBuffer> = (&self.#field_name); + storage_buffers.get(handle).ok_or_else(|| #render_path::render_resource::AsBindGroupError::RetryNextUpdate)?.buffer.clone() + }) + ) + }); } binding_layouts.push(quote! { @@ -254,7 +241,7 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { ty: #render_path::render_resource::BindingType::Buffer { ty: #render_path::render_resource::BufferBindingType::Storage { read_only: #read_only }, has_dynamic_offset: false, - min_binding_size: #min_binding_size, + min_binding_size: None, }, count: None, } @@ -527,6 +514,7 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { type Param = ( #ecs_path::system::lifetimeless::SRes<#render_path::render_asset::RenderAssets<#render_path::texture::GpuImage>>, #ecs_path::system::lifetimeless::SRes<#render_path::texture::FallbackImage>, + #ecs_path::system::lifetimeless::SRes<#render_path::render_asset::RenderAssets<#render_path::storage::GpuShaderStorageBuffer>>, ); fn label() -> Option<&'static str> { @@ -537,7 +525,7 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { &self, layout: &#render_path::render_resource::BindGroupLayout, render_device: &#render_path::renderer::RenderDevice, - (images, fallback_image): &mut #ecs_path::system::SystemParamItem<'_, '_, Self::Param>, + (images, fallback_image, storage_buffers): &mut #ecs_path::system::SystemParamItem<'_, '_, Self::Param>, ) -> Result<#render_path::render_resource::UnpreparedBindGroup, #render_path::render_resource::AsBindGroupError> { let bindings = vec![#(#binding_impls,)*]; diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs index 32dff45d1d..fc26a0c895 100644 --- a/crates/bevy_render/src/lib.rs +++ b/crates/bevy_render/src/lib.rs @@ -35,6 +35,7 @@ pub mod render_resource; pub mod renderer; pub mod settings; mod spatial_bundle; +pub mod storage; pub mod texture; pub mod view; pub mod prelude { @@ -75,6 +76,7 @@ use crate::{ render_resource::{PipelineCache, Shader, ShaderLoader}, renderer::{render_system, RenderInstance}, settings::RenderCreation, + storage::StoragePlugin, view::{ViewPlugin, WindowRenderPlugin}, }; use bevy_app::{App, AppLabel, Plugin, SubApp}; @@ -356,6 +358,7 @@ impl Plugin for RenderPlugin { GlobalsPlugin, MorphPlugin, BatchingPlugin, + StoragePlugin, )); app.init_resource::() diff --git a/crates/bevy_render/src/render_resource/bind_group.rs b/crates/bevy_render/src/render_resource/bind_group.rs index 4380890e17..cb3a6fb347 100644 --- a/crates/bevy_render/src/render_resource/bind_group.rs +++ b/crates/bevy_render/src/render_resource/bind_group.rs @@ -77,6 +77,8 @@ impl Deref for BindGroup { /// # use bevy_render::{render_resource::*, texture::Image}; /// # use bevy_color::LinearRgba; /// # use bevy_asset::Handle; +/// # use bevy_render::storage::ShaderStorageBuffer; +/// /// #[derive(AsBindGroup)] /// struct CoolMaterial { /// #[uniform(0)] @@ -85,9 +87,9 @@ impl Deref for BindGroup { /// #[sampler(2)] /// color_texture: Handle, /// #[storage(3, read_only)] -/// values: Vec, +/// storage_buffer: Handle, /// #[storage(4, read_only, buffer)] -/// buffer: Buffer, +/// raw_buffer: Buffer, /// #[storage_texture(5)] /// storage_texture: Handle, /// } @@ -99,7 +101,8 @@ impl Deref for BindGroup { /// @group(2) @binding(0) var color: vec4; /// @group(2) @binding(1) var color_texture: texture_2d; /// @group(2) @binding(2) var color_sampler: sampler; -/// @group(2) @binding(3) var values: array; +/// @group(2) @binding(3) var storage_buffer: array; +/// @group(2) @binding(4) var raw_buffer: array; /// @group(2) @binding(5) var storage_texture: texture_storage_2d; /// ``` /// Note that the "group" index is determined by the usage context. It is not defined in [`AsBindGroup`]. For example, in Bevy material bind groups @@ -151,15 +154,17 @@ impl Deref for BindGroup { /// |------------------------|-------------------------------------------------------------------------|------------------------| /// | `sampler_type` = "..." | `"filtering"`, `"non_filtering"`, `"comparison"`. | `"filtering"` | /// | `visibility(...)` | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` | `vertex`, `fragment` | -/// /// * `storage(BINDING_INDEX, arguments)` -/// * The field will be converted to a shader-compatible type using the [`ShaderType`] trait, written to a [`Buffer`], and bound as a storage buffer. +/// * The field's [`Handle`](bevy_asset::Handle) will be used to look up the matching [`Buffer`] GPU resource, which +/// will be bound as a storage buffer in shaders. If the `storage` attribute is used, the field is expected a raw +/// buffer, and the buffer will be bound as a storage buffer in shaders. /// * It supports and optional `read_only` parameter. Defaults to false if not present. /// /// | Arguments | Values | Default | /// |------------------------|-------------------------------------------------------------------------|----------------------| /// | `visibility(...)` | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` | `vertex`, `fragment` | /// | `read_only` | if present then value is true, otherwise false | `false` | +/// | `buffer` | if present then the field will be assumed to be a raw wgpu buffer | | /// /// Note that fields without field-level binding attributes will be ignored. /// ``` diff --git a/crates/bevy_render/src/storage.rs b/crates/bevy_render/src/storage.rs new file mode 100644 index 0000000000..4225ee7e28 --- /dev/null +++ b/crates/bevy_render/src/storage.rs @@ -0,0 +1,109 @@ +use crate::render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssetUsages}; +use crate::render_resource::{Buffer, BufferUsages}; +use crate::renderer::RenderDevice; +use bevy_app::{App, Plugin}; +use bevy_asset::{Asset, AssetApp}; +use bevy_ecs::system::lifetimeless::SRes; +use bevy_ecs::system::SystemParamItem; +use bevy_reflect::prelude::ReflectDefault; +use bevy_reflect::Reflect; +use bevy_utils::default; +use wgpu::util::BufferInitDescriptor; + +/// Adds [`ShaderStorageBuffer`] as an asset that is extracted and uploaded to the GPU. +#[derive(Default)] +pub struct StoragePlugin; + +impl Plugin for StoragePlugin { + fn build(&self, app: &mut App) { + app.add_plugins(RenderAssetPlugin::::default()) + .register_type::() + .init_asset::() + .register_asset_reflect::(); + } +} + +/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU. +#[derive(Asset, Reflect, Debug, Clone)] +#[reflect_value(Default)] +pub struct ShaderStorageBuffer { + /// Optional data used to initialize the buffer. + pub data: Option>, + /// The buffer description used to create the buffer. + pub buffer_description: wgpu::BufferDescriptor<'static>, + /// The asset usage of the storage buffer. + pub asset_usage: RenderAssetUsages, +} + +impl Default for ShaderStorageBuffer { + fn default() -> Self { + Self { + data: None, + buffer_description: wgpu::BufferDescriptor { + label: None, + size: 0, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }, + asset_usage: RenderAssetUsages::default(), + } + } +} + +impl ShaderStorageBuffer { + /// Creates a new storage buffer with the given data and asset usage. + pub fn new(data: &[u8], asset_usage: RenderAssetUsages) -> Self { + let mut storage = ShaderStorageBuffer { + data: Some(data.to_vec()), + ..default() + }; + storage.asset_usage = asset_usage; + storage + } + + /// Creates a new storage buffer with the given size and asset usage. + pub fn with_size(size: usize, asset_usage: RenderAssetUsages) -> Self { + let mut storage = ShaderStorageBuffer { + data: None, + ..default() + }; + storage.buffer_description.size = size as u64; + storage.buffer_description.mapped_at_creation = false; + storage.asset_usage = asset_usage; + storage + } +} + +/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU. +pub struct GpuShaderStorageBuffer { + pub buffer: Buffer, +} + +impl RenderAsset for GpuShaderStorageBuffer { + type SourceAsset = ShaderStorageBuffer; + type Param = SRes; + + fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages { + source_asset.asset_usage + } + + fn prepare_asset( + source_asset: Self::SourceAsset, + render_device: &mut SystemParamItem, + ) -> Result> { + match source_asset.data { + Some(data) => { + let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor { + label: source_asset.buffer_description.label, + contents: &data, + usage: source_asset.buffer_description.usage, + }); + Ok(GpuShaderStorageBuffer { buffer }) + } + None => { + let buffer = render_device.create_buffer(&source_asset.buffer_description); + Ok(GpuShaderStorageBuffer { buffer }) + } + } + } +} diff --git a/examples/README.md b/examples/README.md index e592f427dc..f2ee223b37 100644 --- a/examples/README.md +++ b/examples/README.md @@ -404,6 +404,7 @@ Example | Description [Post Processing - Custom Render Pass](../examples/shader/custom_post_processing.rs) | A custom post processing effect, using a custom render pass that runs after the main pass [Shader Defs](../examples/shader/shader_defs.rs) | A shader that uses "shaders defs" (a bevy tool to selectively toggle parts of a shader) [Specialized Mesh Pipeline](../examples/shader/specialized_mesh_pipeline.rs) | Demonstrates how to write a specialized mesh pipeline +[Storage Buffer](../examples/shader/storage_buffer.rs) | A shader that shows how to bind a storage buffer using a custom material. [Texture Binding Array (Bindless Textures)](../examples/shader/texture_binding_array.rs) | A shader that shows how to bind and sample multiple textures as a binding array (a.k.a. bindless textures). ## State diff --git a/examples/shader/storage_buffer.rs b/examples/shader/storage_buffer.rs new file mode 100644 index 0000000000..7661a103a7 --- /dev/null +++ b/examples/shader/storage_buffer.rs @@ -0,0 +1,113 @@ +//! This example demonstrates how to use a storage buffer with `AsBindGroup` in a custom material. +use bevy::{ + prelude::*, + 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"; + +fn main() { + App::new() + .add_plugins((DefaultPlugins, MaterialPlugin::::default())) + .add_systems(Startup, setup) + .add_systems(Update, update) + .run(); +} + +/// set up a simple 3D scene +fn setup( + mut commands: Commands, + mut meshes: ResMut>, + mut buffers: ResMut>, + mut materials: ResMut>, +) { + // Example data for the storage buffer + let color_data: Vec<[f32; 4]> = vec![ + [1.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 1.0], + [1.0, 1.0, 0.0, 1.0], + [0.0, 1.0, 1.0, 1.0], + ]; + + let colors = buffers.add(ShaderStorageBuffer::new( + bytemuck::cast_slice(color_data.as_slice()), + RenderAssetUsages::default(), + )); + + // Create the custom material with the storage buffer + let custom_material = CustomMaterial { colors }; + + let material_handle = materials.add(custom_material); + commands.insert_resource(CustomMaterialHandle(material_handle.clone())); + + // Spawn cubes with the custom material + for i in -6..=6 { + for j in -3..=3 { + commands.spawn(MaterialMeshBundle { + mesh: meshes.add(Cuboid::from_size(Vec3::splat(0.3))), + material: material_handle.clone(), + transform: Transform::from_xyz(i as f32, j as f32, 0.0), + ..default() + }); + } + } + + // Camera + commands.spawn(Camera3dBundle { + transform: Transform::from_xyz(0.0, 0.0, 10.0).looking_at(Vec3::ZERO, Vec3::Y), + ..default() + }); +} + +// Update the material color by time +fn update( + time: Res