diff --git a/Cargo.toml b/Cargo.toml index 6e9daa74a4fba..e8101a191cd44 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 0000000000000..c052411e3f198 --- /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 81fb920f3c444..9e4a4ab9fdc30 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 32dff45d1df7c..fc26a0c895431 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 4380890e17c27..cb3a6fb347631 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 0000000000000..4225ee7e28834 --- /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 e592f427dcdc5..f2ee223b37403 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 0000000000000..7661a103a7533 --- /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