Adds ShaderStorageBuffer asset

Cargo.toml

@@ -2415,6 +2415,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"

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<storage, read> colors: array<vec4<f32>, 5>;
+
+struct Vertex {
+    @builtin(instance_index) instance_index: u32,
+    @location(0) position: vec3<f32>,
+};
+
+struct VertexOutput {
+    @builtin(position) clip_position: vec4<f32>,
+    @location(0) world_position: vec4<f32>,
+    @location(1) color: vec4<f32>,
+};
+
+@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<f32> {
+    return mesh.color;
+}

crates/bevy_render/macros/src/as_bind_group.rs

@@ -212,13 +212,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                         visibility.hygienic_quote(&quote! { #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<TokenStream> {
                             )
                         });
                     } 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<TokenStream> {
                             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<TokenStream> {
             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<TokenStream> {
                 &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<Self::Data>, #render_path::render_resource::AsBindGroupError> {
                 let bindings = vec![#(#binding_impls,)*];
 

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::<RenderAssetBytesPerFrame>()

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<Image>,
 ///     #[storage(3, read_only)]
-///     values: Vec<f32>,
+///     storage_buffer: Handle<ShaderStorageBuffer>,
 ///     #[storage(4, read_only, buffer)]
-///     buffer: Buffer,
+///     raw_buffer: Buffer,
 ///     #[storage_texture(5)]
 ///     storage_texture: Handle<Image>,
 /// }
@@ -99,7 +101,8 @@ impl Deref for BindGroup {
 /// @group(2) @binding(0) var<uniform> color: vec4<f32>;
 /// @group(2) @binding(1) var color_texture: texture_2d<f32>;
 /// @group(2) @binding(2) var color_sampler: sampler;
-/// @group(2) @binding(3) var<storage> values: array<f32>;
+/// @group(2) @binding(3) var<storage> storage_buffer: array<f32>;
+/// @group(2) @binding(4) var<storage> raw_buffer: array<f32>;
 /// @group(2) @binding(5) var storage_texture: texture_storage_2d<rgba8unorm, read_write>;
 /// ```
 /// 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<Storage>`](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.
 /// ```

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::<GpuShaderStorageBuffer>::default())
+            .register_type::<ShaderStorageBuffer>()
+            .init_asset::<ShaderStorageBuffer>()
+            .register_asset_reflect::<ShaderStorageBuffer>();
+    }
+}
+
+/// 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<Vec<u8>>,
+    /// 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<RenderDevice>;
+
+    fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
+        source_asset.asset_usage
+    }
+
+    fn prepare_asset(
+        source_asset: Self::SourceAsset,
+        render_device: &mut SystemParamItem<Self::Param>,
+    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
+        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 })
+            }
+        }
+    }
+}

examples/README.md

@@ -403,6 +403,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