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 <IceSentry@users.noreply.github.com>

Author: tychedelia

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<f32>;
+@group(2) @binding(1) var texture_sampler: sampler;
+
+struct Vertex {
+    @builtin(instance_index) instance_index: u32,
+    @location(0) position: vec3<f32>,
+};
+
+struct VertexOutput {
+    @builtin(position) clip_position: vec4<f32>,
+    @location(0) world_position: vec4<f32>,
+    @location(1) color: vec4<f32>,
+};
+
+@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<u32>(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<f32> {
+    return mesh.color;
+}

assets/shaders/storage_buffer.wgsl

@@ -4,7 +4,6 @@
 }
 
 @group(2) @binding(0) var<storage, read> colors: array<vec4<f32>, 5>;
-@group(2) @binding(1) var<uniform> 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;
 }
 

crates/bevy_pbr/src/meshlet/instance_manager.rs

@@ -125,6 +125,7 @@ impl InstanceManager {
             None,
             None,
             None,
+            None,
         );
 
         // Append instance data

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<u32>,
         previous_skin_index: Option<u32>,
+        tag: Option<u32>,
     ) -> 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<LightmapSlabIndex>,
+    /// 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<NoFrustumCulling>,
             Has<NotShadowReceiver>,
             Has<TransmittedShadowReceiver>,
@@ -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<NoFrustumCulling>,
                 Has<NotShadowReceiver>,
                 Has<TransmittedShadowReceiver>,
@@ -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),
         ))
     }
 

crates/bevy_pbr/src/render/mesh_functions.wgsl

@@ -132,3 +132,7 @@ fn get_visibility_range_dither_level(instance_index: u32, world_position: vec4<f
     return offset + clamp(level, 0, 16);
 }
 #endif
+
+fn get_tag(instance_index: u32) -> u32 {
+    return mesh[instance_index].tag;
+}

crates/bevy_pbr/src/render/mesh_preprocess.wgsl

@@ -345,4 +345,5 @@ fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
     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;
 }

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

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.

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);

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.

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,17 +231,19 @@ 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 {
             world_from_local: mesh_transforms.world_from_local.to_transpose(),
             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<Mesh>,
     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<NoAutomaticBatching>,
         )>,
     >,
 ) {
     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<Self::CompareData>)> {
         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<Self::BufferData> {
         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(

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,
             },
         );
     }