Update the prepass shaders and fix the batching logic for bindless and multidraw. (#16755)

This commit resolves most of the failures seen in #16670. It contains
two major fixes:

1. The prepass shaders weren't updated for bindless mode, so they were
accessing `material` as a single element instead of as an array. I added
the needed `BINDLESS` check.

2. If the mesh didn't support batch set keys (i.e. `get_batch_set_key()`
returns `None`), and multidraw was enabled, the batching logic would try
to multidraw all the meshes in a bin together instead of disabling
multidraw. This is because we checked whether the `Option<BatchSetKey>`
for the previous batch was equal to the `Option<BatchSetKey>` for the
next batch to determine whether objects could be multidrawn together,
which would return true if batch set keys were absent, causing an entire
bin to be multidrawn together. This patch fixes the logic so that
multidraw is only enabled if the batch set keys match *and are `Some`*.

Additionally, this commit adds batch key support for bins that use
`Opaque3dNoLightmapBinKey`, which in practice means prepasses.
Consequently, this patch enables multidraw for the prepass when GPU
culling is enabled.

When testing this patch, try adding `GpuCulling` to the camera in the
`deferred_rendering` and `ssr` examples. You can see that these examples
break without this patch and work properly with it.

---------

Co-authored-by: Alice Cecile <[email protected]>

Author: pcwalton

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -368,7 +368,7 @@ impl PhaseItem for AlphaMask3d {
 
     #[inline]
     fn draw_function(&self) -> DrawFunctionId {
-        self.key.draw_function
+        self.key.batch_set_key.draw_function
     }
 
     #[inline]
@@ -414,7 +414,7 @@ impl BinnedPhaseItem for AlphaMask3d {
 impl CachedRenderPipelinePhaseItem for AlphaMask3d {
     #[inline]
     fn cached_pipeline(&self) -> CachedRenderPipelineId {
-        self.key.pipeline
+        self.key.batch_set_key.pipeline
     }
 }
 

crates/bevy_core_pipeline/src/deferred/mod.rs

@@ -43,7 +43,7 @@ impl PhaseItem for Opaque3dDeferred {
 
     #[inline]
     fn draw_function(&self) -> DrawFunctionId {
-        self.key.draw_function
+        self.key.batch_set_key.draw_function
     }
 
     #[inline]
@@ -89,7 +89,7 @@ impl BinnedPhaseItem for Opaque3dDeferred {
 impl CachedRenderPipelinePhaseItem for Opaque3dDeferred {
     #[inline]
     fn cached_pipeline(&self) -> CachedRenderPipelineId {
-        self.key.pipeline
+        self.key.batch_set_key.pipeline
     }
 }
 
@@ -118,7 +118,7 @@ impl PhaseItem for AlphaMask3dDeferred {
 
     #[inline]
     fn draw_function(&self) -> DrawFunctionId {
-        self.key.draw_function
+        self.key.batch_set_key.draw_function
     }
 
     #[inline]
@@ -163,6 +163,6 @@ impl BinnedPhaseItem for AlphaMask3dDeferred {
 impl CachedRenderPipelinePhaseItem for AlphaMask3dDeferred {
     #[inline]
     fn cached_pipeline(&self) -> CachedRenderPipelineId {
-        self.key.pipeline
+        self.key.batch_set_key.pipeline
     }
 }

crates/bevy_core_pipeline/src/prepass/mod.rs

@@ -149,10 +149,13 @@ pub struct Opaque3dPrepass {
     pub extra_index: PhaseItemExtraIndex,
 }
 
-// TODO: Try interning these.
-/// The data used to bin each opaque 3D object in the prepass and deferred pass.
+/// Information that must be identical in order to place opaque meshes in the
+/// same *batch set* in the prepass and deferred pass.
+///
+/// A batch set is a set of batches that can be multi-drawn together, if
+/// multi-draw is in use.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct OpaqueNoLightmap3dBinKey {
+pub struct OpaqueNoLightmap3dBatchSetKey {
     /// The ID of the GPU pipeline.
     pub pipeline: CachedRenderPipelineId,
 
@@ -163,16 +166,27 @@ pub struct OpaqueNoLightmap3dBinKey {
     ///
     /// In the case of PBR, this is the `MaterialBindGroupIndex`.
     pub material_bind_group_index: Option<u32>,
+}
+
+// TODO: Try interning these.
+/// The data used to bin each opaque 3D object in the prepass and deferred pass.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct OpaqueNoLightmap3dBinKey {
+    /// The key of the *batch set*.
+    ///
+    /// As batches belong to a batch set, meshes in a batch must obviously be
+    /// able to be placed in a single batch set.
+    pub batch_set_key: OpaqueNoLightmap3dBatchSetKey,
 
     /// The ID of the asset.
     pub asset_id: UntypedAssetId,
 }
 
 impl PhaseItemBinKey for OpaqueNoLightmap3dBinKey {
-    type BatchSetKey = ();
+    type BatchSetKey = OpaqueNoLightmap3dBatchSetKey;
 
     fn get_batch_set_key(&self) -> Option<Self::BatchSetKey> {
-        None
+        Some(self.batch_set_key.clone())
     }
 }
 
@@ -188,7 +202,7 @@ impl PhaseItem for Opaque3dPrepass {
 
     #[inline]
     fn draw_function(&self) -> DrawFunctionId {
-        self.key.draw_function
+        self.key.batch_set_key.draw_function
     }
 
     #[inline]
@@ -234,7 +248,7 @@ impl BinnedPhaseItem for Opaque3dPrepass {
 impl CachedRenderPipelinePhaseItem for Opaque3dPrepass {
     #[inline]
     fn cached_pipeline(&self) -> CachedRenderPipelineId {
-        self.key.pipeline
+        self.key.batch_set_key.pipeline
     }
 }
 
@@ -262,7 +276,7 @@ impl PhaseItem for AlphaMask3dPrepass {
 
     #[inline]
     fn draw_function(&self) -> DrawFunctionId {
-        self.key.draw_function
+        self.key.batch_set_key.draw_function
     }
 
     #[inline]
@@ -308,7 +322,7 @@ impl BinnedPhaseItem for AlphaMask3dPrepass {
 impl CachedRenderPipelinePhaseItem for AlphaMask3dPrepass {
     #[inline]
     fn cached_pipeline(&self) -> CachedRenderPipelineId {
-        self.key.pipeline
+        self.key.batch_set_key.pipeline
     }
 }
 

crates/bevy_pbr/src/material.rs

@@ -14,7 +14,8 @@ use bevy_core_pipeline::{
     },
     oit::OrderIndependentTransparencySettings,
     prepass::{
-        DeferredPrepass, DepthPrepass, MotionVectorPrepass, NormalPrepass, OpaqueNoLightmap3dBinKey,
+        DeferredPrepass, DepthPrepass, MotionVectorPrepass, NormalPrepass,
+        OpaqueNoLightmap3dBatchSetKey, OpaqueNoLightmap3dBinKey,
     },
     tonemapping::{DebandDither, Tonemapping},
 };
@@ -906,10 +907,12 @@ pub fn queue_material_meshes<M: Material>(
                         });
                     } else if material.properties.render_method == OpaqueRendererMethod::Forward {
                         let bin_key = OpaqueNoLightmap3dBinKey {
-                            draw_function: draw_alpha_mask_pbr,
-                            pipeline: pipeline_id,
+                            batch_set_key: OpaqueNoLightmap3dBatchSetKey {
+                                draw_function: draw_alpha_mask_pbr,
+                                pipeline: pipeline_id,
+                                material_bind_group_index: Some(material.binding.group.0),
+                            },
                             asset_id: mesh_instance.mesh_asset_id.into(),
-                            material_bind_group_index: Some(material.binding.group.0),
                         };
                         alpha_mask_phase.add(
                             bin_key,

crates/bevy_pbr/src/prepass/mod.rs

@@ -951,21 +951,25 @@ pub fn queue_prepass_material_meshes<M: Material>(
                     if deferred {
                         opaque_deferred_phase.as_mut().unwrap().add(
                             OpaqueNoLightmap3dBinKey {
-                                draw_function: opaque_draw_deferred,
-                                pipeline: pipeline_id,
+                                batch_set_key: OpaqueNoLightmap3dBatchSetKey {
+                                    draw_function: opaque_draw_deferred,
+                                    pipeline: pipeline_id,
+                                    material_bind_group_index: Some(material.binding.group.0),
+                                },
                                 asset_id: mesh_instance.mesh_asset_id.into(),
-                                material_bind_group_index: Some(material.binding.group.0),
                             },
                             (*render_entity, *visible_entity),
                             BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                         );
                     } else if let Some(opaque_phase) = opaque_phase.as_mut() {
                         opaque_phase.add(
                             OpaqueNoLightmap3dBinKey {
-                                draw_function: opaque_draw_prepass,
-                                pipeline: pipeline_id,
+                                batch_set_key: OpaqueNoLightmap3dBatchSetKey {
+                                    draw_function: opaque_draw_prepass,
+                                    pipeline: pipeline_id,
+                                    material_bind_group_index: Some(material.binding.group.0),
+                                },
                                 asset_id: mesh_instance.mesh_asset_id.into(),
-                                material_bind_group_index: Some(material.binding.group.0),
                             },
                             (*render_entity, *visible_entity),
                             BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
@@ -976,10 +980,12 @@ pub fn queue_prepass_material_meshes<M: Material>(
                 MeshPipelineKey::MAY_DISCARD => {
                     if deferred {
                         let bin_key = OpaqueNoLightmap3dBinKey {
-                            pipeline: pipeline_id,
-                            draw_function: alpha_mask_draw_deferred,
+                            batch_set_key: OpaqueNoLightmap3dBatchSetKey {
+                                draw_function: alpha_mask_draw_deferred,
+                                pipeline: pipeline_id,
+                                material_bind_group_index: Some(material.binding.group.0),
+                            },
                             asset_id: mesh_instance.mesh_asset_id.into(),
-                            material_bind_group_index: Some(material.binding.group.0),
                         };
                         alpha_mask_deferred_phase.as_mut().unwrap().add(
                             bin_key,
@@ -988,10 +994,12 @@ pub fn queue_prepass_material_meshes<M: Material>(
                         );
                     } else if let Some(alpha_mask_phase) = alpha_mask_phase.as_mut() {
                         let bin_key = OpaqueNoLightmap3dBinKey {
-                            pipeline: pipeline_id,
-                            draw_function: alpha_mask_draw_prepass,
+                            batch_set_key: OpaqueNoLightmap3dBatchSetKey {
+                                draw_function: alpha_mask_draw_prepass,
+                                pipeline: pipeline_id,
+                                material_bind_group_index: Some(material.binding.group.0),
+                            },
                             asset_id: mesh_instance.mesh_asset_id.into(),
-                            material_bind_group_index: Some(material.binding.group.0),
                         };
                         alpha_mask_phase.add(
                             bin_key,

crates/bevy_pbr/src/render/pbr_prepass.wgsl

@@ -6,6 +6,7 @@
     pbr_functions,
     pbr_functions::SampleBias,
     prepass_io,
+    mesh_bindings::mesh,
     mesh_view_bindings::view,
 }
 
@@ -28,6 +29,15 @@ fn fragment(
     let is_front = true;
 #else   // MESHLET_MESH_MATERIAL_PASS
 
+#ifdef BINDLESS
+    let slot = mesh[in.instance_index].material_bind_group_slot;
+    let flags = pbr_bindings::material[slot].flags;
+    let uv_transform = pbr_bindings::material[slot].uv_transform;
+#else   // BINDLESS
+    let flags = pbr_bindings::material.flags;
+    let uv_transform = pbr_bindings::material.uv_transform;
+#endif  // BINDLESS
+
     // If we're in the crossfade section of a visibility range, conditionally
     // discard the fragment according to the visibility pattern.
 #ifdef VISIBILITY_RANGE_DITHER
@@ -45,8 +55,8 @@ fn fragment(
 
 #ifdef NORMAL_PREPASS
     // NOTE: Unlit bit not set means == 0 is true, so the true case is if lit
-    if (material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u {
-        let double_sided = (material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u;
+    if (flags & pbr_types::STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u {
+        let double_sided = (flags & pbr_types::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u;
 
         let world_normal = pbr_functions::prepare_world_normal(
             in.world_normal,
@@ -62,9 +72,9 @@ fn fragment(
 
 // TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
 #ifdef STANDARD_MATERIAL_NORMAL_MAP_UV_B
-        let uv = (material.uv_transform * vec3(in.uv_b, 1.0)).xy;
+        let uv = (uv_transform * vec3(in.uv_b, 1.0)).xy;
 #else
-        let uv = (material.uv_transform * vec3(in.uv, 1.0)).xy;
+        let uv = (uv_transform * vec3(in.uv, 1.0)).xy;
 #endif
 
         // Fill in the sample bias so we can sample from textures.
@@ -100,7 +110,7 @@ fn fragment(
         let TBN = pbr_functions::calculate_tbn_mikktspace(normal, in.world_tangent);
 
         normal = pbr_functions::apply_normal_mapping(
-            material.flags,
+            flags,
             TBN,
             double_sided,
             is_front,

crates/bevy_render/src/batching/gpu_preprocessing.rs

@@ -683,7 +683,7 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
         // batch sets. This variable stores the last batch set key that we've
         // seen. If our current batch set key is identical to this one, we can
         // merge the current batch into the last batch set.
-        let mut last_multidraw_key = None;
+        let mut maybe_last_multidraw_key = None;
 
         for key in &phase.batchable_mesh_keys {
             let mut batch: Option<BinnedRenderPhaseBatch> = None;
@@ -756,12 +756,16 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
                         // We're in multi-draw mode. Check to see whether our
                         // batch set key is the same as the last one. If so,
                         // merge this batch into the preceding batch set.
-                        let this_multidraw_key = key.get_batch_set_key();
-                        if last_multidraw_key.as_ref() == Some(&this_multidraw_key) {
-                            batch_sets.last_mut().unwrap().push(batch);
-                        } else {
-                            last_multidraw_key = Some(this_multidraw_key);
-                            batch_sets.push(vec![batch]);
+                        match (&maybe_last_multidraw_key, key.get_batch_set_key()) {
+                            (Some(ref last_multidraw_key), Some(this_multidraw_key))
+                                if *last_multidraw_key == this_multidraw_key =>
+                            {
+                                batch_sets.last_mut().unwrap().push(batch);
+                            }
+                            (_, maybe_this_multidraw_key) => {
+                                maybe_last_multidraw_key = maybe_this_multidraw_key;
+                                batch_sets.push(vec![batch]);
+                            }
                         }
                     }
                 }