Allow phase items not associated with meshes to be binned. (#14029)

pcwalton · mockersf · commit 26f73132126d · 2024-06-27T19:41:42.000+02:00
As reported in #14004, many third-party plugins, such as Hanabi, enqueue entities that don't have meshes into render phases. However, the introduction of indirect mode added a dependency on mesh-specific data, breaking this workflow. This is because GPU preprocessing requires that the render phases manage indirect draw parameters, which don't apply to objects that aren't meshes. The existing code skips over binned entities that don't have indirect draw parameters, which causes the rendering to be skipped for such objects. To support this workflow, this commit adds a new field, `non_mesh_items`, to `BinnedRenderPhase`. This field contains a simple list of (bin key, entity) pairs. After drawing batchable and unbatchable objects, the non-mesh items are drawn one after another. Bevy itself doesn't enqueue any items into this list; it exists solely for the application and/or plugins to use. Additionally, this commit switches the asset ID in the standard bin keys to be an untyped asset ID rather than that of a mesh. This allows more flexibility, allowing bins to be keyed off any type of asset. This patch adds a new example, `custom_phase_item`, which simultaneously serves to demonstrate how to use this new feature and to act as a regression test so this doesn't break again. Fixes #14004. ## Changelog ### Added * `BinnedRenderPhase` now contains a `non_mesh_items` field for plugins to add custom items to.
diff --git a/Cargo.toml b/Cargo.toml
@@ -3219,6 +3219,17 @@ description = "Displays an example model with anisotropy"
 category = "3D Rendering"
 wasm = false
 
+[[example]]
+name = "custom_phase_item"
+path = "examples/shader/custom_phase_item.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.custom_phase_item]
+name = "Custom phase item"
+description = "Demonstrates how to enqueue custom draw commands in a render phase"
+category = "Shaders"
+wasm = true
+
 [profile.wasm-release]
 inherits = "release"
 opt-level = "z"
diff --git a/assets/shaders/custom_phase_item.wgsl b/assets/shaders/custom_phase_item.wgsl
@@ -0,0 +1,36 @@
+// `custom_phase_item.wgsl`
+//
+// This shader goes with the `custom_phase_item` example. It demonstrates how to
+// enqueue custom rendering logic in a `RenderPhase`.
+
+// The GPU-side vertex structure.
+struct Vertex {
+    // The world-space position of the vertex.
+    @location(0) position: vec3<f32>,
+    // The color of the vertex.
+    @location(1) color: vec3<f32>,
+};
+
+// Information passed from the vertex shader to the fragment shader.
+struct VertexOutput {
+    // The clip-space position of the vertex.
+    @builtin(position) clip_position: vec4<f32>,
+    // The color of the vertex.
+    @location(0) color: vec3<f32>,
+};
+
+// The vertex shader entry point.
+@vertex
+fn vertex(vertex: Vertex) -> VertexOutput {
+    // Use an orthographic projection.
+    var vertex_output: VertexOutput;
+    vertex_output.clip_position = vec4(vertex.position.xyz, 1.0);
+    vertex_output.color = vertex.color;
+    return vertex_output;
+}
+
+// The fragment shader entry point.
+@fragment
+fn fragment(vertex_output: VertexOutput) -> @location(0) vec4<f32> {
+    return vec4(vertex_output.color, 1.0);
+}
diff --git a/crates/bevy_asset/src/id.rs b/crates/bevy_asset/src/id.rs
@@ -288,13 +288,17 @@ impl Hash for UntypedAssetId {
     }
 }
 
+impl Ord for UntypedAssetId {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.type_id()
+            .cmp(&other.type_id())
+            .then_with(|| self.internal().cmp(&other.internal()))
+    }
+}
+
 impl PartialOrd for UntypedAssetId {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        if self.type_id() != other.type_id() {
-            None
-        } else {
-            Some(self.internal().cmp(&other.internal()))
-        }
+        Some(self.cmp(other))
     }
 }
 
diff --git a/crates/bevy_core_pipeline/src/core_3d/mod.rs b/crates/bevy_core_pipeline/src/core_3d/mod.rs
@@ -64,7 +64,7 @@ pub const DEPTH_TEXTURE_SAMPLING_SUPPORTED: bool = true;
 
 use std::ops::Range;
 
-use bevy_asset::AssetId;
+use bevy_asset::{AssetId, UntypedAssetId};
 use bevy_color::LinearRgba;
 pub use camera_3d::*;
 pub use main_opaque_pass_3d_node::*;
@@ -76,7 +76,6 @@ use bevy_math::FloatOrd;
 use bevy_render::{
     camera::{Camera, ExtractedCamera},
     extract_component::ExtractComponentPlugin,
-    mesh::Mesh,
     prelude::Msaa,
     render_graph::{EmptyNode, RenderGraphApp, ViewNodeRunner},
     render_phase::{
@@ -221,7 +220,7 @@ pub struct Opaque3d {
     pub extra_index: PhaseItemExtraIndex,
 }
 
-/// Data that must be identical in order to batch meshes together.
+/// Data that must be identical in order to batch phase items together.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Opaque3dBinKey {
     /// The identifier of the render pipeline.
@@ -230,8 +229,11 @@ pub struct Opaque3dBinKey {
     /// The function used to draw.
     pub draw_function: DrawFunctionId,
 
-    /// The mesh.
-    pub asset_id: AssetId<Mesh>,
+    /// The asset that this phase item is associated with.
+    ///
+    /// Normally, this is the ID of the mesh, but for non-mesh items it might be
+    /// the ID of another type of asset.
+    pub asset_id: UntypedAssetId,
 
     /// The ID of a bind group specific to the material.
     ///
diff --git a/crates/bevy_core_pipeline/src/deferred/node.rs b/crates/bevy_core_pipeline/src/deferred/node.rs
@@ -144,8 +144,8 @@ impl ViewNode for DeferredGBufferPrepassNode {
             }
 
             // Opaque draws
-            if !opaque_deferred_phase.batchable_keys.is_empty()
-                || !opaque_deferred_phase.unbatchable_keys.is_empty()
+            if !opaque_deferred_phase.batchable_mesh_keys.is_empty()
+                || !opaque_deferred_phase.unbatchable_mesh_keys.is_empty()
             {
                 #[cfg(feature = "trace")]
                 let _opaque_prepass_span = info_span!("opaque_deferred_prepass").entered();
diff --git a/crates/bevy_core_pipeline/src/prepass/mod.rs b/crates/bevy_core_pipeline/src/prepass/mod.rs
@@ -29,12 +29,11 @@ pub mod node;
 
 use std::ops::Range;
 
-use bevy_asset::AssetId;
+use bevy_asset::UntypedAssetId;
 use bevy_ecs::prelude::*;
 use bevy_math::Mat4;
 use bevy_reflect::Reflect;
 use bevy_render::{
-    mesh::Mesh,
     render_phase::{
         BinnedPhaseItem, CachedRenderPipelinePhaseItem, DrawFunctionId, PhaseItem,
         PhaseItemExtraIndex,
@@ -147,7 +146,7 @@ pub struct Opaque3dPrepass {
 }
 
 // TODO: Try interning these.
-/// The data used to bin each opaque 3D mesh in the prepass and deferred pass.
+/// 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 ID of the GPU pipeline.
@@ -156,8 +155,8 @@ pub struct OpaqueNoLightmap3dBinKey {
     /// The function used to draw the mesh.
     pub draw_function: DrawFunctionId,
 
-    /// The ID of the mesh.
-    pub asset_id: AssetId<Mesh>,
+    /// The ID of the asset.
+    pub asset_id: UntypedAssetId,
 
     /// The ID of a bind group specific to the material.
     ///
diff --git a/crates/bevy_core_pipeline/src/prepass/node.rs b/crates/bevy_core_pipeline/src/prepass/node.rs
@@ -120,8 +120,8 @@ impl ViewNode for PrepassNode {
             }
 
             // Opaque draws
-            if !opaque_prepass_phase.batchable_keys.is_empty()
-                || !opaque_prepass_phase.unbatchable_keys.is_empty()
+            if !opaque_prepass_phase.batchable_mesh_keys.is_empty()
+                || !opaque_prepass_phase.unbatchable_mesh_keys.is_empty()
             {
                 #[cfg(feature = "trace")]
                 let _opaque_prepass_span = info_span!("opaque_prepass").entered();
diff --git a/crates/bevy_pbr/src/material.rs b/crates/bevy_pbr/src/material.rs
@@ -763,11 +763,15 @@ pub fn queue_material_meshes<M: Material>(
                         let bin_key = Opaque3dBinKey {
                             draw_function: draw_opaque_pbr,
                             pipeline: pipeline_id,
-                            asset_id: mesh_instance.mesh_asset_id,
+                            asset_id: mesh_instance.mesh_asset_id.into(),
                             material_bind_group_id: material.get_bind_group_id().0,
                             lightmap_image,
                         };
-                        opaque_phase.add(bin_key, *visible_entity, mesh_instance.should_batch());
+                        opaque_phase.add(
+                            bin_key,
+                            *visible_entity,
+                            BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
+                        );
                     }
                 }
                 // Alpha mask
@@ -787,13 +791,13 @@ pub fn queue_material_meshes<M: Material>(
                         let bin_key = OpaqueNoLightmap3dBinKey {
                             draw_function: draw_alpha_mask_pbr,
                             pipeline: pipeline_id,
-                            asset_id: mesh_instance.mesh_asset_id,
+                            asset_id: mesh_instance.mesh_asset_id.into(),
                             material_bind_group_id: material.get_bind_group_id().0,
                         };
                         alpha_mask_phase.add(
                             bin_key,
                             *visible_entity,
-                            mesh_instance.should_batch(),
+                            BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                         );
                     }
                 }
diff --git a/crates/bevy_pbr/src/prepass/mod.rs b/crates/bevy_pbr/src/prepass/mod.rs
@@ -860,22 +860,22 @@ pub fn queue_prepass_material_meshes<M: Material>(
                             OpaqueNoLightmap3dBinKey {
                                 draw_function: opaque_draw_deferred,
                                 pipeline: pipeline_id,
-                                asset_id: mesh_instance.mesh_asset_id,
+                                asset_id: mesh_instance.mesh_asset_id.into(),
                                 material_bind_group_id: material.get_bind_group_id().0,
                             },
                             *visible_entity,
-                            mesh_instance.should_batch(),
+                            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,
-                                asset_id: mesh_instance.mesh_asset_id,
+                                asset_id: mesh_instance.mesh_asset_id.into(),
                                 material_bind_group_id: material.get_bind_group_id().0,
                             },
                             *visible_entity,
-                            mesh_instance.should_batch(),
+                            BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                         );
                     }
                 }
@@ -885,25 +885,25 @@ pub fn queue_prepass_material_meshes<M: Material>(
                         let bin_key = OpaqueNoLightmap3dBinKey {
                             pipeline: pipeline_id,
                             draw_function: alpha_mask_draw_deferred,
-                            asset_id: mesh_instance.mesh_asset_id,
+                            asset_id: mesh_instance.mesh_asset_id.into(),
                             material_bind_group_id: material.get_bind_group_id().0,
                         };
                         alpha_mask_deferred_phase.as_mut().unwrap().add(
                             bin_key,
                             *visible_entity,
-                            mesh_instance.should_batch(),
+                            BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                         );
                     } 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,
-                            asset_id: mesh_instance.mesh_asset_id,
+                            asset_id: mesh_instance.mesh_asset_id.into(),
                             material_bind_group_id: material.get_bind_group_id().0,
                         };
                         alpha_mask_phase.add(
                             bin_key,
                             *visible_entity,
-                            mesh_instance.should_batch(),
+                            BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                         );
                     }
                 }
diff --git a/crates/bevy_pbr/src/render/light.rs b/crates/bevy_pbr/src/render/light.rs
@@ -1,11 +1,10 @@
-use bevy_asset::AssetId;
+use bevy_asset::UntypedAssetId;
 use bevy_color::ColorToComponents;
 use bevy_core_pipeline::core_3d::CORE_3D_DEPTH_FORMAT;
 use bevy_ecs::entity::EntityHashSet;
 use bevy_ecs::prelude::*;
 use bevy_ecs::{entity::EntityHashMap, system::lifetimeless::Read};
 use bevy_math::{Mat4, UVec4, Vec2, Vec3, Vec3Swizzles, Vec4, Vec4Swizzles};
-use bevy_render::mesh::Mesh;
 use bevy_render::{
     diagnostic::RecordDiagnostics,
     mesh::GpuMesh,
@@ -1286,10 +1285,10 @@ pub fn queue_shadows<M: Material>(
                     ShadowBinKey {
                         draw_function: draw_shadow_mesh,
                         pipeline: pipeline_id,
-                        asset_id: mesh_instance.mesh_asset_id,
+                        asset_id: mesh_instance.mesh_asset_id.into(),
                     },
                     entity,
-                    mesh_instance.should_batch(),
+                    BinnedRenderPhaseType::mesh(mesh_instance.should_batch()),
                 );
             }
         }
@@ -1303,6 +1302,7 @@ pub struct Shadow {
     pub extra_index: PhaseItemExtraIndex,
 }
 
+/// Data used to bin each object in the shadow map phase.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct ShadowBinKey {
     /// The identifier of the render pipeline.
@@ -1311,8 +1311,8 @@ pub struct ShadowBinKey {
     /// The function used to draw.
     pub draw_function: DrawFunctionId,
 
-    /// The mesh.
-    pub asset_id: AssetId<Mesh>,
+    /// The object.
+    pub asset_id: UntypedAssetId,
 }
 
 impl PhaseItem for Shadow {
diff --git a/crates/bevy_render/src/batching/gpu_preprocessing.rs b/crates/bevy_render/src/batching/gpu_preprocessing.rs
@@ -523,9 +523,9 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
 
         // Prepare batchables.
 
-        for key in &phase.batchable_keys {
+        for key in &phase.batchable_mesh_keys {
             let mut batch: Option<BinnedRenderPhaseBatch> = None;
-            for &entity in &phase.batchable_values[key] {
+            for &entity in &phase.batchable_mesh_values[key] {
                 let Some(input_index) = GFBD::get_binned_index(&system_param_item, entity) else {
                     continue;
                 };
@@ -583,8 +583,8 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
         }
 
         // Prepare unbatchables.
-        for key in &phase.unbatchable_keys {
-            let unbatchables = phase.unbatchable_values.get_mut(key).unwrap();
+        for key in &phase.unbatchable_mesh_keys {
+            let unbatchables = phase.unbatchable_mesh_values.get_mut(key).unwrap();
             for &entity in &unbatchables.entities {
                 let Some(input_index) = GFBD::get_binned_index(&system_param_item, entity) else {
                     continue;
diff --git a/crates/bevy_render/src/batching/mod.rs b/crates/bevy_render/src/batching/mod.rs
@@ -156,8 +156,8 @@ where
     BPI: BinnedPhaseItem,
 {
     for phase in phases.values_mut() {
-        phase.batchable_keys.sort_unstable();
-        phase.unbatchable_keys.sort_unstable();
+        phase.batchable_mesh_keys.sort_unstable();
+        phase.unbatchable_mesh_keys.sort_unstable();
     }
 }
 
diff --git a/crates/bevy_render/src/batching/no_gpu_preprocessing.rs b/crates/bevy_render/src/batching/no_gpu_preprocessing.rs
@@ -104,9 +104,9 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
     for phase in phases.values_mut() {
         // Prepare batchables.
 
-        for key in &phase.batchable_keys {
+        for key in &phase.batchable_mesh_keys {
             let mut batch_set: SmallVec<[BinnedRenderPhaseBatch; 1]> = smallvec![];
-            for &entity in &phase.batchable_values[key] {
+            for &entity in &phase.batchable_mesh_values[key] {
                 let Some(buffer_data) = GFBD::get_binned_batch_data(&system_param_item, entity)
                 else {
                     continue;
@@ -141,8 +141,8 @@ pub fn batch_and_prepare_binned_render_phase<BPI, GFBD>(
         }
 
         // Prepare unbatchables.
-        for key in &phase.unbatchable_keys {
-            let unbatchables = phase.unbatchable_values.get_mut(key).unwrap();
+        for key in &phase.unbatchable_mesh_keys {
+            let unbatchables = phase.unbatchable_mesh_values.get_mut(key).unwrap();
             for &entity in &unbatchables.entities {
                 let Some(buffer_data) = GFBD::get_binned_batch_data(&system_param_item, entity)
                 else {
diff --git a/crates/bevy_render/src/render_phase/mod.rs b/crates/bevy_render/src/render_phase/mod.rs
diff --git a/crates/bevy_render/src/view/mod.rs b/crates/bevy_render/src/view/mod.rs
diff --git a/examples/README.md b/examples/README.md
diff --git a/examples/shader/custom_phase_item.rs b/examples/shader/custom_phase_item.rs

Original file line number	Diff line number	Diff line change
`@@ -288,13 +288,17 @@ impl Hash for UntypedAssetId {`
`288`	`288`	`}`
`289`	`289`	`}`
`290`	`290`
	`291`	`+impl Ord for UntypedAssetId {`
	`292`	`+ fn cmp(&self, other: &Self) -> std::cmp::Ordering {`
	`293`	`+ self.type_id()`
	`294`	`+ .cmp(&other.type_id())`
	`295`	`+ .then_with(\|\| self.internal().cmp(&other.internal()))`
	`296`	`+ }`
	`297`	`+}`
	`298`	`+`
`291`	`299`	`impl PartialOrd for UntypedAssetId {`
`292`	`300`	`fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {`
`293`		`- if self.type_id() != other.type_id() {`
`294`		`- None`
`295`		`- } else {`
`296`		`- Some(self.internal().cmp(&other.internal()))`
`297`		`- }`
	`301`	`+ Some(self.cmp(other))`
`298`	`302`	`}`
`299`	`303`	`}`
`300`	`304`
Original file line number	Diff line number	Diff line change
`@@ -144,8 +144,8 @@ impl ViewNode for DeferredGBufferPrepassNode {`
`144`	`144`	`}`
`145`	`145`
`146`	`146`	`// Opaque draws`
`147`		`- if !opaque_deferred_phase.batchable_keys.is_empty()`
`148`		`- \|\| !opaque_deferred_phase.unbatchable_keys.is_empty()`
	`147`	`+ if !opaque_deferred_phase.batchable_mesh_keys.is_empty()`
	`148`	`+ \|\| !opaque_deferred_phase.unbatchable_mesh_keys.is_empty()`
`149`	`149`	`{`
`150`	`150`	`#[cfg(feature = "trace")]`
`151`	`151`	`let _opaque_prepass_span = info_span!("opaque_deferred_prepass").entered();`
Original file line number	Diff line number	Diff line change
`@@ -120,8 +120,8 @@ impl ViewNode for PrepassNode {`
`120`	`120`	`}`
`121`	`121`
`122`	`122`	`// Opaque draws`
`123`		`- if !opaque_prepass_phase.batchable_keys.is_empty()`
`124`		`- \|\| !opaque_prepass_phase.unbatchable_keys.is_empty()`
	`123`	`+ if !opaque_prepass_phase.batchable_mesh_keys.is_empty()`
	`124`	`+ \|\| !opaque_prepass_phase.unbatchable_mesh_keys.is_empty()`
`125`	`125`	`{`
`126`	`126`	`#[cfg(feature = "trace")]`
`127`	`127`	`let _opaque_prepass_span = info_span!("opaque_prepass").entered();`