Adds ShaderStorageBuffer asset (#14663)

Adds a new `Handle<Storage>` 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<Storage>` asset type that is prepared into a
`GpuStorageBuffer` render asset. This asset can either be initialized
with a `Vec<u8>` 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<Storage>` asset instead. Usages of Vec` should be converted
into assets instead.

---------

Co-authored-by: IceSentry <IceSentry@users.noreply.github.com>

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"

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

@@ -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

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::<CustomMaterial>::default()))
+        .add_systems(Startup, setup)
+        .add_systems(Update, update)
+        .run();
+}
+
+/// set up a simple 3D scene
+fn setup(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
+    mut materials: ResMut<Assets<CustomMaterial>>,
+) {
+    // 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<Time>,
+    material_handle: Res<CustomMaterialHandle>,
+    mut materials: ResMut<Assets<CustomMaterial>>,
+    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
+) {
+    let material = materials.get_mut(&material_handle.0).unwrap();
+    let buffer = buffers.get_mut(&material.colors).unwrap();
+    buffer.data = Some(
+        bytemuck::cast_slice(
+            (0..5)
+                .map(|i| {
+                    let t = time.elapsed_seconds() * 5.0;
+                    [
+                        (t + i as f32).sin() / 2.0 + 0.5,
+                        (t + i as f32 + 2.0).sin() / 2.0 + 0.5,
+                        (t + i as f32 + 4.0).sin() / 2.0 + 0.5,
+                        1.0,
+                    ]
+                })
+                .collect::<Vec<[f32; 4]>>()
+                .as_slice(),
+        )
+        .to_vec(),
+    );
+}
+
+// Holds a handle to the custom material
+#[derive(Resource)]
+struct CustomMaterialHandle(Handle<CustomMaterial>);
+
+// This struct defines the data that will be passed to your shader
+#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
+struct CustomMaterial {
+    #[storage(0, read_only)]
+    colors: Handle<ShaderStorageBuffer>,
+}
+
+impl Material for CustomMaterial {
+    fn vertex_shader() -> ShaderRef {
+        SHADER_ASSET_PATH.into()
+    }
+
+    fn fragment_shader() -> ShaderRef {
+        SHADER_ASSET_PATH.into()
+    }
+}