Skip to content

Commit

Permalink
Adds ShaderStorageBuffer asset (#14663)
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
tychedelia and IceSentry committed Sep 2, 2024
1 parent 3a8d559 commit a464004
Show file tree
Hide file tree
Showing 8 changed files with 297 additions and 29 deletions.
11 changes: 11 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
38 changes: 38 additions & 0 deletions assets/shaders/storage_buffer.wgsl
Original file line number Diff line number Diff line change
@@ -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;
}
36 changes: 12 additions & 24 deletions crates/bevy_render/macros/src/as_bind_group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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! {
Expand All @@ -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! {
Expand All @@ -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,
}
Expand Down Expand Up @@ -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> {
Expand All @@ -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,)*];

Expand Down
3 changes: 3 additions & 0 deletions crates/bevy_render/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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};
Expand Down Expand Up @@ -356,6 +358,7 @@ impl Plugin for RenderPlugin {
GlobalsPlugin,
MorphPlugin,
BatchingPlugin,
StoragePlugin,
));

app.init_resource::<RenderAssetBytesPerFrame>()
Expand Down
15 changes: 10 additions & 5 deletions crates/bevy_render/src/render_resource/bind_group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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>,
/// }
Expand All @@ -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
Expand Down Expand Up @@ -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.
/// ```
Expand Down
109 changes: 109 additions & 0 deletions crates/bevy_render/src/storage.rs
Original file line number Diff line number Diff line change
@@ -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 })
}
}
}
}
1 change: 1 addition & 0 deletions examples/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit a464004

Please sign in to comment.