From a861452d68ad2a3aedc1291319c9bd1b5b29f31b Mon Sep 17 00:00:00 2001 From: charlotte Date: Mon, 10 Feb 2025 14:38:13 -0800 Subject: [PATCH] Add user supplied mesh tag (#17648) # Objective Because of mesh preprocessing, users cannot rely on `@builtin(instance_index)` in order to reference external data, as the instance index is not stable, either from frame to frame or relative to the total spawn order of mesh instances. ## Solution Add a user supplied mesh index that can be used for referencing external data when drawing instanced meshes. Closes #13373 ## Testing Benchmarked `many_cubes` showing no difference in total frame time. ## Showcase https://github.com/user-attachments/assets/80620147-aafc-4d9d-a8ee-e2149f7c8f3b --------- Co-authored-by: IceSentry --- assets/shaders/automatic_instancing.wgsl | 43 +++++++ assets/shaders/storage_buffer.wgsl | 4 +- .../bevy_pbr/src/meshlet/instance_manager.rs | 1 + crates/bevy_pbr/src/render/mesh.rs | 26 +++- .../bevy_pbr/src/render/mesh_functions.wgsl | 4 + .../bevy_pbr/src/render/mesh_preprocess.wgsl | 1 + crates/bevy_pbr/src/render/mesh_types.wgsl | 2 + .../mesh_preprocess_types.wgsl | 5 +- crates/bevy_render/src/mesh/components.rs | 5 + crates/bevy_render/src/mesh/mod.rs | 2 +- crates/bevy_sprite/src/mesh2d/mesh.rs | 19 ++- examples/2d/mesh2d_manual.rs | 1 + examples/shader/automatic_instancing.rs | 115 +++++++++++++----- examples/shader/custom_render_phase.rs | 2 + examples/shader/storage_buffer.rs | 34 ++---- 15 files changed, 196 insertions(+), 68 deletions(-) create mode 100644 assets/shaders/automatic_instancing.wgsl diff --git a/assets/shaders/automatic_instancing.wgsl b/assets/shaders/automatic_instancing.wgsl new file mode 100644 index 0000000000..35276246b0 --- /dev/null +++ b/assets/shaders/automatic_instancing.wgsl @@ -0,0 +1,43 @@ +#import bevy_pbr::{ + mesh_functions, + view_transformations::position_world_to_clip +} + +@group(2) @binding(0) var texture: texture_2d; +@group(2) @binding(1) var texture_sampler: sampler; + +struct Vertex { + @builtin(instance_index) instance_index: u32, + @location(0) position: vec3, +}; + +struct VertexOutput { + @builtin(position) clip_position: vec4, + @location(0) world_position: vec4, + @location(1) color: vec4, +}; + +@vertex +fn vertex(vertex: Vertex) -> VertexOutput { + var out: VertexOutput; + + // Lookup the tag for the given mesh + let tag = mesh_functions::get_tag(vertex.instance_index); + var world_from_local = mesh_functions::get_world_from_local(vertex.instance_index); + out.world_position = mesh_functions::mesh_position_local_to_world(world_from_local, vec4(vertex.position, 1.0)); + out.clip_position = position_world_to_clip(out.world_position.xyz); + + let tex_dim = textureDimensions(texture); + // Find the texel coordinate as derived from the tag + let texel_coord = vec2(tag % tex_dim.x, tag / tex_dim.x); + + out.color = textureLoad(texture, texel_coord, 0); + return out; +} + +@fragment +fn fragment( + mesh: VertexOutput, +) -> @location(0) vec4 { + return mesh.color; +} \ No newline at end of file diff --git a/assets/shaders/storage_buffer.wgsl b/assets/shaders/storage_buffer.wgsl index 1859e8dde2..c27053b9a2 100644 --- a/assets/shaders/storage_buffer.wgsl +++ b/assets/shaders/storage_buffer.wgsl @@ -4,7 +4,6 @@ } @group(2) @binding(0) var colors: array, 5>; -@group(2) @binding(1) var color_id: u32; struct Vertex { @builtin(instance_index) instance_index: u32, @@ -20,11 +19,12 @@ struct VertexOutput { @vertex fn vertex(vertex: Vertex) -> VertexOutput { var out: VertexOutput; + let tag = mesh_functions::get_tag(vertex.instance_index); var world_from_local = mesh_functions::get_world_from_local(vertex.instance_index); out.world_position = mesh_functions::mesh_position_local_to_world(world_from_local, vec4(vertex.position, 1.0)); out.clip_position = position_world_to_clip(out.world_position.xyz); - out.color = colors[color_id]; + out.color = colors[tag]; return out; } diff --git a/crates/bevy_pbr/src/meshlet/instance_manager.rs b/crates/bevy_pbr/src/meshlet/instance_manager.rs index 26f6432a1f..f615114d41 100644 --- a/crates/bevy_pbr/src/meshlet/instance_manager.rs +++ b/crates/bevy_pbr/src/meshlet/instance_manager.rs @@ -125,6 +125,7 @@ impl InstanceManager { None, None, None, + None, ); // Append instance data diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs index 996cc03dae..a02b56e441 100644 --- a/crates/bevy_pbr/src/render/mesh.rs +++ b/crates/bevy_pbr/src/render/mesh.rs @@ -485,6 +485,8 @@ pub struct MeshUniform { /// Low 16 bits: index of the material inside the bind group data. /// High 16 bits: index of the lightmap in the binding array. pub material_and_lightmap_bind_group_slot: u32, + /// User supplied tag to identify this mesh instance. + pub tag: u32, } /// Information that has to be transferred from CPU to GPU in order to produce @@ -541,10 +543,10 @@ pub struct MeshInputUniform { /// Low 16 bits: index of the material inside the bind group data. /// High 16 bits: index of the lightmap in the binding array. pub material_and_lightmap_bind_group_slot: u32, + /// User supplied tag to identify this mesh instance. + pub tag: u32, /// Padding. - pub pad_a: u32, - /// Padding. - pub pad_b: u32, + pub pad: u32, } /// Information about each mesh instance needed to cull it on GPU. @@ -578,6 +580,7 @@ impl MeshUniform { maybe_lightmap: Option<(LightmapSlotIndex, Rect)>, current_skin_index: Option, previous_skin_index: Option, + tag: Option, ) -> Self { let (local_from_world_transpose_a, local_from_world_transpose_b) = mesh_transforms.world_from_local.inverse_transpose_3x3(); @@ -598,6 +601,7 @@ impl MeshUniform { previous_skin_index: previous_skin_index.unwrap_or(u32::MAX), material_and_lightmap_bind_group_slot: u32::from(material_bind_group_slot) | ((lightmap_bind_group_slot as u32) << 16), + tag: tag.unwrap_or(0), } } } @@ -729,6 +733,8 @@ pub struct RenderMeshInstanceShared { /// Index of the slab that the lightmap resides in, if a lightmap is /// present. pub lightmap_slab_index: Option, + /// User supplied tag to identify this mesh instance. + pub tag: u32, } /// Information that is gathered during the parallel portion of mesh extraction @@ -808,6 +814,7 @@ impl RenderMeshInstanceShared { fn from_components( previous_transform: Option<&PreviousGlobalTransform>, mesh: &Mesh3d, + tag: Option<&MeshTag>, not_shadow_caster: bool, no_automatic_batching: bool, ) -> Self { @@ -828,6 +835,7 @@ impl RenderMeshInstanceShared { // This gets filled in later, during `RenderMeshGpuBuilder::update`. material_bindings_index: default(), lightmap_slab_index: None, + tag: tag.map_or(0, |i| **i), } } @@ -1160,8 +1168,8 @@ impl RenderMeshInstanceGpuBuilder { material_and_lightmap_bind_group_slot: u32::from( self.shared.material_bindings_index.slot, ) | ((lightmap_slot as u32) << 16), - pad_a: 0, - pad_b: 0, + tag: self.shared.tag, + pad: 0, }; // Did the last frame contain this entity as well? @@ -1296,6 +1304,7 @@ pub fn extract_meshes_for_cpu_building( &GlobalTransform, Option<&PreviousGlobalTransform>, &Mesh3d, + Option<&MeshTag>, Has, Has, Has, @@ -1314,6 +1323,7 @@ pub fn extract_meshes_for_cpu_building( transform, previous_transform, mesh, + tag, no_frustum_culling, not_shadow_receiver, transmitted_receiver, @@ -1341,6 +1351,7 @@ pub fn extract_meshes_for_cpu_building( let shared = RenderMeshInstanceShared::from_components( previous_transform, mesh, + tag, not_shadow_caster, no_automatic_batching, ); @@ -1402,6 +1413,7 @@ pub fn extract_meshes_for_gpu_building( Option<&Lightmap>, Option<&Aabb>, &Mesh3d, + Option<&MeshTag>, Has, Has, Has, @@ -1459,6 +1471,7 @@ pub fn extract_meshes_for_gpu_building( lightmap, aabb, mesh, + tag, no_frustum_culling, not_shadow_receiver, transmitted_receiver, @@ -1487,6 +1500,7 @@ pub fn extract_meshes_for_gpu_building( let shared = RenderMeshInstanceShared::from_components( previous_transform, mesh, + tag, not_shadow_caster, no_automatic_batching, ); @@ -1840,6 +1854,7 @@ impl GetBatchData for MeshPipeline { maybe_lightmap.map(|lightmap| (lightmap.slot_index, lightmap.uv_rect)), current_skin_index, previous_skin_index, + Some(mesh_instance.tag), ), mesh_instance.should_batch().then_some(( material_bind_group_index.group, @@ -1907,6 +1922,7 @@ impl GetFullBatchData for MeshPipeline { maybe_lightmap.map(|lightmap| (lightmap.slot_index, lightmap.uv_rect)), current_skin_index, previous_skin_index, + Some(mesh_instance.tag), )) } diff --git a/crates/bevy_pbr/src/render/mesh_functions.wgsl b/crates/bevy_pbr/src/render/mesh_functions.wgsl index 23857bc6aa..93392862fc 100644 --- a/crates/bevy_pbr/src/render/mesh_functions.wgsl +++ b/crates/bevy_pbr/src/render/mesh_functions.wgsl @@ -132,3 +132,7 @@ fn get_visibility_range_dither_level(instance_index: u32, world_position: vec4 u32 { + return mesh[instance_index].tag; +} \ No newline at end of file diff --git a/crates/bevy_pbr/src/render/mesh_preprocess.wgsl b/crates/bevy_pbr/src/render/mesh_preprocess.wgsl index 315dd13d3f..9bbe96e80d 100644 --- a/crates/bevy_pbr/src/render/mesh_preprocess.wgsl +++ b/crates/bevy_pbr/src/render/mesh_preprocess.wgsl @@ -345,4 +345,5 @@ fn main(@builtin(global_invocation_id) global_invocation_id: vec3) { output[mesh_output_index].previous_skin_index = current_input[input_index].previous_skin_index; output[mesh_output_index].material_and_lightmap_bind_group_slot = current_input[input_index].material_and_lightmap_bind_group_slot; + output[mesh_output_index].tag = current_input[input_index].tag; } diff --git a/crates/bevy_pbr/src/render/mesh_types.wgsl b/crates/bevy_pbr/src/render/mesh_types.wgsl index f0258770c6..ebac14915e 100644 --- a/crates/bevy_pbr/src/render/mesh_types.wgsl +++ b/crates/bevy_pbr/src/render/mesh_types.wgsl @@ -22,6 +22,8 @@ struct Mesh { // Low 16 bits: index of the material inside the bind group data. // High 16 bits: index of the lightmap in the binding array. material_and_lightmap_bind_group_slot: u32, + // User supplied index to identify the mesh instance + tag: u32, }; #ifdef SKINNED diff --git a/crates/bevy_render/src/experimental/occlusion_culling/mesh_preprocess_types.wgsl b/crates/bevy_render/src/experimental/occlusion_culling/mesh_preprocess_types.wgsl index af9deab8b8..6d1199adb4 100644 --- a/crates/bevy_render/src/experimental/occlusion_culling/mesh_preprocess_types.wgsl +++ b/crates/bevy_render/src/experimental/occlusion_culling/mesh_preprocess_types.wgsl @@ -19,8 +19,9 @@ struct MeshInput { // Low 16 bits: index of the material inside the bind group data. // High 16 bits: index of the lightmap in the binding array. material_and_lightmap_bind_group_slot: u32, - pad_a: u32, - pad_b: u32, + // User supplied index to identify the mesh instance + tag: u32, + pad: u32, } // The `wgpu` indirect parameters structure. This is a union of two structures. diff --git a/crates/bevy_render/src/mesh/components.rs b/crates/bevy_render/src/mesh/components.rs index 7c0d87cf51..b5b03ac2b8 100644 --- a/crates/bevy_render/src/mesh/components.rs +++ b/crates/bevy_render/src/mesh/components.rs @@ -150,3 +150,8 @@ pub fn mark_3d_meshes_as_changed_if_their_assets_changed( } } } + +/// A component that stores an arbitrary index used to identify the mesh instance when rendering. +#[derive(Component, Clone, Debug, Default, Deref, DerefMut, Reflect, PartialEq, Eq)] +#[reflect(Component, Default)] +pub struct MeshTag(pub u32); diff --git a/crates/bevy_render/src/mesh/mod.rs b/crates/bevy_render/src/mesh/mod.rs index cb89aaf6bd..40fb08d987 100644 --- a/crates/bevy_render/src/mesh/mod.rs +++ b/crates/bevy_render/src/mesh/mod.rs @@ -21,7 +21,7 @@ use bevy_ecs::{ SystemParamItem, }, }; -pub use components::{mark_3d_meshes_as_changed_if_their_assets_changed, Mesh2d, Mesh3d}; +pub use components::{mark_3d_meshes_as_changed_if_their_assets_changed, Mesh2d, Mesh3d, MeshTag}; use wgpu::IndexFormat; /// Registers all [`MeshBuilder`] types. diff --git a/crates/bevy_sprite/src/mesh2d/mesh.rs b/crates/bevy_sprite/src/mesh2d/mesh.rs index 59db71e8bc..13ad5981e5 100644 --- a/crates/bevy_sprite/src/mesh2d/mesh.rs +++ b/crates/bevy_sprite/src/mesh2d/mesh.rs @@ -19,6 +19,7 @@ use bevy_ecs::{ }; use bevy_image::{BevyDefault, Image, ImageSampler, TextureFormatPixelInfo}; use bevy_math::{Affine3, Vec4}; +use bevy_render::mesh::MeshTag; use bevy_render::prelude::Msaa; use bevy_render::RenderSet::PrepareAssets; use bevy_render::{ @@ -230,10 +231,11 @@ pub struct Mesh2dUniform { pub local_from_world_transpose_a: [Vec4; 2], pub local_from_world_transpose_b: f32, pub flags: u32, + pub tag: u32, } -impl From<&Mesh2dTransforms> for Mesh2dUniform { - fn from(mesh_transforms: &Mesh2dTransforms) -> Self { +impl Mesh2dUniform { + fn from_components(mesh_transforms: &Mesh2dTransforms, tag: u32) -> Self { let (local_from_world_transpose_a, local_from_world_transpose_b) = mesh_transforms.world_from_local.inverse_transpose_3x3(); Self { @@ -241,6 +243,7 @@ impl From<&Mesh2dTransforms> for Mesh2dUniform { local_from_world_transpose_a, local_from_world_transpose_b, flags: mesh_transforms.flags, + tag, } } } @@ -259,6 +262,7 @@ pub struct RenderMesh2dInstance { pub mesh_asset_id: AssetId, pub material_bind_group_id: Material2dBindGroupId, pub automatic_batching: bool, + pub tag: u32, } #[derive(Default, Resource, Deref, DerefMut)] @@ -275,13 +279,14 @@ pub fn extract_mesh2d( &ViewVisibility, &GlobalTransform, &Mesh2d, + Option<&MeshTag>, Has, )>, >, ) { render_mesh_instances.clear(); - for (entity, view_visibility, transform, handle, no_automatic_batching) in &query { + for (entity, view_visibility, transform, handle, tag, no_automatic_batching) in &query { if !view_visibility.get() { continue; } @@ -295,6 +300,7 @@ pub fn extract_mesh2d( mesh_asset_id: handle.0.id(), material_bind_group_id: Material2dBindGroupId::default(), automatic_batching: !no_automatic_batching, + tag: tag.map_or(0, |i| **i), }, ); } @@ -422,7 +428,7 @@ impl GetBatchData for Mesh2dPipeline { ) -> Option<(Self::BufferData, Option)> { let mesh_instance = mesh_instances.get(&main_entity)?; Some(( - (&mesh_instance.transforms).into(), + Mesh2dUniform::from_components(&mesh_instance.transforms, mesh_instance.tag), mesh_instance.automatic_batching.then_some(( mesh_instance.material_bind_group_id, mesh_instance.mesh_asset_id, @@ -439,7 +445,10 @@ impl GetFullBatchData for Mesh2dPipeline { main_entity: MainEntity, ) -> Option { let mesh_instance = mesh_instances.get(&main_entity)?; - Some((&mesh_instance.transforms).into()) + Some(Mesh2dUniform::from_components( + &mesh_instance.transforms, + mesh_instance.tag, + )) } fn get_index_and_compare_data( diff --git a/examples/2d/mesh2d_manual.rs b/examples/2d/mesh2d_manual.rs index 31ac71c094..ca99a71ba0 100644 --- a/examples/2d/mesh2d_manual.rs +++ b/examples/2d/mesh2d_manual.rs @@ -352,6 +352,7 @@ pub fn extract_colored_mesh2d( transforms, material_bind_group_id: Material2dBindGroupId::default(), automatic_batching: false, + tag: 0, }, ); } diff --git a/examples/shader/automatic_instancing.rs b/examples/shader/automatic_instancing.rs index d3cd7c45cc..854e08a4dd 100644 --- a/examples/shader/automatic_instancing.rs +++ b/examples/shader/automatic_instancing.rs @@ -1,51 +1,104 @@ //! Shows that multiple instances of a cube are automatically instanced in one draw call //! Try running this example in a graphics profiler and all the cubes should be only a single draw call. +//! Also demonstrates how to use `MeshTag` to use external data in a custom material. -use bevy::prelude::*; +use bevy::{ + prelude::*, + reflect::TypePath, + render::{ + mesh::MeshTag, + render_resource::{AsBindGroup, ShaderRef}, + }, +}; + +const SHADER_ASSET_PATH: &str = "shaders/automatic_instancing.wgsl"; fn main() { App::new() - .add_plugins(DefaultPlugins) + .add_plugins((DefaultPlugins, MaterialPlugin::::default())) .add_systems(Startup, setup) + .add_systems(Update, update) .run(); } -/// set up a simple 3D scene +/// Sets up an instanced grid of cubes, where each cube is colored based on an image that is +/// sampled in the vertex shader. The cubes are then animated in a spiral pattern. +/// +/// This example demonstrates one use of automatic instancing and how to use `MeshTag` to use +/// external data in a custom material. For example, here we use the "index" of each cube to +/// determine the texel coordinate to sample from the image in the shader. fn setup( mut commands: Commands, + assets: Res, mut meshes: ResMut>, - mut materials: ResMut>, + mut materials: ResMut>, ) { - // camera + // We will use this image as our external data for our material to sample from in the vertex shader + let image = assets.load("branding/icon.png"); + + // Our single mesh handle that will be instanced + let mesh_handle = meshes.add(Cuboid::from_size(Vec3::splat(0.01))); + + // Create the custom material with a reference to our texture + // Automatic instancing works with any Material, including the `StandardMaterial`. + // This custom material is used to demonstrate the optional `MeshTag` feature. + let material_handle = materials.add(CustomMaterial { + image: image.clone(), + }); + + // We're hardcoding the image dimensions for simplicity + let image_dims = UVec2::new(256, 256); + let total_pixels = image_dims.x * image_dims.y; + + for index in 0..total_pixels { + // Get x,y from index - x goes left to right, y goes top to bottom + let x = index % image_dims.x; + let y = index / image_dims.x; + + // Convert to centered world coordinates + let world_x = (x as f32 - image_dims.x as f32 / 2.0) / 50.0; + let world_y = -((y as f32 - image_dims.y as f32 / 2.0) / 50.0); // Still need negative for world space + + commands.spawn(( + // For automatic instancing to take effect you need to + // use the same mesh handle and material handle for each instance + Mesh3d(mesh_handle.clone()), + MeshMaterial3d(material_handle.clone()), + // This is an optional component that can be used to help tie external data to a mesh instance + MeshTag(index), + Transform::from_xyz(world_x, world_y, 0.0), + )); + } + + // Camera commands.spawn(( Camera3d::default(), - Transform::from_xyz(0.0, 48.0, 20.0).looking_at(Vec3::ZERO, Vec3::Y), - )); - // light - commands.spawn(( - PointLight { - shadows_enabled: true, - ..default() - }, - Transform::from_xyz(0.0, 16.0, 8.0), + Transform::from_xyz(0.0, 0.0, 5.0).looking_at(Vec3::ZERO, Vec3::Y), )); +} - let mesh = meshes.add(Cuboid::from_size(Vec3::splat(0.5))); - // This example uses the StandardMaterial but it can work with most custom material too - let material = materials.add(Color::srgb_u8(124, 144, 255)); - - // spawn 1000 cubes - for x in -5..5 { - for y in -5..5 { - for z in -5..5 { - commands.spawn(( - // For automatic instancing to take effect you need to - // use the same mesh handle and material handle for each instance - Mesh3d(mesh.clone()), - MeshMaterial3d(material.clone()), - Transform::from_xyz(x as f32, y as f32, z as f32), - )); - } - } +// Animate the transform +fn update(time: Res