From f4bd56d26d0d02b227fb0faceb6d4519649e16f3 Mon Sep 17 00:00:00 2001 From: Sarthak Singh <35749450+SarthakSingh31@users.noreply.github.com> Date: Sat, 31 May 2025 01:06:56 +0530 Subject: [PATCH] Fixed memory leak in bindless material (#19041) # Objective Fixed #19035. Fixed #18882. It consisted of two different bugs: - The allocations where being incremented even when a Data binding was created. - The ref counting on the binding was broken. ## Solution - Stopped incrementing the allocations when a data binding was created. - Rewrote the ref counting code to more reliably track the ref count. ## Testing Tested my fix for 10 minutes with the `examples/3d/animated_material.rs` example. I changed the example to spawn 51x51 meshes instead of 3x3 meshes to heighten the effects of the bug. My branch: (After 10 minutes of running the modified example) GPU: 172 MB CPU: ~700 MB Main branch: (After 2 minutes of running the modified example, my computer started to stutter so I had to end it early) GPU: 376 MB CPU: ~1300 MB --- crates/bevy_pbr/src/material_bind_groups.rs | 122 ++++++++++---------- 1 file changed, 59 insertions(+), 63 deletions(-) diff --git a/crates/bevy_pbr/src/material_bind_groups.rs b/crates/bevy_pbr/src/material_bind_groups.rs index b539d2098f..735bc77c99 100644 --- a/crates/bevy_pbr/src/material_bind_groups.rs +++ b/crates/bevy_pbr/src/material_bind_groups.rs @@ -1048,63 +1048,14 @@ where for (bindless_index, owned_binding_resource) in binding_resources.drain(..) { let bindless_index = BindlessIndex(bindless_index); - // If this is an other reference to an object we've already - // allocated, just bump its reference count. - if let Some(pre_existing_resource_slot) = allocation_candidate + + let pre_existing_slot = allocation_candidate .pre_existing_resources - .get(&bindless_index) - { - allocated_resource_slots.insert(bindless_index, *pre_existing_resource_slot); - - match owned_binding_resource { - OwnedBindingResource::Buffer(_) => { - self.buffers - .get_mut(&bindless_index) - .expect("Buffer binding array should exist") - .bindings - .get_mut(*pre_existing_resource_slot as usize) - .and_then(|binding| binding.as_mut()) - .expect("Slot should exist") - .ref_count += 1; - } - - OwnedBindingResource::Data(_) => { - panic!("Data buffers can't be deduplicated") - } - - OwnedBindingResource::TextureView(texture_view_dimension, _) => { - let bindless_resource_type = - BindlessResourceType::from(texture_view_dimension); - self.textures - .get_mut(&bindless_resource_type) - .expect("Texture binding array should exist") - .bindings - .get_mut(*pre_existing_resource_slot as usize) - .and_then(|binding| binding.as_mut()) - .expect("Slot should exist") - .ref_count += 1; - } - - OwnedBindingResource::Sampler(sampler_binding_type, _) => { - let bindless_resource_type = - BindlessResourceType::from(sampler_binding_type); - self.samplers - .get_mut(&bindless_resource_type) - .expect("Sampler binding array should exist") - .bindings - .get_mut(*pre_existing_resource_slot as usize) - .and_then(|binding| binding.as_mut()) - .expect("Slot should exist") - .ref_count += 1; - } - } - - continue; - } + .get(&bindless_index); // Otherwise, we need to insert it anew. let binding_resource_id = BindingResourceId::from(&owned_binding_resource); - match owned_binding_resource { + let increment_allocated_resource_count = match owned_binding_resource { OwnedBindingResource::Buffer(buffer) => { let slot = self .buffers @@ -1112,14 +1063,27 @@ where .expect("Buffer binding array should exist") .insert(binding_resource_id, buffer); allocated_resource_slots.insert(bindless_index, slot); + + if let Some(pre_existing_slot) = pre_existing_slot { + assert_eq!(*pre_existing_slot, slot); + + false + } else { + true + } } OwnedBindingResource::Data(data) => { + if pre_existing_slot.is_some() { + panic!("Data buffers can't be deduplicated") + } + let slot = self .data_buffers .get_mut(&bindless_index) .expect("Data buffer binding array should exist") .insert(&data); allocated_resource_slots.insert(bindless_index, slot); + false } OwnedBindingResource::TextureView(texture_view_dimension, texture_view) => { let bindless_resource_type = BindlessResourceType::from(texture_view_dimension); @@ -1129,6 +1093,14 @@ where .expect("Texture array should exist") .insert(binding_resource_id, texture_view); allocated_resource_slots.insert(bindless_index, slot); + + if let Some(pre_existing_slot) = pre_existing_slot { + assert_eq!(*pre_existing_slot, slot); + + false + } else { + true + } } OwnedBindingResource::Sampler(sampler_binding_type, sampler) => { let bindless_resource_type = BindlessResourceType::from(sampler_binding_type); @@ -1138,11 +1110,21 @@ where .expect("Sampler should exist") .insert(binding_resource_id, sampler); allocated_resource_slots.insert(bindless_index, slot); + + if let Some(pre_existing_slot) = pre_existing_slot { + assert_eq!(*pre_existing_slot, slot); + + false + } else { + true + } } - } + }; // Bump the allocated resource count. - self.allocated_resource_count += 1; + if increment_allocated_resource_count { + self.allocated_resource_count += 1; + } } allocated_resource_slots @@ -1626,16 +1608,30 @@ where /// Inserts a bindless resource into a binding array and returns the index /// of the slot it was inserted into. fn insert(&mut self, binding_resource_id: BindingResourceId, resource: R) -> u32 { - let slot = self.free_slots.pop().unwrap_or(self.len); - self.resource_to_slot.insert(binding_resource_id, slot); + match self.resource_to_slot.entry(binding_resource_id) { + bevy_platform::collections::hash_map::Entry::Occupied(o) => { + let slot = *o.get(); - if self.bindings.len() < slot as usize + 1 { - self.bindings.resize_with(slot as usize + 1, || None); + self.bindings[slot as usize] + .as_mut() + .expect("A slot in the resource_to_slot map should have a value") + .ref_count += 1; + + slot + } + bevy_platform::collections::hash_map::Entry::Vacant(v) => { + let slot = self.free_slots.pop().unwrap_or(self.len); + v.insert(slot); + + if self.bindings.len() < slot as usize + 1 { + self.bindings.resize_with(slot as usize + 1, || None); + } + self.bindings[slot as usize] = Some(MaterialBindlessBinding::new(resource)); + + self.len += 1; + slot + } } - self.bindings[slot as usize] = Some(MaterialBindlessBinding::new(resource)); - - self.len += 1; - slot } /// Removes a reference to an object from the slot.