Reimplement bindless storage buffers. (#17994)

Support for bindless storage buffers was temporarily removed with the
bindless revamp. This commit restores that support.

assets/shaders/bindless_material.wgsl

@@ -15,7 +15,7 @@ struct MaterialBindings {
 }
 
 #ifdef BINDLESS
-@group(2) @binding(0) var<storage> materials: binding_array<MaterialBindings>;
+@group(2) @binding(0) var<storage> materials: array<MaterialBindings>;
 @group(2) @binding(10) var<storage> material_color: binding_array<Color>;
 #else   // BINDLESS
 @group(2) @binding(0) var<uniform> material_color: Color;

crates/bevy_pbr/src/material_bind_groups.rs

@@ -401,6 +401,12 @@ where
     dirty: BufferDirtyState,
 }
 
+/// The size of the buffer that we assign to unused buffer slots, in bytes.
+///
+/// This is essentially arbitrary, as it doesn't seem to matter to `wgpu` what
+/// the size is.
+const DEFAULT_BINDLESS_FALLBACK_BUFFER_SIZE: u64 = 16;
+
 impl From<u32> for MaterialBindGroupSlot {
     fn from(value: u32) -> Self {
         MaterialBindGroupSlot(value)
@@ -699,7 +705,10 @@ where
                     bindless_buffer_descriptor.bindless_index,
                     render_device.create_buffer(&BufferDescriptor {
                         label: Some("bindless fallback buffer"),
-                        size: bindless_buffer_descriptor.size as u64,
+                        size: match bindless_buffer_descriptor.size {
+                            Some(size) => size as u64,
+                            None => DEFAULT_BINDLESS_FALLBACK_BUFFER_SIZE,
+                        },
                         usage: BufferUsages::STORAGE,
                         mapped_at_creation: false,
                     }),
@@ -1607,7 +1616,10 @@ where
                         bindless_index,
                         MaterialDataBuffer::new(
                             buffer_descriptor.binding_number,
-                            buffer_descriptor.size as u32,
+                            buffer_descriptor
+                                .size
+                                .expect("Data buffers should have a size")
+                                as u32,
                         ),
                     );
                 }

crates/bevy_render/macros/src/as_bind_group.rs

@@ -242,9 +242,12 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                         ),
                         bindless_index:
                             #render_path::render_resource::BindlessIndex(#binding_index),
-                        size: <#converted_shader_type as
+                                size: Some(
-                            #render_path::render_resource::ShaderType>::min_size().get() as
+                                    <
-                            usize,
+                                        #converted_shader_type as
+                                        #render_path::render_resource::ShaderType
+                                    >::min_size().get() as usize
+                                ),
                     }
                 });
 
@@ -375,15 +378,9 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                 }
 
                 BindingType::Storage => {
-                    if attr_bindless_count.is_some() {
-                        return Err(Error::new_spanned(
-                            attr,
-                            "Storage buffers are unsupported in bindless mode",
-                        ));
-                    }
-
                     let StorageAttrs {
                         visibility,
+                        binding_array: binding_array_binding,
                         read_only,
                         buffer,
                     } = get_storage_binding_attr(nested_meta_items)?;
@@ -413,8 +410,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                         });
                     }
 
-                    // TODO: Support bindless buffers that aren't
-                    // structure-level `#[uniform]` attributes.
                     non_bindless_binding_layouts.push(quote! {
                         #bind_group_layout_entries.push(
                             #render_path::render_resource::BindGroupLayoutEntry {
@@ -429,6 +424,54 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                             }
                         );
                     });
+
+                    if let Some(binding_array_binding) = binding_array_binding {
+                        // Add the storage buffer to the `BindlessResourceType` list
+                        // in the bindless descriptor.
+                        let bindless_resource_type = quote! {
+                            #render_path::render_resource::BindlessResourceType::Buffer
+                        };
+                        add_bindless_resource_type(
+                            &render_path,
+                            &mut bindless_resource_types,
+                            binding_index,
+                            bindless_resource_type,
+                        );
+
+                        // Push the buffer descriptor.
+                        bindless_buffer_descriptors.push(quote! {
+                            #render_path::render_resource::BindlessBufferDescriptor {
+                                // Note that, because this is bindless, *binding
+                                // index* here refers to the index in the bindless
+                                // index table (`bindless_index`), and the actual
+                                // binding number is the *binding array binding*.
+                                binding_number: #render_path::render_resource::BindingNumber(
+                                    #binding_array_binding
+                                ),
+                                bindless_index:
+                                    #render_path::render_resource::BindlessIndex(#binding_index),
+                                size: None,
+                            }
+                        });
+
+                        // Declare the binding array.
+                        bindless_binding_layouts.push(quote!{
+                            #bind_group_layout_entries.push(
+                                #render_path::render_resource::BindGroupLayoutEntry {
+                                    binding: #binding_array_binding,
+                                    visibility: #render_path::render_resource::ShaderStages::all(),
+                                    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: #actual_bindless_slot_count,
+                                }
+                            );
+                        });
+                    }
                 }
 
                 BindingType::StorageTexture => {
@@ -1270,6 +1313,14 @@ fn get_visibility_flag_value(meta_list: &MetaList) -> Result<ShaderStageVisibili
     Ok(ShaderStageVisibility::Flags(visibility))
 }
 
+// Returns the `binding_array(10)` part of a field-level declaration like
+// `#[storage(binding_array(10))]`.
+fn get_binding_array_flag_value(meta_list: &MetaList) -> Result<u32> {
+    meta_list
+        .parse_args_with(|input: ParseStream| input.parse::<LitInt>())?
+        .base10_parse()
+}
+
 #[derive(Clone, Copy, Default)]
 enum BindingTextureDimension {
     D1,
@@ -1610,6 +1661,7 @@ fn get_sampler_binding_type_value(lit_str: &LitStr) -> Result<SamplerBindingType
 #[derive(Default)]
 struct StorageAttrs {
     visibility: ShaderStageVisibility,
+    binding_array: Option<u32>,
     read_only: bool,
     buffer: bool,
 }
@@ -1619,6 +1671,7 @@ const BUFFER: Symbol = Symbol("buffer");
 
 fn get_storage_binding_attr(metas: Vec<Meta>) -> Result<StorageAttrs> {
     let mut visibility = ShaderStageVisibility::vertex_fragment();
+    let mut binding_array = None;
     let mut read_only = false;
     let mut buffer = false;
 
@@ -1629,6 +1682,10 @@ fn get_storage_binding_attr(metas: Vec<Meta>) -> Result<StorageAttrs> {
             List(m) if m.path == VISIBILITY => {
                 visibility = get_visibility_flag_value(&m)?;
             }
+            // Parse #[storage(0, binding_array(...))] for bindless mode.
+            List(m) if m.path == BINDING_ARRAY_MODIFIER_NAME => {
+                binding_array = Some(get_binding_array_flag_value(&m)?);
+            }
             Path(path) if path == READ_ONLY => {
                 read_only = true;
             }
@@ -1646,6 +1703,7 @@ fn get_storage_binding_attr(metas: Vec<Meta>) -> Result<StorageAttrs> {
 
     Ok(StorageAttrs {
         visibility,
+        binding_array,
         read_only,
         buffer,
     })

crates/bevy_render/src/render_resource/bind_group.rs

@@ -196,15 +196,19 @@ impl Deref for BindGroup {
 ///
 /// ## `storage(BINDING_INDEX, arguments)`
 ///
-/// * The field's [`Handle<Storage>`](bevy_asset::Handle) will be used to look up the matching [`Buffer`] GPU resource, which
+/// * The field's [`Handle<Storage>`](bevy_asset::Handle) will be used to look
-///   will be bound as a storage buffer in shaders. If the `storage` attribute is used, the field is expected a raw
+///   up the matching [`Buffer`] GPU resource, which will be bound as a storage
-///   buffer, and the buffer will be bound as a storage buffer in shaders.
+///   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.
+///   In bindless mode, `binding_array()` argument that specifies the binding
+///   number of the resulting storage buffer binding array must be present.
 ///
-/// | Arguments              | Values                                                                  | Default              |
+/// | Arguments              | Values                                                                  | Default                |
-/// |------------------------|-------------------------------------------------------------------------|----------------------|
+/// |------------------------|-------------------------------------------------------------------------|------------------------|
-/// | `visibility(...)`      | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` | `vertex`, `fragment` |
+/// | `visibility(...)`      | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` | `vertex`, `fragment`   |
-/// | `read_only`            | if present then value is true, otherwise false                          | `false`              |
+/// | `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       |                      |
+/// | `buffer`               | if present then the field will be assumed to be a raw wgpu buffer       |                        |
+/// | `binding_array(...)`   | the binding number of the binding array, for bindless mode              | bindless mode disabled |
 ///
 /// Note that fields without field-level binding attributes will be ignored.
 /// ```

crates/bevy_render/src/render_resource/bindless.rs

@@ -181,8 +181,8 @@ pub struct BindlessBufferDescriptor {
     /// `BINDLESS_INDEX` in `#[uniform(BINDLESS_INDEX, StandardMaterialUniform,
     /// bindless(BINDING_NUMBER)]`.
     pub bindless_index: BindlessIndex,
-    /// The size of the buffer in bytes.
+    /// The size of the buffer in bytes, if known.
-    pub size: usize,
+    pub size: Option<usize>,
 }
 
 /// The index of the actual binding in the bind group.