Move storage buffer specialization out of key and into "pipeline type"

Author: cart

crates/bevy_pbr/src/material.rs

@@ -39,7 +39,7 @@ use std::marker::PhantomData;
 /// way to render [`Mesh`] entities with custom shader logic. For materials that can specialize their [`RenderPipelineDescriptor`]
 /// based on specific material values, see [`SpecializedMaterial`]. [`Material`] automatically implements [`SpecializedMaterial`]
 /// and can be used anywhere that type is used (such as [`MaterialPlugin`]).
-pub trait Material: Asset + RenderAsset {
+pub trait Material: Asset + RenderAsset + Sized {
     /// Returns this material's [`BindGroup`]. This should match the layout returned by [`Material::bind_group_layout`].
     fn bind_group(material: &<Self as RenderAsset>::PreparedAsset) -> &BindGroup;
 
@@ -78,6 +78,7 @@ pub trait Material: Asset + RenderAsset {
     #[allow(unused_variables)]
     #[inline]
     fn specialize(
+        pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         layout: &MeshVertexBufferLayout,
     ) -> Result<(), SpecializedMeshPipelineError> {
@@ -89,19 +90,16 @@ impl<M: Material> SpecializedMaterial for M {
     type Key = ();
 
     #[inline]
-    fn key(
-        _render_device: &RenderDevice,
-        _material: &<Self as RenderAsset>::PreparedAsset,
-    ) -> Self::Key {
-    }
+    fn key(_material: &<Self as RenderAsset>::PreparedAsset) -> Self::Key {}
 
     #[inline]
     fn specialize(
+        pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         _key: Self::Key,
         layout: &MeshVertexBufferLayout,
     ) -> Result<(), SpecializedMeshPipelineError> {
-        <M as Material>::specialize(descriptor, layout)
+        <M as Material>::specialize(pipeline, descriptor, layout)
     }
 
     #[inline]
@@ -141,20 +139,18 @@ impl<M: Material> SpecializedMaterial for M {
 /// way to render [`Mesh`] entities with custom shader logic. [`SpecializedMaterials`](SpecializedMaterial) use their [`SpecializedMaterial::Key`]
 /// to customize their [`RenderPipelineDescriptor`] based on specific material values. The slightly simpler [`Material`] trait
 /// should be used for materials that do not need specialization. [`Material`] types automatically implement [`SpecializedMaterial`].
-pub trait SpecializedMaterial: Asset + RenderAsset {
+pub trait SpecializedMaterial: Asset + RenderAsset + Sized {
     /// The key used to specialize this material's [`RenderPipelineDescriptor`].
     type Key: PartialEq + Eq + Hash + Clone + Send + Sync;
 
     /// Extract the [`SpecializedMaterial::Key`] for the "prepared" version of this material. This key will be
     /// passed in to the [`SpecializedMaterial::specialize`] function when compiling the [`RenderPipeline`](bevy_render::render_resource::RenderPipeline)
     /// for a given entity's material.
-    fn key(
-        render_device: &RenderDevice,
-        material: &<Self as RenderAsset>::PreparedAsset,
-    ) -> Self::Key;
+    fn key(material: &<Self as RenderAsset>::PreparedAsset) -> Self::Key;
 
     /// Specializes the given `descriptor` according to the given `key`.
     fn specialize(
+        pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         key: Self::Key,
         layout: &MeshVertexBufferLayout,
@@ -258,7 +254,7 @@ impl<M: SpecializedMaterial> SpecializedMeshPipeline for MaterialPipeline<M> {
         let descriptor_layout = descriptor.layout.as_mut().unwrap();
         descriptor_layout.insert(1, self.material_layout.clone());
 
-        M::specialize(&mut descriptor, key.material_key, layout)?;
+        M::specialize(self, &mut descriptor, key.material_key, layout)?;
         Ok(descriptor)
     }
 }
@@ -315,7 +311,6 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
     material_pipeline: Res<MaterialPipeline<M>>,
     mut pipelines: ResMut<SpecializedMeshPipelines<MaterialPipeline<M>>>,
     mut pipeline_cache: ResMut<PipelineCache>,
-    render_device: Res<RenderDevice>,
     msaa: Res<Msaa>,
     render_meshes: Res<RenderAssets<Mesh>>,
     render_materials: Res<RenderAssets<M>>,
@@ -328,7 +323,6 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
         &mut RenderPhase<Transparent3d>,
     )>,
 ) {
-    let render_device = render_device.into_inner();
     for (view, visible_entities, mut opaque_phase, mut alpha_mask_phase, mut transparent_phase) in
         views.iter_mut()
     {
@@ -363,7 +357,7 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
                             mesh_key |= MeshPipelineKey::TRANSPARENT_MAIN_PASS;
                         }
 
-                        let material_key = M::key(render_device, material);
+                        let material_key = M::key(material);
 
                         let pipeline_id = pipelines.specialize(
                             &mut pipeline_cache,

crates/bevy_pbr/src/pbr_material.rs

@@ -1,7 +1,4 @@
-use crate::{
-    AlphaMode, MaterialPipeline, SpecializedMaterial, CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT,
-    PBR_SHADER_HANDLE,
-};
+use crate::{AlphaMode, MaterialPipeline, SpecializedMaterial, PBR_SHADER_HANDLE};
 use bevy_asset::{AssetServer, Handle};
 use bevy_ecs::system::{lifetimeless::SRes, SystemParamItem};
 use bevy_math::Vec4;
@@ -360,28 +357,20 @@ impl RenderAsset for StandardMaterial {
 pub struct StandardMaterialKey {
     normal_map: bool,
     cull_mode: Option<Face>,
-    use_storage_buffers: bool,
 }
 
 impl SpecializedMaterial for StandardMaterial {
     type Key = StandardMaterialKey;
 
-    fn key(
-        render_device: &RenderDevice,
-        render_asset: &<Self as RenderAsset>::PreparedAsset,
-    ) -> Self::Key {
+    fn key(render_asset: &<Self as RenderAsset>::PreparedAsset) -> Self::Key {
         StandardMaterialKey {
             normal_map: render_asset.has_normal_map,
             cull_mode: render_asset.cull_mode,
-            use_storage_buffers: matches!(
-                render_device
-                    .get_supported_read_only_binding_type(CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT),
-                BufferBindingType::Storage { .. }
-            ),
         }
     }
 
     fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         key: Self::Key,
         _layout: &MeshVertexBufferLayout,
@@ -395,14 +384,6 @@ impl SpecializedMaterial for StandardMaterial {
                 .push(String::from("STANDARDMATERIAL_NORMAL_MAP"));
         }
         descriptor.primitive.cull_mode = key.cull_mode;
-        if !key.use_storage_buffers {
-            descriptor
-                .fragment
-                .as_mut()
-                .unwrap()
-                .shader_defs
-                .push(String::from("NO_STORAGE_BUFFERS_SUPPORT"));
-        }
         if let Some(label) = &mut descriptor.label {
             *label = format!("pbr_{}", *label).into();
         }

crates/bevy_pbr/src/render/mesh.rs

@@ -560,6 +560,18 @@ impl SpecializedMeshPipeline for MeshPipeline {
             vertex_attributes.push(Mesh::ATTRIBUTE_TANGENT.at_shader_location(3));
         }
 
+        // TODO: consider exposing this in shaders in a more generally useful way, such as:
+        // # if AVAILABLE_STORAGE_BUFFER_BINDINGS == 3
+        // /* use storage buffers here */
+        // # elif
+        // /* use uniforms here */
+        if !matches!(
+            self.clustered_forward_buffer_binding_type,
+            BufferBindingType::Storage { .. }
+        ) {
+            shader_defs.push(String::from("NO_STORAGE_BUFFERS_SUPPORT"));
+        }
+
         let mut bind_group_layout = vec![self.view_layout.clone()];
         if layout.contains(Mesh::ATTRIBUTE_JOINT_INDEX)
             && layout.contains(Mesh::ATTRIBUTE_JOINT_WEIGHT)

examples/shader/custom_vertex_attribute.rs

@@ -137,6 +137,7 @@ impl Material for CustomMaterial {
     }
 
     fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         layout: &MeshVertexBufferLayout,
     ) -> Result<(), SpecializedMeshPipelineError> {

examples/shader/shader_material_glsl.rs

@@ -94,9 +94,10 @@ impl RenderAsset for CustomMaterial {
 impl SpecializedMaterial for CustomMaterial {
     type Key = ();
 
-    fn key(_: &RenderDevice, _: &<CustomMaterial as RenderAsset>::PreparedAsset) -> Self::Key {}
+    fn key(_: &<CustomMaterial as RenderAsset>::PreparedAsset) -> Self::Key {}
 
     fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,
         _: Self::Key,
         _layout: &MeshVertexBufferLayout,