Storage Support in AsBindGroup

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"

assets/shaders/storage_buffer.wgsl

@@ -0,0 +1,19 @@
+// read-only storage buffer
+@group(1) @binding(0)
+var<storage, read> color_buffer: array<f32>;
+
+// read-write storage buffer
+@group(1) @binding(1)
+var<storage, read_write> writable_buffer: array<f32>;
+
+@fragment
+fn fragment(
+    #import bevy_pbr::mesh_vertex_output
+) -> @location(0) vec4<f32> {
+    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]);
+}

crates/bevy_render/macros/src/as_bind_group.rs

@@ -11,13 +11,15 @@ 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)]
 enum BindingType {
     Uniform,
     Texture,
     Sampler,
+    Storage,
 }
 
 #[derive(Clone)]
@@ -126,6 +128,8 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                 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<TokenStream> {
                         }
                     });
                 }
+                BindingType::Storage => {
+                    let StorageAttrs {
+                        read_only,
+                        visibility,
+                    } = get_storage_attrs(nested_meta_items)?;
+
+                    let visibility =
+                        visibility.hygenic_quote(&quote! { #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<SamplerBindingType
         )),
     }
 }
+
+const READ_ONLY: Symbol = Symbol("read_only");
+
+#[derive(Default)]
+struct StorageAttrs {
+    read_only: bool,
+    visibility: ShaderStageVisibility,
+}
+
+fn get_storage_attrs(metas: Vec<NestedMeta>) -> Result<StorageAttrs> {
+    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,
+    })
+}

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);
 

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 {

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
 

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::<WritableBuffer>::default())
+    .add_plugin(MaterialPlugin::<CustomMaterial>::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<Assets<Mesh>>,
+    mut materials: ResMut<Assets<CustomMaterial>>,
+    render_device: ResMut<RenderDevice>,
+) {
+    // backing storage buffer
+    let color_buffer = render_device.create_buffer(&BufferDescriptor {
+        label: None,
+        size: (std::mem::size_of::<f32>() * 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::<f32>() * 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<ColorBuffer>,
+    render_queue: Res<RenderQueue>,
+    time: Res<Time>,
+) {
+    let blueness = (time.elapsed_seconds() * 5.).sin() / 2.0 + 0.5;
+    render_queue.write_buffer(&color_buffer, 0, cast_slice(&[0.0, blueness, 0.0, 1.0]));
+}
+
+fn read_buffer_system(
+    writable_buffer: ResMut<WritableBuffer>,
+    render_device: ResMut<RenderDevice>,
+) {
+    let data = writable_buffer.slice(..);
+
+    render_device.map_buffer(&data, MapMode::Read, move |val| {
+        assert!(val.is_ok());
+    });
+
+    let device = render_device.wgpu_device();
+    device.poll(Maintain::Wait);
+
+    let data = data.get_mapped_range();
+    let slice = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, BUFFER_SIZE) };
+
+    info!("writable_buffer: {:?}", slice);
+
+    drop(data);
+    writable_buffer.unmap();
+}
+
+impl Material for CustomMaterial {
+    fn fragment_shader() -> ShaderRef {
+        "shaders/storage_buffer.wgsl".into()
+    }
+}
+
+#[derive(AsBindGroup, TypeUuid, Debug, Clone)]
+#[uuid = "b83c887a-ad95-4eb4-b150-60f95ef2028a"]
+pub struct CustomMaterial {
+    // read only storage buffer
+    #[storage(0, read_only)]
+    color: Buffer,
+
+    // writable storage buffer (default)
+    #[storage(1)]
+    writable_buffer: Buffer,
+}