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
This commit is contained in:
Sarthak Singh 2025-05-31 01:06:56 +05:30 committed by François Mockers
parent 75d92f4434
commit f4bd56d26d

View File

@ -1048,63 +1048,14 @@ where
for (bindless_index, owned_binding_resource) in binding_resources.drain(..) { for (bindless_index, owned_binding_resource) in binding_resources.drain(..) {
let bindless_index = BindlessIndex(bindless_index); let bindless_index = BindlessIndex(bindless_index);
// If this is an other reference to an object we've already
// allocated, just bump its reference count. let pre_existing_slot = allocation_candidate
if let Some(pre_existing_resource_slot) = allocation_candidate
.pre_existing_resources .pre_existing_resources
.get(&bindless_index) .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;
}
// Otherwise, we need to insert it anew. // Otherwise, we need to insert it anew.
let binding_resource_id = BindingResourceId::from(&owned_binding_resource); 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) => { OwnedBindingResource::Buffer(buffer) => {
let slot = self let slot = self
.buffers .buffers
@ -1112,14 +1063,27 @@ where
.expect("Buffer binding array should exist") .expect("Buffer binding array should exist")
.insert(binding_resource_id, buffer); .insert(binding_resource_id, buffer);
allocated_resource_slots.insert(bindless_index, slot); 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) => { OwnedBindingResource::Data(data) => {
if pre_existing_slot.is_some() {
panic!("Data buffers can't be deduplicated")
}
let slot = self let slot = self
.data_buffers .data_buffers
.get_mut(&bindless_index) .get_mut(&bindless_index)
.expect("Data buffer binding array should exist") .expect("Data buffer binding array should exist")
.insert(&data); .insert(&data);
allocated_resource_slots.insert(bindless_index, slot); allocated_resource_slots.insert(bindless_index, slot);
false
} }
OwnedBindingResource::TextureView(texture_view_dimension, texture_view) => { OwnedBindingResource::TextureView(texture_view_dimension, texture_view) => {
let bindless_resource_type = BindlessResourceType::from(texture_view_dimension); let bindless_resource_type = BindlessResourceType::from(texture_view_dimension);
@ -1129,6 +1093,14 @@ where
.expect("Texture array should exist") .expect("Texture array should exist")
.insert(binding_resource_id, texture_view); .insert(binding_resource_id, texture_view);
allocated_resource_slots.insert(bindless_index, slot); 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) => { OwnedBindingResource::Sampler(sampler_binding_type, sampler) => {
let bindless_resource_type = BindlessResourceType::from(sampler_binding_type); let bindless_resource_type = BindlessResourceType::from(sampler_binding_type);
@ -1138,12 +1110,22 @@ where
.expect("Sampler should exist") .expect("Sampler should exist")
.insert(binding_resource_id, sampler); .insert(binding_resource_id, sampler);
allocated_resource_slots.insert(bindless_index, slot); 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. // Bump the allocated resource count.
if increment_allocated_resource_count {
self.allocated_resource_count += 1; self.allocated_resource_count += 1;
} }
}
allocated_resource_slots allocated_resource_slots
} }
@ -1626,8 +1608,20 @@ where
/// Inserts a bindless resource into a binding array and returns the index /// Inserts a bindless resource into a binding array and returns the index
/// of the slot it was inserted into. /// of the slot it was inserted into.
fn insert(&mut self, binding_resource_id: BindingResourceId, resource: R) -> u32 { fn insert(&mut self, binding_resource_id: BindingResourceId, resource: R) -> u32 {
match self.resource_to_slot.entry(binding_resource_id) {
bevy_platform::collections::hash_map::Entry::Occupied(o) => {
let slot = *o.get();
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); let slot = self.free_slots.pop().unwrap_or(self.len);
self.resource_to_slot.insert(binding_resource_id, slot); v.insert(slot);
if self.bindings.len() < slot as usize + 1 { if self.bindings.len() < slot as usize + 1 {
self.bindings.resize_with(slot as usize + 1, || None); self.bindings.resize_with(slot as usize + 1, || None);
@ -1637,6 +1631,8 @@ where
self.len += 1; self.len += 1;
slot slot
} }
}
}
/// Removes a reference to an object from the slot. /// Removes a reference to an object from the slot.
/// ///