shadow performance

crates/bevy_pbr/src/extended_material.rs

@@ -67,6 +67,14 @@ pub trait MaterialExtension: Asset + AsBindGroup + Clone + Sized {
         ShaderRef::Default
     }
 
+    /// Specifies the shadow batch key (see [`Material::shadow_material_key`]). By default, extended materials
+    /// do not batch for shadows.
+    /// If the extension doesn't modify vertices or apply discards based on extended material properties, you
+    /// may be able to improve shadow performance by returning the base material key.
+    fn shadow_material_key(&self, _base_key: Option<u64>) -> Option<u64> {
+        None
+    }
+
     /// Customizes the default [`RenderPipelineDescriptor`] for a specific entity using the entity's
     /// [`MaterialPipelineKey`] and [`MeshVertexBufferLayout`] as input.
     /// Specialization for the base material is applied before this function is called.
@@ -211,6 +219,11 @@ impl<B: Material, E: MaterialExtension> Material for ExtendedMaterial<B, E> {
         }
     }
 
+    fn shadow_material_key(&self) -> Option<u64> {
+        self.extension
+            .shadow_material_key(self.base.shadow_material_key())
+    }
+
     fn specialize(
         pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,

crates/bevy_pbr/src/material.rs

@@ -180,6 +180,13 @@ pub trait Material: Asset + AsBindGroup + Clone + Sized {
     ) -> Result<(), SpecializedMeshPipelineError> {
         Ok(())
     }
+
+    /// Specify a batch key to use for shadows. Returning Some(value) allows batching of distinct materials which will produce
+    /// the same depth-prepass output (such as opaque [`StandardMaterial`]s with different textures), rendering with the first
+    /// material of the batch for all meshes.
+    fn shadow_material_key(&self) -> Option<u64> {
+        None
+    }
 }
 
 /// Adds the necessary ECS resources and render logic to enable rendering entities using the given [`Material`]
@@ -773,16 +780,33 @@ pub struct MaterialProperties {
 pub struct PreparedMaterial<T: Material> {
     pub bindings: Vec<(u32, OwnedBindingResource)>,
     pub bind_group: BindGroup,
+    pub shadow_bind_group_id: MaterialBindGroupId,
     pub key: T::Data,
     pub properties: MaterialProperties,
 }
 
-#[derive(Component, Clone, Copy, Default, PartialEq, Eq, Deref, DerefMut)]
-pub struct MaterialBindGroupId(Option<BindGroupId>);
+#[derive(Component, Clone, Copy, Default)]
+pub enum MaterialBindGroupId {
+    #[default]
+    Unbatched,
+    Instance(BindGroupId),
+    Group(u64),
+}
+
+impl PartialEq for MaterialBindGroupId {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Instance(l0), Self::Instance(r0)) => l0 == r0,
+            (Self::Group(l0), Self::Group(r0)) => l0 == r0,
+            // unbatched should never report equal to prevent batching
+            _ => false,
+        }
+    }
+}
 
 impl<T: Material> PreparedMaterial<T> {
     pub fn get_bind_group_id(&self) -> MaterialBindGroupId {
-        MaterialBindGroupId(Some(self.bind_group.id()))
+        MaterialBindGroupId::Instance(self.bind_group.id())
     }
 }
 
@@ -937,9 +961,14 @@ fn prepare_material<M: Material>(
         OpaqueRendererMethod::Deferred => OpaqueRendererMethod::Deferred,
         OpaqueRendererMethod::Auto => default_opaque_render_method,
     };
+    let shadow_bind_group_id = match material.shadow_material_key() {
+        Some(key) => MaterialBindGroupId::Group(key),
+        None => MaterialBindGroupId::Instance(prepared.bind_group.id()),
+    };
     Ok(PreparedMaterial {
         bindings: prepared.bindings,
         bind_group: prepared.bind_group,
+        shadow_bind_group_id,
         key: prepared.data,
         properties: MaterialProperties {
             alpha_mode: material.alpha_mode(),

crates/bevy_pbr/src/pbr_material.rs

@@ -801,6 +801,11 @@ impl Material for StandardMaterial {
         PBR_SHADER_HANDLE.into()
     }
 
+    fn shadow_material_key(&self) -> Option<u64> {
+        // we can batch all pure opaque materials together for shadow rendering
+        (self.alpha_mode == AlphaMode::Opaque).then_some(0)
+    }
+
     fn specialize(
         _pipeline: &MaterialPipeline<Self>,
         descriptor: &mut RenderPipelineDescriptor,

crates/bevy_pbr/src/render/light.rs

@@ -1,3 +1,4 @@
+use bevy_asset::AssetId;
 use bevy_core_pipeline::core_3d::{Transparent3d, CORE_3D_DEPTH_FORMAT};
 use bevy_ecs::prelude::*;
 use bevy_math::{Mat4, UVec3, UVec4, Vec2, Vec3, Vec3Swizzles, Vec4, Vec4Swizzles};
@@ -1590,7 +1591,7 @@ pub fn queue_shadows<M: Material>(
     shadow_draw_functions: Res<DrawFunctions<Shadow>>,
     prepass_pipeline: Res<PrepassPipeline<M>>,
     render_meshes: Res<RenderAssets<Mesh>>,
-    render_mesh_instances: Res<RenderMeshInstances>,
+    mut render_mesh_instances: ResMut<RenderMeshInstances>,
     render_materials: Res<RenderMaterials<M>>,
     render_material_instances: Res<RenderMaterialInstances<M>>,
     mut pipelines: ResMut<SpecializedMeshPipelines<PrepassPipeline<M>>>,
@@ -1635,7 +1636,7 @@ pub fn queue_shadows<M: Material>(
             // NOTE: Lights with shadow mapping disabled will have no visible entities
             // so no meshes will be queued
             for entity in visible_entities.iter().copied() {
-                let Some(mesh_instance) = render_mesh_instances.get(&entity) else {
+                let Some(mesh_instance) = render_mesh_instances.get_mut(&entity) else {
                     continue;
                 };
                 if !mesh_instance.shadow_caster {
@@ -1685,11 +1686,13 @@ pub fn queue_shadows<M: Material>(
                     }
                 };
 
+                mesh_instance.shadow_material_bind_group_id = material.shadow_bind_group_id;
+
                 shadow_phase.add(Shadow {
                     draw_function: draw_shadow_mesh,
                     pipeline: pipeline_id,
                     entity,
-                    distance: 0.0, // TODO: sort front-to-back
+                    asset_id: mesh_instance.mesh_asset_id,
                     batch_range: 0..1,
                     dynamic_offset: None,
                 });
@@ -1699,7 +1702,7 @@ pub fn queue_shadows<M: Material>(
 }
 
 pub struct Shadow {
-    pub distance: f32,
+    pub asset_id: AssetId<Mesh>,
     pub entity: Entity,
     pub pipeline: CachedRenderPipelineId,
     pub draw_function: DrawFunctionId,
@@ -1708,7 +1711,7 @@ pub struct Shadow {
 }
 
 impl PhaseItem for Shadow {
-    type SortKey = usize;
+    type SortKey = (usize, AssetId<Mesh>);
 
     #[inline]
     fn entity(&self) -> Entity {
@@ -1717,7 +1720,7 @@ impl PhaseItem for Shadow {
 
     #[inline]
     fn sort_key(&self) -> Self::SortKey {
-        self.pipeline.id()
+        (self.pipeline.id(), self.asset_id)
     }
 
     #[inline]
@@ -1727,10 +1730,11 @@ impl PhaseItem for Shadow {
 
     #[inline]
     fn sort(items: &mut [Self]) {
-        // The shadow phase is sorted by pipeline id for performance reasons.
+        // The shadow phase is sorted by pipeline id and mesh id for performance reasons.
         // Grouping all draw commands using the same pipeline together performs
-        // better than rebinding everything at a high rate.
-        radsort::sort_by_key(items, |item| item.sort_key());
+        // better than rebinding everything at a high rate, and grouping by mesh
+        // allows batching to reduce draw calls.
+        items.sort_unstable_by_key(Self::sort_key);
     }
 
     #[inline]

crates/bevy_pbr/src/render/mesh.rs

@@ -145,7 +145,7 @@ impl Plugin for MeshRenderPlugin {
                             batch_and_prepare_render_phase::<Transmissive3d, MeshPipeline>,
                             batch_and_prepare_render_phase::<Transparent3d, MeshPipeline>,
                             batch_and_prepare_render_phase::<AlphaMask3d, MeshPipeline>,
-                            batch_and_prepare_render_phase::<Shadow, MeshPipeline>,
+                            batch_and_prepare_render_phase::<Shadow, ShadowMeshPipeline>,
                             batch_and_prepare_render_phase::<Opaque3dDeferred, MeshPipeline>,
                             batch_and_prepare_render_phase::<AlphaMask3dDeferred, MeshPipeline>,
                         )
@@ -257,6 +257,7 @@ pub struct RenderMeshInstance {
     pub transforms: MeshTransforms,
     pub mesh_asset_id: AssetId<Mesh>,
     pub material_bind_group_id: MaterialBindGroupId,
+    pub shadow_material_bind_group_id: MaterialBindGroupId,
     pub shadow_caster: bool,
     pub automatic_batching: bool,
 }
@@ -328,6 +329,7 @@ pub fn extract_meshes(
                     transforms,
                     shadow_caster: !not_shadow_caster,
                     material_bind_group_id: MaterialBindGroupId::default(),
+                    shadow_material_bind_group_id: MaterialBindGroupId::default(),
                     automatic_batching: !no_automatic_batching,
                 },
             ));
@@ -503,6 +505,36 @@ impl GetBatchData for MeshPipeline {
     }
 }
 
+pub struct ShadowMeshPipeline;
+impl GetBatchData for ShadowMeshPipeline {
+    type Param = (SRes<RenderMeshInstances>, SRes<RenderLightmaps>);
+    // The material bind group ID, the mesh ID, and the lightmap ID,
+    // respectively.
+    type CompareData = (MaterialBindGroupId, AssetId<Mesh>, Option<AssetId<Image>>);
+
+    type BufferData = MeshUniform;
+
+    fn get_batch_data(
+        (mesh_instances, lightmaps): &SystemParamItem<Self::Param>,
+        entity: Entity,
+    ) -> Option<(Self::BufferData, Option<Self::CompareData>)> {
+        let mesh_instance = mesh_instances.get(&entity)?;
+        let maybe_lightmap = lightmaps.render_lightmaps.get(&entity);
+
+        Some((
+            MeshUniform::new(
+                &mesh_instance.transforms,
+                maybe_lightmap.map(|lightmap| lightmap.uv_rect),
+            ),
+            mesh_instance.automatic_batching.then_some((
+                mesh_instance.shadow_material_bind_group_id,
+                mesh_instance.mesh_asset_id,
+                maybe_lightmap.map(|lightmap| lightmap.image),
+            )),
+        ))
+    }
+}
+
 bitflags::bitflags! {
     #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
     #[repr(transparent)]

examples/stress_tests/many_cubes.rs

@@ -41,6 +41,10 @@ struct Args {
     /// the number of different textures from which to randomly select the material base color. 0 means no textures.
     #[argh(option, default = "0")]
     material_texture_count: usize,
+
+    /// enable cascaded shadow map
+    #[argh(switch)]
+    shadows: bool,
 }
 
 #[derive(Default, Clone)]
@@ -185,7 +189,13 @@ fn setup(
         }
     }
 
-    commands.spawn(DirectionalLightBundle { ..default() });
+    commands.spawn(DirectionalLightBundle {
+        directional_light: DirectionalLight {
+            shadows_enabled: args.shadows,
+            ..default()
+        },
+        ..default()
+    });
 }
 
 fn init_textures(args: &Args, images: &mut Assets<Image>) -> Vec<Handle<Image>> {