Skip to content

Unify RenderMaterialInstances and RenderMeshMaterialIds, and fix an associated race condition. #18734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 9, 2025
Merged

Unify RenderMaterialInstances and RenderMeshMaterialIds, and fix an associated race condition. #18734

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
64dc97b
Unify `RenderMaterialInstances` and `RenderMeshMaterialIds`, and fix an
pcwalton Apr 6, 2025
f73eda0
Unbreak meshlets
pcwalton Apr 6, 2025
e4bca83
Merge remote-tracking branch 'origin/main' into remove-render-mesh-ma…
pcwalton Apr 9, 2025
0210def
Address review comments
pcwalton Apr 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 130 additions & 35 deletions crates/bevy_pbr/src/material.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,7 @@ where
.init_resource::<EntitySpecializationTicks<M>>()
.init_resource::<SpecializedMaterialPipelineCache<M>>()
.init_resource::<DrawFunctions<Shadow>>()
.init_resource::<RenderMaterialInstances<M>>()
.init_resource::<RenderMaterialInstances>()
.add_render_command::<Shadow, DrawPrepass<M>>()
.add_render_command::<Transmissive3d, DrawMaterial<M>>()
.add_render_command::<Transparent3d, DrawMaterial<M>>()
Expand All @@ -316,7 +316,10 @@ where
.add_systems(
ExtractSchedule,
(
extract_mesh_materials::<M>.before(ExtractMeshesSet),
early_sweep_material_instances::<M>
.before(late_sweep_material_instances)
.before(ExtractMeshesSet),
extract_mesh_materials::<M>.before(early_sweep_material_instances::<M>),
extract_entities_needs_specialization::<M>.after(extract_cameras),
),
)
Expand Down Expand Up @@ -542,7 +545,7 @@ pub struct SetMaterialBindGroup<M: Material, const I: usize>(PhantomData<M>);
impl<P: PhaseItem, M: Material, const I: usize> RenderCommand<P> for SetMaterialBindGroup<M, I> {
type Param = (
SRes<RenderAssets<PreparedMaterial<M>>>,
SRes<RenderMaterialInstances<M>>,
SRes<RenderMaterialInstances>,
SRes<MaterialBindGroupAllocator<M>>,
);
type ViewQuery = ();
Expand All @@ -564,10 +567,13 @@ impl<P: PhaseItem, M: Material, const I: usize> RenderCommand<P> for SetMaterial
let material_instances = material_instances.into_inner();
let material_bind_group_allocator = material_bind_group_allocator.into_inner();

let Some(material_asset_id) = material_instances.get(&item.main_entity()) else {
let Some(material_instance) = material_instances.instances.get(&item.main_entity()) else {
return RenderCommandResult::Skip;
};
let Some(material) = materials.get(*material_asset_id) else {
let Ok(material_asset_id) = material_instance.asset_id.try_typed::<M>() else {
Copy link
Member

@tychedelia tychedelia Apr 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a note for myself, but this is basically guaranteed to succeed because the phase item was queued in a queue_material_meshes::<M> and so we would never look up a material instance whose asset id's material type doesn't match the draw command's,

Copy link
Contributor Author

@pcwalton pcwalton Apr 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you want me to add an error! there?

return RenderCommandResult::Skip;
};
let Some(material) = materials.get(material_asset_id) else {
return RenderCommandResult::Skip;
};
let Some(material_bind_group) = material_bind_group_allocator.get(material.binding.group)
Expand All @@ -582,16 +588,45 @@ impl<P: PhaseItem, M: Material, const I: usize> RenderCommand<P> for SetMaterial
}
}

/// Stores all extracted instances of a [`Material`] in the render world.
#[derive(Resource, Deref, DerefMut)]
pub struct RenderMaterialInstances<M: Material>(pub MainEntityHashMap<AssetId<M>>);
/// Stores all extracted instances of all [`Material`]s in the render world.
#[derive(Resource, Default)]
pub struct RenderMaterialInstances {
/// Maps from each entity in the main world to the
/// [`RenderMaterialInstance`] associated with it.
pub instances: MainEntityHashMap<RenderMaterialInstance>,
/// A monotonically-increasing counter, which we use to sweep
/// [`RenderMaterialInstances::instances`] when the entities and/or required
/// components are removed.
current_change_tick: Tick,
}

impl<M: Material> Default for RenderMaterialInstances<M> {
fn default() -> Self {
Self(Default::default())
impl RenderMaterialInstances {
/// Returns the mesh material ID for the entity with the given mesh, or a
/// dummy mesh material ID if the mesh has no material ID.
///
/// Meshes almost always have materials, but in very specific circumstances
/// involving custom pipelines they won't. (See the
/// `specialized_mesh_pipelines` example.)
pub(crate) fn mesh_material(&self, entity: MainEntity) -> UntypedAssetId {
match self.instances.get(&entity) {
Some(render_instance) => render_instance.asset_id,
None => AssetId::<StandardMaterial>::invalid().into(),
Copy link
Member

@tychedelia tychedelia Apr 6, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should extract this to a constant (in another PR) as we're using this pattern in a few places and the use of StandardMaterial is potentially confusing.

Copy link
Contributor Author

@pcwalton pcwalton Apr 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

}
}
}

/// The material associated with a single mesh instance in the main world.
///
/// Note that this uses an [`UntypedAssetId`] and isn't generic over the
/// material type, for simplicity.
pub struct RenderMaterialInstance {
/// The material asset.
pub(crate) asset_id: UntypedAssetId,
/// The [`RenderMaterialInstances::current_change_tick`] at which this
/// material instance was last modified.
last_change_tick: Tick,
}

pub const fn alpha_mode_pipeline_key(alpha_mode: AlphaMode, msaa: &Msaa) -> MeshPipelineKey {
match alpha_mode {
// Premultiplied and Add share the same pipeline key
Expand Down Expand Up @@ -672,39 +707,91 @@ fn mark_meshes_as_changed_if_their_materials_changed<M>(
/// Fills the [`RenderMaterialInstances`] and [`RenderMeshMaterialIds`]
/// resources from the meshes in the scene.
fn extract_mesh_materials<M: Material>(
mut material_instances: ResMut<RenderMaterialInstances<M>>,
mut material_ids: ResMut<RenderMeshMaterialIds>,
mut material_instances: ResMut<RenderMaterialInstances>,
changed_meshes_query: Extract<
Query<
(Entity, &ViewVisibility, &MeshMaterial3d<M>),
Or<(Changed<ViewVisibility>, Changed<MeshMaterial3d<M>>)>,
>,
>,
mut removed_visibilities_query: Extract<RemovedComponents<ViewVisibility>>,
mut removed_materials_query: Extract<RemovedComponents<MeshMaterial3d<M>>>,
) {
let last_change_tick = material_instances.current_change_tick;

for (entity, view_visibility, material) in &changed_meshes_query {
if view_visibility.get() {
material_instances.insert(entity.into(), material.id());
material_ids.insert(entity.into(), material.id().into());
material_instances.instances.insert(
entity.into(),
RenderMaterialInstance {
asset_id: material.id().untyped(),
last_change_tick,
},
);
} else {
material_instances.remove(&MainEntity::from(entity));
material_ids.remove(entity.into());
material_instances
.instances
.remove(&MainEntity::from(entity));
}
}
}

/// Removes mesh materials from [`RenderMaterialInstances`] when their
/// [`MeshMaterial3d`] components are removed.
///
/// This is tricky because we have to deal with the case in which a material of
/// type A was removed and replaced with a material of type B in the same frame
/// (which is actually somewhat common of an operation). In this case, even
/// though an entry will be present in `RemovedComponents<MeshMaterial3d<A>>`,
/// we must not remove the entry in `RenderMaterialInstances` which corresponds
/// to material B. To handle this case, we use change ticks to avoid removing
/// the entry if it was updated this frame.
///
/// This is the first of two sweep phases. Because this phase runs once per
/// material type, we need a second phase in order to guarantee that we only
/// bump [`RenderMaterialInstances::current_change_tick`] once.
fn early_sweep_material_instances<M>(
mut material_instances: ResMut<RenderMaterialInstances>,
mut removed_materials_query: Extract<RemovedComponents<MeshMaterial3d<M>>>,
) where
M: Material,
{
let last_change_tick = material_instances.current_change_tick;

for entity in removed_materials_query.read() {
if let Entry::Occupied(occupied_entry) = material_instances.instances.entry(entity.into()) {
// Only sweep the entry if it wasn't updated this frame.
if occupied_entry.get().last_change_tick != last_change_tick {
occupied_entry.remove();
}
}
}
}

for entity in removed_visibilities_query
.read()
.chain(removed_materials_query.read())
{
// Only queue a mesh for removal if we didn't pick it up above.
// It's possible that a necessary component was removed and re-added in
// the same frame.
if !changed_meshes_query.contains(entity) {
material_instances.remove(&MainEntity::from(entity));
material_ids.remove(entity.into());
/// Removes mesh materials from [`RenderMaterialInstances`] when their
/// [`ViewVisibility`] components are removed.
///
/// This runs after all invocations of [`early_sweep_material_instances`] and is
/// responsible for bumping [`RenderMaterialInstances::current_change_tick`] in
/// preparation for a new frame.
fn late_sweep_material_instances(
mut material_instances: ResMut<RenderMaterialInstances>,
mut removed_visibilities_query: Extract<RemovedComponents<ViewVisibility>>,
) {
let last_change_tick = material_instances.current_change_tick;

for entity in removed_visibilities_query.read() {
if let Entry::Occupied(occupied_entry) = material_instances.instances.entry(entity.into()) {
// Only sweep the entry if it wasn't updated this frame. It's
// possible that a `ViewVisibility` component was removed and
// re-added in the same frame.
if occupied_entry.get().last_change_tick != last_change_tick {
occupied_entry.remove();
}
}
}

material_instances
.current_change_tick
.set(last_change_tick.get() + 1);
}

pub fn extract_entities_needs_specialization<M>(
Expand Down Expand Up @@ -852,7 +939,7 @@ pub fn specialize_material_meshes<M: Material>(
render_meshes: Res<RenderAssets<RenderMesh>>,
render_materials: Res<RenderAssets<PreparedMaterial<M>>>,
render_mesh_instances: Res<RenderMeshInstances>,
render_material_instances: Res<RenderMaterialInstances<M>>,
render_material_instances: Res<RenderMaterialInstances>,
render_lightmaps: Res<RenderLightmaps>,
render_visibility_ranges: Res<RenderVisibilityRanges>,
(
Expand Down Expand Up @@ -907,7 +994,11 @@ pub fn specialize_material_meshes<M: Material>(
.or_default();

for (_, visible_entity) in visible_entities.iter::<Mesh3d>() {
let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
let Some(material_instance) = render_material_instances.instances.get(visible_entity)
else {
continue;
};
let Ok(material_asset_id) = material_instance.asset_id.try_typed::<M>() else {
continue;
};
let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
Expand All @@ -928,7 +1019,7 @@ pub fn specialize_material_meshes<M: Material>(
let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) else {
continue;
};
let Some(material) = render_materials.get(*material_asset_id) else {
let Some(material) = render_materials.get(material_asset_id) else {
continue;
};
let Some(material_bind_group) =
Expand Down Expand Up @@ -1004,7 +1095,7 @@ pub fn specialize_material_meshes<M: Material>(
pub fn queue_material_meshes<M: Material>(
render_materials: Res<RenderAssets<PreparedMaterial<M>>>,
render_mesh_instances: Res<RenderMeshInstances>,
render_material_instances: Res<RenderMaterialInstances<M>>,
render_material_instances: Res<RenderMaterialInstances>,
mesh_allocator: Res<MeshAllocator>,
gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
mut opaque_render_phases: ResMut<ViewBinnedRenderPhases<Opaque3d>>,
Expand Down Expand Up @@ -1054,14 +1145,18 @@ pub fn queue_material_meshes<M: Material>(
continue;
}

let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
let Some(material_instance) = render_material_instances.instances.get(visible_entity)
else {
continue;
};
let Ok(material_asset_id) = material_instance.asset_id.try_typed::<M>() else {
continue;
};
let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
else {
continue;
};
let Some(material) = render_materials.get(*material_asset_id) else {
let Some(material) = render_materials.get(material_asset_id) else {
continue;
};

Expand Down
28 changes: 15 additions & 13 deletions crates/bevy_pbr/src/meshlet/instance_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use super::{meshlet_mesh_manager::MeshletMeshManager, MeshletMesh, MeshletMesh3d
use crate::{
Material, MaterialBindingId, MeshFlags, MeshTransforms, MeshUniform, NotShadowCaster,
NotShadowReceiver, PreviousGlobalTransform, RenderMaterialBindings, RenderMaterialInstances,
RenderMeshMaterialIds, StandardMaterial,
StandardMaterial,
};
use bevy_asset::{AssetEvent, AssetId, AssetServer, Assets, UntypedAssetId};
use bevy_ecs::{
Expand Down Expand Up @@ -90,7 +90,7 @@ impl InstanceManager {
transform: &GlobalTransform,
previous_transform: Option<&PreviousGlobalTransform>,
render_layers: Option<&RenderLayers>,
mesh_material_ids: &RenderMeshMaterialIds,
mesh_material_ids: &RenderMaterialInstances,
render_material_bindings: &RenderMaterialBindings,
not_shadow_receiver: bool,
not_shadow_caster: bool,
Expand Down Expand Up @@ -193,7 +193,7 @@ pub fn extract_meshlet_mesh_entities(
mut instance_manager: ResMut<InstanceManager>,
// TODO: Replace main_world and system_state when Extract<ResMut<Assets<MeshletMesh>>> is possible
mut main_world: ResMut<MainWorld>,
mesh_material_ids: Res<RenderMeshMaterialIds>,
mesh_material_ids: Res<RenderMaterialInstances>,
render_material_bindings: Res<RenderMaterialBindings>,
mut system_state: Local<
Option<
Expand Down Expand Up @@ -275,20 +275,22 @@ pub fn extract_meshlet_mesh_entities(
/// and note that the material is used by at least one entity in the scene.
pub fn queue_material_meshlet_meshes<M: Material>(
mut instance_manager: ResMut<InstanceManager>,
render_material_instances: Res<RenderMaterialInstances<M>>,
render_material_instances: Res<RenderMaterialInstances>,
) {
let instance_manager = instance_manager.deref_mut();

for (i, (instance, _, _)) in instance_manager.instances.iter().enumerate() {
if let Some(material_asset_id) = render_material_instances.get(instance) {
if let Some(material_id) = instance_manager
.material_id_lookup
.get(&material_asset_id.untyped())
{
instance_manager
.material_ids_present_in_scene
.insert(*material_id);
instance_manager.instance_material_ids.get_mut()[i] = *material_id;
if let Some(material_instance) = render_material_instances.instances.get(instance) {
if let Ok(material_asset_id) = material_instance.asset_id.try_typed::<M>() {
if let Some(material_id) = instance_manager
.material_id_lookup
.get(&material_asset_id.untyped())
{
instance_manager
.material_ids_present_in_scene
.insert(*material_id);
instance_manager.instance_material_ids.get_mut()[i] = *material_id;
}
}
}
}
Expand Down
22 changes: 16 additions & 6 deletions crates/bevy_pbr/src/meshlet/material_pipeline_prepare.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ pub fn prepare_material_meshlet_meshes_main_opaque_pass<M: Material>(
material_pipeline: Res<MaterialPipeline<M>>,
mesh_pipeline: Res<MeshPipeline>,
render_materials: Res<RenderAssets<PreparedMaterial<M>>>,
render_material_instances: Res<RenderMaterialInstances<M>>,
render_material_instances: Res<RenderMaterialInstances>,
material_bind_group_allocator: Res<MaterialBindGroupAllocator<M>>,
asset_server: Res<AssetServer>,
mut mesh_vertex_buffer_layouts: ResMut<MeshVertexBufferLayouts>,
Expand Down Expand Up @@ -148,8 +148,13 @@ pub fn prepare_material_meshlet_meshes_main_opaque_pass<M: Material>(

view_key |= MeshPipelineKey::from_primitive_topology(PrimitiveTopology::TriangleList);

for material_id in render_material_instances.values().collect::<HashSet<_>>() {
let Some(material) = render_materials.get(*material_id) else {
for material_id in render_material_instances
.instances
.values()
.flat_map(|instance| instance.asset_id.try_typed::<M>().ok())
.collect::<HashSet<_>>()
{
let Some(material) = render_materials.get(material_id) else {
continue;
};
let Some(material_bind_group) =
Expand Down Expand Up @@ -256,7 +261,7 @@ pub fn prepare_material_meshlet_meshes_prepass<M: Material>(
pipeline_cache: Res<PipelineCache>,
prepass_pipeline: Res<PrepassPipeline<M>>,
render_materials: Res<RenderAssets<PreparedMaterial<M>>>,
render_material_instances: Res<RenderMaterialInstances<M>>,
render_material_instances: Res<RenderMaterialInstances>,
mut mesh_vertex_buffer_layouts: ResMut<MeshVertexBufferLayouts>,
material_bind_group_allocator: Res<MaterialBindGroupAllocator<M>>,
asset_server: Res<AssetServer>,
Expand Down Expand Up @@ -293,8 +298,13 @@ pub fn prepare_material_meshlet_meshes_prepass<M: Material>(

view_key |= MeshPipelineKey::from_primitive_topology(PrimitiveTopology::TriangleList);

for material_id in render_material_instances.values().collect::<HashSet<_>>() {
let Some(material) = render_materials.get(*material_id) else {
for material_id in render_material_instances
.instances
.values()
.flat_map(|instance| instance.asset_id.try_typed::<M>().ok())
.collect::<HashSet<_>>()
{
let Some(material) = render_materials.get(material_id) else {
continue;
};
let Some(material_bind_group) =
Expand Down
Loading
Loading