diff --git a/Cargo.toml b/Cargo.toml index c6ec1f14d6bd0..a387231d69ec6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1237,6 +1237,16 @@ description = "A shader that renders a mesh multiple times in one draw call" category = "Shaders" wasm = true +[[example]] +name = "storage_buffer" +path = "examples/shader/storage_buffer.rs" + +[package.metadata.example.storage_buffer] +name = "Storage Buffer" +description = "A shader that uses storage buffer" +category = "Shaders" +wasm = false + [[example]] name = "animate_shader" path = "examples/shader/animate_shader.rs" diff --git a/assets/shaders/storage_buffer.wgsl b/assets/shaders/storage_buffer.wgsl new file mode 100644 index 0000000000000..62ee30768fbc5 --- /dev/null +++ b/assets/shaders/storage_buffer.wgsl @@ -0,0 +1,19 @@ +// read-only storage buffer +@group(1) @binding(0) +var color_buffer: array; + +// read-write storage buffer +@group(1) @binding(1) +var writable_buffer: array; + +@fragment +fn fragment( + #import bevy_pbr::mesh_vertex_output +) -> @location(0) vec4 { + writable_buffer[0] += 10.; + writable_buffer[1] += 5.; + writable_buffer[2] += 2.; + writable_buffer[3] += 1.; + + return vec4(color_buffer[0], color_buffer[1], color_buffer[2], color_buffer[3]); +} diff --git a/crates/bevy_render/macros/src/as_bind_group.rs b/crates/bevy_render/macros/src/as_bind_group.rs index 4c8a2b14b2ac9..85a0aec7a802d 100644 --- a/crates/bevy_render/macros/src/as_bind_group.rs +++ b/crates/bevy_render/macros/src/as_bind_group.rs @@ -11,6 +11,7 @@ use syn::{ const UNIFORM_ATTRIBUTE_NAME: Symbol = Symbol("uniform"); const TEXTURE_ATTRIBUTE_NAME: Symbol = Symbol("texture"); const SAMPLER_ATTRIBUTE_NAME: Symbol = Symbol("sampler"); +const STORAGE_ATTRIBUTE_NAME: Symbol = Symbol("storage"); const BIND_GROUP_DATA_ATTRIBUTE_NAME: Symbol = Symbol("bind_group_data"); #[derive(Copy, Clone, Debug)] @@ -18,6 +19,7 @@ enum BindingType { Uniform, Texture, Sampler, + Storage, } #[derive(Clone)] @@ -126,6 +128,8 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { BindingType::Texture } else if attr_ident == SAMPLER_ATTRIBUTE_NAME { BindingType::Sampler + } else if attr_ident == STORAGE_ATTRIBUTE_NAME { + BindingType::Storage } else { continue; }; @@ -256,6 +260,34 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result { } }); } + BindingType::Storage => { + let StorageAttrs { + read_only, + visibility, + } = get_storage_attrs(nested_meta_items)?; + + let visibility = + visibility.hygenic_quote("e! { #render_path::render_resource }); + + binding_impls.push(quote! { + #render_path::render_resource::OwnedBindingResource::Buffer({ + self.#field_name.clone() + }) + }); + + binding_layouts.push(quote! { + #render_path::render_resource::BindGroupLayoutEntry { + binding: #binding_index, + visibility: #visibility, + 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: None, + }, + count: None, + } + }); + } } } } @@ -861,3 +893,50 @@ fn get_sampler_binding_type_value(lit_str: &LitStr) -> Result) -> Result { + let mut visibility = ShaderStageVisibility::vertex_fragment(); + let mut read_only = false; + + for meta in metas { + use syn::{ + Meta::{List, NameValue}, + NestedMeta::Meta, + }; + match meta { + // Parse #[storage(0, visibility(...))]. + Meta(List(m)) if m.path == VISIBILITY => { + visibility = get_visibility_flag_value(&m.nested)?; + } + // Parse #[storage(0, read_only)]. + Meta(syn::Meta::Path(path)) if path == READ_ONLY => { + read_only = true; + } + Meta(NameValue(m)) => { + return Err(Error::new_spanned( + m.path, + "Not a valid name. Available attributes: `visibility, read_only`.", + )); + } + _ => { + return Err(Error::new_spanned( + meta, + "Not a name value pair: `foo = \"...\"`", + )); + } + } + } + + Ok(StorageAttrs { + read_only, + visibility, + }) +} diff --git a/crates/bevy_render/macros/src/lib.rs b/crates/bevy_render/macros/src/lib.rs index 863b48baf5fa0..13e3ee4387db1 100644 --- a/crates/bevy_render/macros/src/lib.rs +++ b/crates/bevy_render/macros/src/lib.rs @@ -17,7 +17,10 @@ pub fn derive_extract_resource(input: TokenStream) -> TokenStream { extract_resource::derive_extract_resource(input) } -#[proc_macro_derive(AsBindGroup, attributes(uniform, texture, sampler, bind_group_data))] +#[proc_macro_derive( + AsBindGroup, + attributes(uniform, texture, sampler, storage, bind_group_data) +)] pub fn derive_as_bind_group(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); diff --git a/crates/bevy_render/src/render_resource/mod.rs b/crates/bevy_render/src/render_resource/mod.rs index 33a010ff502d1..feccbf9863a5a 100644 --- a/crates/bevy_render/src/render_resource/mod.rs +++ b/crates/bevy_render/src/render_resource/mod.rs @@ -33,16 +33,16 @@ pub use wgpu::{ ComputePipelineDescriptor as RawComputePipelineDescriptor, DepthBiasState, DepthStencilState, Extent3d, Face, Features as WgpuFeatures, FilterMode, FragmentState as RawFragmentState, FrontFace, ImageCopyBuffer, ImageCopyBufferBase, ImageCopyTexture, ImageCopyTextureBase, - ImageDataLayout, ImageSubresourceRange, IndexFormat, Limits as WgpuLimits, LoadOp, MapMode, - MultisampleState, Operations, Origin3d, PipelineLayout, PipelineLayoutDescriptor, PolygonMode, - PrimitiveState, PrimitiveTopology, RenderPassColorAttachment, RenderPassDepthStencilAttachment, - RenderPassDescriptor, RenderPipelineDescriptor as RawRenderPipelineDescriptor, - SamplerBindingType, SamplerDescriptor, ShaderModule, ShaderModuleDescriptor, ShaderSource, - ShaderStages, StencilFaceState, StencilOperation, StencilState, StorageTextureAccess, - TextureAspect, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType, - TextureUsages, TextureViewDescriptor, TextureViewDimension, VertexAttribute, - VertexBufferLayout as RawVertexBufferLayout, VertexFormat, VertexState as RawVertexState, - VertexStepMode, + ImageDataLayout, ImageSubresourceRange, IndexFormat, Limits as WgpuLimits, LoadOp, Maintain, + MapMode, MultisampleState, Operations, Origin3d, PipelineLayout, PipelineLayoutDescriptor, + PolygonMode, PrimitiveState, PrimitiveTopology, RenderPassColorAttachment, + RenderPassDepthStencilAttachment, RenderPassDescriptor, + RenderPipelineDescriptor as RawRenderPipelineDescriptor, SamplerBindingType, SamplerDescriptor, + ShaderModule, ShaderModuleDescriptor, ShaderSource, ShaderStages, StencilFaceState, + StencilOperation, StencilState, StorageTextureAccess, TextureAspect, TextureDescriptor, + TextureDimension, TextureFormat, TextureSampleType, TextureUsages, TextureViewDescriptor, + TextureViewDimension, VertexAttribute, VertexBufferLayout as RawVertexBufferLayout, + VertexFormat, VertexState as RawVertexState, VertexStepMode, }; pub mod encase { diff --git a/examples/README.md b/examples/README.md index e354e6f283120..9d5968d6e9f7f 100644 --- a/examples/README.md +++ b/examples/README.md @@ -268,6 +268,7 @@ Example | Description [Material - Screenspace Texture](../examples/shader/shader_material_screenspace_texture.rs) | A shader that samples a texture with view-independent UV coordinates [Post Processing](../examples/shader/post_processing.rs) | A custom post processing effect, using two cameras, with one reusing the render texture of the first one [Shader Defs](../examples/shader/shader_defs.rs) | A shader that uses "shaders defs" (a bevy tool to selectively toggle parts of a shader) +[Storage Buffer](../examples/shader/storage_buffer.rs) | A shader that uses storage buffer ## Stress Tests diff --git a/examples/shader/storage_buffer.rs b/examples/shader/storage_buffer.rs new file mode 100644 index 0000000000000..8f435685955d4 --- /dev/null +++ b/examples/shader/storage_buffer.rs @@ -0,0 +1,137 @@ +//! A shader that uses storage buffer +use bevy::{ + prelude::*, + reflect::TypeUuid, + render::{ + extract_resource::{ExtractResource, ExtractResourcePlugin}, + render_resource::{ + AsBindGroup, Buffer, BufferDescriptor, BufferUsages, Maintain, MapMode, ShaderRef, + }, + renderer::{RenderDevice, RenderQueue}, + settings::{WgpuFeatures, WgpuSettings}, + RenderApp, RenderStage, + }, +}; +use bytemuck::cast_slice; + +const BUFFER_SIZE: usize = 4; + +fn main() { + let mut app = App::new(); + + app.insert_resource(WgpuSettings { + features: WgpuFeatures::MAPPABLE_PRIMARY_BUFFERS, + ..Default::default() + }) + .add_plugins(DefaultPlugins) + .add_plugin(ExtractResourcePlugin::::default()) + .add_plugin(MaterialPlugin::::default()) + .add_startup_system(setup) + .add_system(update_buffer_system); + + let render_app = app.sub_app_mut(RenderApp); + render_app.add_system_to_stage(RenderStage::Cleanup, read_buffer_system); + + app.run(); +} + +/// set up a simple 3D scene +fn setup( + mut commands: Commands, + mut meshes: ResMut>, + mut materials: ResMut>, + render_device: ResMut, +) { + // backing storage buffer + let color_buffer = render_device.create_buffer(&BufferDescriptor { + label: None, + size: (std::mem::size_of::() * BUFFER_SIZE) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + // writable buffer + let writable_buffer = render_device.create_buffer(&BufferDescriptor { + label: None, + size: (std::mem::size_of::() * BUFFER_SIZE) as u64, + usage: BufferUsages::STORAGE | BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + + // store as wrapped resource for later usage + commands.insert_resource(ColorBuffer(color_buffer.clone())); + commands.insert_resource(WritableBuffer(writable_buffer.clone())); + + // cube + commands.spawn(MaterialMeshBundle { + mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })), + transform: Transform::from_xyz(0.0, 0.5, 0.0), + material: materials.add(CustomMaterial { + color: color_buffer, + writable_buffer, + }), + ..default() + }); + + // camera + commands.spawn(Camera3dBundle { + transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y), + ..default() + }); +} + +#[derive(Deref, DerefMut, Resource)] +struct ColorBuffer(Buffer); + +#[derive(Deref, DerefMut, Resource, Clone, ExtractResource)] +struct WritableBuffer(Buffer); + +// updates the storage buffer +fn update_buffer_system( + color_buffer: ResMut, + render_queue: Res, + time: Res