5e8dec8a90
# Objective Fixes https://github.com/bevyengine/bevy/issues/16586. ## Solution - Free meshes before allocating new ones (so hopefully the existing allocation is used, but it's not guaranteed since it might end up getting used by a smaller mesh first). - Keep track of modified render assets, and have the mesh allocator free their allocations. - Cleaned up some render asset code to make it more understandable, since it took me several minutes to reverse engineer/remember how it was supposed to work. Long term we'll probably want to explicitly reusing allocations for modified meshes that haven't grown in size, or do delta uploads using a compute shader or something, but this is an easy fix for the near term. ## Testing Ran the example provided in the issue. No crash after a few minutes, and memory usage remains steady.
1021 lines
36 KiB
Rust
1021 lines
36 KiB
Rust
//! Manages mesh vertex and index buffers.
|
||
|
||
use alloc::vec::Vec;
|
||
use core::{
|
||
fmt::{self, Display, Formatter},
|
||
ops::Range,
|
||
};
|
||
use nonmax::NonMaxU32;
|
||
|
||
use bevy_app::{App, Plugin};
|
||
use bevy_asset::AssetId;
|
||
use bevy_derive::{Deref, DerefMut};
|
||
use bevy_ecs::{
|
||
resource::Resource,
|
||
schedule::IntoScheduleConfigs as _,
|
||
system::{Res, ResMut},
|
||
world::{FromWorld, World},
|
||
};
|
||
use bevy_platform_support::collections::{hash_map::Entry, HashMap, HashSet};
|
||
use bevy_utils::default;
|
||
use offset_allocator::{Allocation, Allocator};
|
||
use tracing::error;
|
||
use wgpu::{
|
||
BufferDescriptor, BufferSize, BufferUsages, CommandEncoderDescriptor, DownlevelFlags,
|
||
COPY_BUFFER_ALIGNMENT,
|
||
};
|
||
|
||
use crate::{
|
||
mesh::{Indices, Mesh, MeshVertexBufferLayouts, RenderMesh},
|
||
render_asset::{prepare_assets, ExtractedAssets},
|
||
render_resource::Buffer,
|
||
renderer::{RenderAdapter, RenderDevice, RenderQueue},
|
||
Render, RenderApp, RenderSet,
|
||
};
|
||
|
||
/// A plugin that manages GPU memory for mesh data.
|
||
pub struct MeshAllocatorPlugin;
|
||
|
||
/// Manages the assignment of mesh data to GPU buffers.
|
||
///
|
||
/// The Bevy renderer tries to pack vertex and index data for multiple meshes
|
||
/// together so that multiple meshes can be drawn back-to-back without any
|
||
/// rebinding. This resource manages these buffers.
|
||
///
|
||
/// Within each slab, or hardware buffer, the underlying allocation algorithm is
|
||
/// [`offset-allocator`], a Rust port of Sebastian Aaltonen's hard-real-time C++
|
||
/// `OffsetAllocator`. Slabs start small and then grow as their contents fill
|
||
/// up, up to a maximum size limit. To reduce fragmentation, vertex and index
|
||
/// buffers that are too large bypass this system and receive their own buffers.
|
||
///
|
||
/// The [`MeshAllocatorSettings`] allows you to tune the behavior of the
|
||
/// allocator for better performance with your application. Most applications
|
||
/// won't need to change the settings from their default values.
|
||
#[derive(Resource)]
|
||
pub struct MeshAllocator {
|
||
/// Holds all buffers and allocators.
|
||
slabs: HashMap<SlabId, Slab>,
|
||
|
||
/// Maps a layout to the slabs that hold elements of that layout.
|
||
///
|
||
/// This is used when allocating, so that we can find the appropriate slab
|
||
/// to place an object in.
|
||
slab_layouts: HashMap<ElementLayout, Vec<SlabId>>,
|
||
|
||
/// Maps mesh asset IDs to the ID of the slabs that hold their vertex data.
|
||
mesh_id_to_vertex_slab: HashMap<AssetId<Mesh>, SlabId>,
|
||
|
||
/// Maps mesh asset IDs to the ID of the slabs that hold their index data.
|
||
mesh_id_to_index_slab: HashMap<AssetId<Mesh>, SlabId>,
|
||
|
||
/// The next slab ID to assign.
|
||
next_slab_id: SlabId,
|
||
|
||
/// Whether we can pack multiple vertex arrays into a single slab on this
|
||
/// platform.
|
||
///
|
||
/// This corresponds to [`DownlevelFlags::BASE_VERTEX`], which is unset on
|
||
/// WebGL 2. On this platform, we must give each vertex array its own
|
||
/// buffer, because we can't adjust the first vertex when we perform a draw.
|
||
general_vertex_slabs_supported: bool,
|
||
}
|
||
|
||
/// Tunable parameters that customize the behavior of the allocator.
|
||
///
|
||
/// Generally, these parameters adjust the tradeoff between memory fragmentation
|
||
/// and performance. You can adjust them as desired for your application. Most
|
||
/// applications can stick with the default values.
|
||
#[derive(Resource)]
|
||
pub struct MeshAllocatorSettings {
|
||
/// The minimum size of a slab (hardware buffer), in bytes.
|
||
///
|
||
/// The default value is 1 MiB.
|
||
pub min_slab_size: u64,
|
||
|
||
/// The maximum size of a slab (hardware buffer), in bytes.
|
||
///
|
||
/// When a slab reaches this limit, a new slab is created.
|
||
///
|
||
/// The default value is 512 MiB.
|
||
pub max_slab_size: u64,
|
||
|
||
/// The maximum size of vertex or index data that can be placed in a general
|
||
/// slab, in bytes.
|
||
///
|
||
/// If a mesh has vertex or index data that exceeds this size limit, that
|
||
/// data is placed in its own slab. This reduces fragmentation, but incurs
|
||
/// more CPU-side binding overhead when drawing the mesh.
|
||
///
|
||
/// The default value is 256 MiB.
|
||
pub large_threshold: u64,
|
||
|
||
/// The factor by which we scale a slab when growing it.
|
||
///
|
||
/// This value must be greater than 1. Higher values result in more
|
||
/// fragmentation but fewer expensive copy operations when growing the
|
||
/// buffer.
|
||
///
|
||
/// The default value is 1.5.
|
||
pub growth_factor: f64,
|
||
}
|
||
|
||
impl Default for MeshAllocatorSettings {
|
||
fn default() -> Self {
|
||
Self {
|
||
// 1 MiB
|
||
min_slab_size: 1024 * 1024,
|
||
// 512 MiB
|
||
max_slab_size: 1024 * 1024 * 512,
|
||
// 256 MiB
|
||
large_threshold: 1024 * 1024 * 256,
|
||
// 1.5× growth
|
||
growth_factor: 1.5,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The hardware buffer that mesh data lives in, as well as the range within
|
||
/// that buffer.
|
||
pub struct MeshBufferSlice<'a> {
|
||
/// The buffer that the mesh data resides in.
|
||
pub buffer: &'a Buffer,
|
||
|
||
/// The range of elements within this buffer that the mesh data resides in,
|
||
/// measured in elements.
|
||
///
|
||
/// This is not a byte range; it's an element range. For vertex data, this
|
||
/// is measured in increments of a single vertex. (Thus, if a vertex is 32
|
||
/// bytes long, then this range is in units of 32 bytes each.) For index
|
||
/// data, this is measured in increments of a single index value (2 or 4
|
||
/// bytes). Draw commands generally take their ranges in elements, not
|
||
/// bytes, so this is the most convenient unit in this case.
|
||
pub range: Range<u32>,
|
||
}
|
||
|
||
/// The index of a single slab.
|
||
#[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
|
||
#[repr(transparent)]
|
||
pub struct SlabId(pub NonMaxU32);
|
||
|
||
/// Data for a single slab.
|
||
enum Slab {
|
||
/// A slab that can contain multiple objects.
|
||
General(GeneralSlab),
|
||
/// A slab that contains a single object.
|
||
LargeObject(LargeObjectSlab),
|
||
}
|
||
|
||
/// A resizable slab that can contain multiple objects.
|
||
///
|
||
/// This is the normal type of slab used for objects that are below the
|
||
/// [`MeshAllocatorSettings::large_threshold`]. Slabs are divided into *slots*,
|
||
/// which are described in detail in the [`ElementLayout`] documentation.
|
||
struct GeneralSlab {
|
||
/// The [`Allocator`] that manages the objects in this slab.
|
||
allocator: Allocator,
|
||
|
||
/// The GPU buffer that backs this slab.
|
||
///
|
||
/// This may be `None` if the buffer hasn't been created yet. We delay
|
||
/// creation of buffers until allocating all the meshes for a single frame,
|
||
/// so that we don't needlessly create and resize buffers when many meshes
|
||
/// load all at once.
|
||
buffer: Option<Buffer>,
|
||
|
||
/// Allocations that are on the GPU.
|
||
///
|
||
/// The range is in slots.
|
||
resident_allocations: HashMap<AssetId<Mesh>, SlabAllocation>,
|
||
|
||
/// Allocations that are waiting to be uploaded to the GPU.
|
||
///
|
||
/// The range is in slots.
|
||
pending_allocations: HashMap<AssetId<Mesh>, SlabAllocation>,
|
||
|
||
/// The layout of a single element (vertex or index).
|
||
element_layout: ElementLayout,
|
||
|
||
/// The size of this slab in slots.
|
||
current_slot_capacity: u32,
|
||
}
|
||
|
||
/// A slab that contains a single object.
|
||
///
|
||
/// Typically, this is for objects that exceed the
|
||
/// [`MeshAllocatorSettings::large_threshold`]. This is also for objects that
|
||
/// would ordinarily receive their own slab but can't because of platform
|
||
/// limitations, most notably vertex arrays on WebGL 2.
|
||
struct LargeObjectSlab {
|
||
/// The GPU buffer that backs this slab.
|
||
///
|
||
/// This may be `None` if the buffer hasn't been created yet.
|
||
buffer: Option<Buffer>,
|
||
|
||
/// The layout of a single element (vertex or index).
|
||
element_layout: ElementLayout,
|
||
}
|
||
|
||
/// The type of element that a slab can store.
|
||
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
|
||
enum ElementClass {
|
||
/// Data for a vertex.
|
||
Vertex,
|
||
/// A vertex index.
|
||
Index,
|
||
}
|
||
|
||
/// The results of [`GeneralSlab::grow_if_necessary`].
|
||
enum SlabGrowthResult {
|
||
/// The mesh data already fits in the slab; the slab doesn't need to grow.
|
||
NoGrowthNeeded,
|
||
/// The slab needed to grow.
|
||
///
|
||
/// The [`SlabToReallocate`] contains the old capacity of the slab.
|
||
NeededGrowth(SlabToReallocate),
|
||
/// The slab wanted to grow but couldn't because it hit its maximum size.
|
||
CantGrow,
|
||
}
|
||
|
||
/// Information about the size of individual elements (vertices or indices)
|
||
/// within a slab.
|
||
///
|
||
/// Slab objects are allocated in units of *slots*. Usually, each element takes
|
||
/// up one slot, and so elements and slots are equivalent. Occasionally,
|
||
/// however, a slot may consist of 2 or even 4 elements. This occurs when the
|
||
/// size of an element isn't divisible by [`COPY_BUFFER_ALIGNMENT`]. When we
|
||
/// resize buffers, we perform GPU-to-GPU copies to shuffle the existing
|
||
/// elements into their new positions, and such copies must be on
|
||
/// [`COPY_BUFFER_ALIGNMENT`] boundaries. Slots solve this problem by
|
||
/// guaranteeing that the size of an allocation quantum is divisible by both the
|
||
/// size of an element and [`COPY_BUFFER_ALIGNMENT`], so we can relocate it
|
||
/// freely.
|
||
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
|
||
struct ElementLayout {
|
||
/// Either a vertex or an index.
|
||
class: ElementClass,
|
||
|
||
/// The size in bytes of a single element (vertex or index).
|
||
size: u64,
|
||
|
||
/// The number of elements that make up a single slot.
|
||
///
|
||
/// Usually, this is 1, but it can be different if [`ElementLayout::size`]
|
||
/// isn't divisible by 4. See the comment in [`ElementLayout`] for more
|
||
/// details.
|
||
elements_per_slot: u32,
|
||
}
|
||
|
||
/// The location of an allocation and the slab it's contained in.
|
||
struct MeshAllocation {
|
||
/// The ID of the slab.
|
||
slab_id: SlabId,
|
||
/// Holds the actual allocation.
|
||
slab_allocation: SlabAllocation,
|
||
}
|
||
|
||
/// An allocation within a slab.
|
||
#[derive(Clone)]
|
||
struct SlabAllocation {
|
||
/// The actual [`Allocator`] handle, needed to free the allocation.
|
||
allocation: Allocation,
|
||
/// The number of slots that this allocation takes up.
|
||
slot_count: u32,
|
||
}
|
||
|
||
/// Holds information about all slabs scheduled to be allocated or reallocated.
|
||
#[derive(Default, Deref, DerefMut)]
|
||
struct SlabsToReallocate(HashMap<SlabId, SlabToReallocate>);
|
||
|
||
/// Holds information about a slab that's scheduled to be allocated or
|
||
/// reallocated.
|
||
#[derive(Default)]
|
||
struct SlabToReallocate {
|
||
/// The capacity of the slab before we decided to grow it.
|
||
old_slot_capacity: u32,
|
||
}
|
||
|
||
impl Display for SlabId {
|
||
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
|
||
self.0.fmt(f)
|
||
}
|
||
}
|
||
|
||
impl Plugin for MeshAllocatorPlugin {
|
||
fn build(&self, app: &mut App) {
|
||
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
|
||
return;
|
||
};
|
||
|
||
render_app
|
||
.init_resource::<MeshAllocatorSettings>()
|
||
.add_systems(
|
||
Render,
|
||
allocate_and_free_meshes
|
||
.in_set(RenderSet::PrepareAssets)
|
||
.before(prepare_assets::<RenderMesh>),
|
||
);
|
||
}
|
||
|
||
fn finish(&self, app: &mut App) {
|
||
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
|
||
return;
|
||
};
|
||
|
||
// The `RenderAdapter` isn't available until now, so we can't do this in
|
||
// [`Plugin::build`].
|
||
render_app.init_resource::<MeshAllocator>();
|
||
}
|
||
}
|
||
|
||
impl FromWorld for MeshAllocator {
|
||
fn from_world(world: &mut World) -> Self {
|
||
// Note whether we're on WebGL 2. In this case, we must give every
|
||
// vertex array its own slab.
|
||
let render_adapter = world.resource::<RenderAdapter>();
|
||
let general_vertex_slabs_supported = render_adapter
|
||
.get_downlevel_capabilities()
|
||
.flags
|
||
.contains(DownlevelFlags::BASE_VERTEX);
|
||
|
||
Self {
|
||
slabs: HashMap::default(),
|
||
slab_layouts: HashMap::default(),
|
||
mesh_id_to_vertex_slab: HashMap::default(),
|
||
mesh_id_to_index_slab: HashMap::default(),
|
||
next_slab_id: default(),
|
||
general_vertex_slabs_supported,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// A system that processes newly-extracted or newly-removed meshes and writes
|
||
/// their data into buffers or frees their data as appropriate.
|
||
pub fn allocate_and_free_meshes(
|
||
mut mesh_allocator: ResMut<MeshAllocator>,
|
||
mesh_allocator_settings: Res<MeshAllocatorSettings>,
|
||
extracted_meshes: Res<ExtractedAssets<RenderMesh>>,
|
||
mut mesh_vertex_buffer_layouts: ResMut<MeshVertexBufferLayouts>,
|
||
render_device: Res<RenderDevice>,
|
||
render_queue: Res<RenderQueue>,
|
||
) {
|
||
// Process removed or modified meshes.
|
||
mesh_allocator.free_meshes(&extracted_meshes);
|
||
|
||
// Process newly-added or modified meshes.
|
||
mesh_allocator.allocate_meshes(
|
||
&mesh_allocator_settings,
|
||
&extracted_meshes,
|
||
&mut mesh_vertex_buffer_layouts,
|
||
&render_device,
|
||
&render_queue,
|
||
);
|
||
}
|
||
|
||
impl MeshAllocator {
|
||
/// Returns the buffer and range within that buffer of the vertex data for
|
||
/// the mesh with the given ID.
|
||
///
|
||
/// If the mesh wasn't allocated, returns None.
|
||
pub fn mesh_vertex_slice(&self, mesh_id: &AssetId<Mesh>) -> Option<MeshBufferSlice> {
|
||
self.mesh_slice_in_slab(mesh_id, *self.mesh_id_to_vertex_slab.get(mesh_id)?)
|
||
}
|
||
|
||
/// Returns the buffer and range within that buffer of the index data for
|
||
/// the mesh with the given ID.
|
||
///
|
||
/// If the mesh has no index data or wasn't allocated, returns None.
|
||
pub fn mesh_index_slice(&self, mesh_id: &AssetId<Mesh>) -> Option<MeshBufferSlice> {
|
||
self.mesh_slice_in_slab(mesh_id, *self.mesh_id_to_index_slab.get(mesh_id)?)
|
||
}
|
||
|
||
/// Returns the IDs of the vertex buffer and index buffer respectively for
|
||
/// the mesh with the given ID.
|
||
///
|
||
/// If the mesh wasn't allocated, or has no index data in the case of the
|
||
/// index buffer, the corresponding element in the returned tuple will be
|
||
/// None.
|
||
pub fn mesh_slabs(&self, mesh_id: &AssetId<Mesh>) -> (Option<SlabId>, Option<SlabId>) {
|
||
(
|
||
self.mesh_id_to_vertex_slab.get(mesh_id).cloned(),
|
||
self.mesh_id_to_index_slab.get(mesh_id).cloned(),
|
||
)
|
||
}
|
||
|
||
/// Given a slab and a mesh with data located with it, returns the buffer
|
||
/// and range of that mesh data within the slab.
|
||
fn mesh_slice_in_slab(
|
||
&self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
slab_id: SlabId,
|
||
) -> Option<MeshBufferSlice> {
|
||
match self.slabs.get(&slab_id)? {
|
||
Slab::General(general_slab) => {
|
||
let slab_allocation = general_slab.resident_allocations.get(mesh_id)?;
|
||
Some(MeshBufferSlice {
|
||
buffer: general_slab.buffer.as_ref()?,
|
||
range: (slab_allocation.allocation.offset
|
||
* general_slab.element_layout.elements_per_slot)
|
||
..((slab_allocation.allocation.offset + slab_allocation.slot_count)
|
||
* general_slab.element_layout.elements_per_slot),
|
||
})
|
||
}
|
||
|
||
Slab::LargeObject(large_object_slab) => {
|
||
let buffer = large_object_slab.buffer.as_ref()?;
|
||
Some(MeshBufferSlice {
|
||
buffer,
|
||
range: 0..((buffer.size() / large_object_slab.element_layout.size) as u32),
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Processes newly-loaded meshes, allocating room in the slabs for their
|
||
/// mesh data and performing upload operations as appropriate.
|
||
fn allocate_meshes(
|
||
&mut self,
|
||
mesh_allocator_settings: &MeshAllocatorSettings,
|
||
extracted_meshes: &ExtractedAssets<RenderMesh>,
|
||
mesh_vertex_buffer_layouts: &mut MeshVertexBufferLayouts,
|
||
render_device: &RenderDevice,
|
||
render_queue: &RenderQueue,
|
||
) {
|
||
let mut slabs_to_grow = SlabsToReallocate::default();
|
||
|
||
// Allocate.
|
||
for (mesh_id, mesh) in &extracted_meshes.extracted {
|
||
// Allocate vertex data. Note that we can only pack mesh vertex data
|
||
// together if the platform supports it.
|
||
let vertex_element_layout = ElementLayout::vertex(mesh_vertex_buffer_layouts, mesh);
|
||
if self.general_vertex_slabs_supported {
|
||
self.allocate(
|
||
mesh_id,
|
||
mesh.get_vertex_buffer_size() as u64,
|
||
vertex_element_layout,
|
||
&mut slabs_to_grow,
|
||
mesh_allocator_settings,
|
||
);
|
||
} else {
|
||
self.allocate_large(mesh_id, vertex_element_layout);
|
||
}
|
||
|
||
// Allocate index data.
|
||
if let (Some(index_buffer_data), Some(index_element_layout)) =
|
||
(mesh.get_index_buffer_bytes(), ElementLayout::index(mesh))
|
||
{
|
||
self.allocate(
|
||
mesh_id,
|
||
index_buffer_data.len() as u64,
|
||
index_element_layout,
|
||
&mut slabs_to_grow,
|
||
mesh_allocator_settings,
|
||
);
|
||
}
|
||
}
|
||
|
||
// Perform growth.
|
||
for (slab_id, slab_to_grow) in slabs_to_grow.0 {
|
||
self.reallocate_slab(render_device, render_queue, slab_id, slab_to_grow);
|
||
}
|
||
|
||
// Copy new mesh data in.
|
||
for (mesh_id, mesh) in &extracted_meshes.extracted {
|
||
self.copy_mesh_vertex_data(mesh_id, mesh, render_device, render_queue);
|
||
self.copy_mesh_index_data(mesh_id, mesh, render_device, render_queue);
|
||
}
|
||
}
|
||
|
||
/// Copies vertex array data from a mesh into the appropriate spot in the
|
||
/// slab.
|
||
fn copy_mesh_vertex_data(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
mesh: &Mesh,
|
||
render_device: &RenderDevice,
|
||
render_queue: &RenderQueue,
|
||
) {
|
||
let Some(&slab_id) = self.mesh_id_to_vertex_slab.get(mesh_id) else {
|
||
return;
|
||
};
|
||
|
||
// Call the generic function.
|
||
self.copy_element_data(
|
||
mesh_id,
|
||
mesh.get_vertex_buffer_size(),
|
||
|slice| mesh.write_packed_vertex_buffer_data(slice),
|
||
BufferUsages::VERTEX,
|
||
slab_id,
|
||
render_device,
|
||
render_queue,
|
||
);
|
||
}
|
||
|
||
/// Copies index array data from a mesh into the appropriate spot in the
|
||
/// slab.
|
||
fn copy_mesh_index_data(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
mesh: &Mesh,
|
||
render_device: &RenderDevice,
|
||
render_queue: &RenderQueue,
|
||
) {
|
||
let Some(&slab_id) = self.mesh_id_to_index_slab.get(mesh_id) else {
|
||
return;
|
||
};
|
||
let Some(index_data) = mesh.get_index_buffer_bytes() else {
|
||
return;
|
||
};
|
||
|
||
// Call the generic function.
|
||
self.copy_element_data(
|
||
mesh_id,
|
||
index_data.len(),
|
||
|slice| slice.copy_from_slice(index_data),
|
||
BufferUsages::INDEX,
|
||
slab_id,
|
||
render_device,
|
||
render_queue,
|
||
);
|
||
}
|
||
|
||
/// A generic function that copies either vertex or index data into a slab.
|
||
fn copy_element_data(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
len: usize,
|
||
fill_data: impl Fn(&mut [u8]),
|
||
buffer_usages: BufferUsages,
|
||
slab_id: SlabId,
|
||
render_device: &RenderDevice,
|
||
render_queue: &RenderQueue,
|
||
) {
|
||
let Some(slab) = self.slabs.get_mut(&slab_id) else {
|
||
return;
|
||
};
|
||
|
||
match *slab {
|
||
Slab::General(ref mut general_slab) => {
|
||
let (Some(buffer), Some(allocated_range)) = (
|
||
&general_slab.buffer,
|
||
general_slab.pending_allocations.remove(mesh_id),
|
||
) else {
|
||
return;
|
||
};
|
||
|
||
let slot_size = general_slab.element_layout.slot_size();
|
||
|
||
// round up size to a multiple of the slot size to satisfy wgpu alignment requirements
|
||
if let Some(size) = BufferSize::new((len as u64).next_multiple_of(slot_size)) {
|
||
// Write the data in.
|
||
if let Some(mut buffer) = render_queue.write_buffer_with(
|
||
buffer,
|
||
allocated_range.allocation.offset as u64 * slot_size,
|
||
size,
|
||
) {
|
||
let slice = &mut buffer.as_mut()[..len];
|
||
fill_data(slice);
|
||
}
|
||
}
|
||
|
||
// Mark the allocation as resident.
|
||
general_slab
|
||
.resident_allocations
|
||
.insert(*mesh_id, allocated_range);
|
||
}
|
||
|
||
Slab::LargeObject(ref mut large_object_slab) => {
|
||
debug_assert!(large_object_slab.buffer.is_none());
|
||
|
||
// Create the buffer and its data in one go.
|
||
let buffer = render_device.create_buffer(&BufferDescriptor {
|
||
label: Some(&format!(
|
||
"large mesh slab {} ({}buffer)",
|
||
slab_id,
|
||
buffer_usages_to_str(buffer_usages)
|
||
)),
|
||
size: len as u64,
|
||
usage: buffer_usages | BufferUsages::COPY_DST,
|
||
mapped_at_creation: true,
|
||
});
|
||
{
|
||
let slice = &mut buffer.slice(..).get_mapped_range_mut()[..len];
|
||
fill_data(slice);
|
||
}
|
||
buffer.unmap();
|
||
large_object_slab.buffer = Some(buffer);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Frees allocations for meshes that were removed or modified this frame.
|
||
fn free_meshes(&mut self, extracted_meshes: &ExtractedAssets<RenderMesh>) {
|
||
let mut empty_slabs = <HashSet<_>>::default();
|
||
|
||
// TODO: Consider explicitly reusing allocations for changed meshes of the same size
|
||
let meshes_to_free = extracted_meshes
|
||
.removed
|
||
.iter()
|
||
.chain(extracted_meshes.modified.iter());
|
||
|
||
for mesh_id in meshes_to_free {
|
||
if let Some(slab_id) = self.mesh_id_to_vertex_slab.remove(mesh_id) {
|
||
self.free_allocation_in_slab(mesh_id, slab_id, &mut empty_slabs);
|
||
}
|
||
if let Some(slab_id) = self.mesh_id_to_index_slab.remove(mesh_id) {
|
||
self.free_allocation_in_slab(mesh_id, slab_id, &mut empty_slabs);
|
||
}
|
||
}
|
||
|
||
for empty_slab in empty_slabs {
|
||
self.slab_layouts.values_mut().for_each(|slab_ids| {
|
||
let idx = slab_ids.iter().position(|&slab_id| slab_id == empty_slab);
|
||
if let Some(idx) = idx {
|
||
slab_ids.remove(idx);
|
||
}
|
||
});
|
||
self.slabs.remove(&empty_slab);
|
||
}
|
||
}
|
||
|
||
/// Given a slab and the ID of a mesh containing data in it, marks the
|
||
/// allocation as free.
|
||
///
|
||
/// If this results in the slab becoming empty, this function adds the slab
|
||
/// to the `empty_slabs` set.
|
||
fn free_allocation_in_slab(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
slab_id: SlabId,
|
||
empty_slabs: &mut HashSet<SlabId>,
|
||
) {
|
||
let Some(slab) = self.slabs.get_mut(&slab_id) else {
|
||
return;
|
||
};
|
||
|
||
match *slab {
|
||
Slab::General(ref mut general_slab) => {
|
||
let Some(slab_allocation) = general_slab
|
||
.resident_allocations
|
||
.remove(mesh_id)
|
||
.or_else(|| general_slab.pending_allocations.remove(mesh_id))
|
||
else {
|
||
return;
|
||
};
|
||
|
||
general_slab.allocator.free(slab_allocation.allocation);
|
||
|
||
if general_slab.is_empty() {
|
||
empty_slabs.insert(slab_id);
|
||
}
|
||
}
|
||
Slab::LargeObject(_) => {
|
||
empty_slabs.insert(slab_id);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Allocates space for mesh data with the given byte size and layout in the
|
||
/// appropriate slab, creating that slab if necessary.
|
||
fn allocate(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
data_byte_len: u64,
|
||
layout: ElementLayout,
|
||
slabs_to_grow: &mut SlabsToReallocate,
|
||
settings: &MeshAllocatorSettings,
|
||
) {
|
||
let data_element_count = data_byte_len.div_ceil(layout.size) as u32;
|
||
let data_slot_count = data_element_count.div_ceil(layout.elements_per_slot);
|
||
|
||
// If the mesh data is too large for a slab, give it a slab of its own.
|
||
if data_slot_count as u64 * layout.slot_size()
|
||
>= settings.large_threshold.min(settings.max_slab_size)
|
||
{
|
||
self.allocate_large(mesh_id, layout);
|
||
} else {
|
||
self.allocate_general(mesh_id, data_slot_count, layout, slabs_to_grow, settings);
|
||
}
|
||
}
|
||
|
||
/// Allocates space for mesh data with the given slot size and layout in the
|
||
/// appropriate general slab.
|
||
fn allocate_general(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
data_slot_count: u32,
|
||
layout: ElementLayout,
|
||
slabs_to_grow: &mut SlabsToReallocate,
|
||
settings: &MeshAllocatorSettings,
|
||
) {
|
||
let candidate_slabs = self.slab_layouts.entry(layout).or_default();
|
||
|
||
// Loop through the slabs that accept elements of the appropriate type
|
||
// and try to allocate the mesh inside them. We go with the first one
|
||
// that succeeds.
|
||
let mut mesh_allocation = None;
|
||
for &slab_id in &*candidate_slabs {
|
||
let Some(Slab::General(slab)) = self.slabs.get_mut(&slab_id) else {
|
||
unreachable!("Slab not found")
|
||
};
|
||
|
||
let Some(allocation) = slab.allocator.allocate(data_slot_count) else {
|
||
continue;
|
||
};
|
||
|
||
// Try to fit the object in the slab, growing if necessary.
|
||
match slab.grow_if_necessary(allocation.offset + data_slot_count, settings) {
|
||
SlabGrowthResult::NoGrowthNeeded => {}
|
||
SlabGrowthResult::NeededGrowth(slab_to_reallocate) => {
|
||
// If we already grew the slab this frame, don't replace the
|
||
// `SlabToReallocate` entry. We want to keep the entry
|
||
// corresponding to the size that the slab had at the start
|
||
// of the frame, so that we can copy only the used portion
|
||
// of the initial buffer to the new one.
|
||
if let Entry::Vacant(vacant_entry) = slabs_to_grow.entry(slab_id) {
|
||
vacant_entry.insert(slab_to_reallocate);
|
||
}
|
||
}
|
||
SlabGrowthResult::CantGrow => continue,
|
||
}
|
||
|
||
mesh_allocation = Some(MeshAllocation {
|
||
slab_id,
|
||
slab_allocation: SlabAllocation {
|
||
allocation,
|
||
slot_count: data_slot_count,
|
||
},
|
||
});
|
||
break;
|
||
}
|
||
|
||
// If we still have no allocation, make a new slab.
|
||
if mesh_allocation.is_none() {
|
||
let new_slab_id = self.next_slab_id;
|
||
self.next_slab_id.0 = NonMaxU32::new(self.next_slab_id.0.get() + 1).unwrap_or_default();
|
||
|
||
let new_slab = GeneralSlab::new(
|
||
new_slab_id,
|
||
&mut mesh_allocation,
|
||
settings,
|
||
layout,
|
||
data_slot_count,
|
||
);
|
||
|
||
self.slabs.insert(new_slab_id, Slab::General(new_slab));
|
||
candidate_slabs.push(new_slab_id);
|
||
slabs_to_grow.insert(new_slab_id, SlabToReallocate::default());
|
||
}
|
||
|
||
let mesh_allocation = mesh_allocation.expect("Should have been able to allocate");
|
||
|
||
// Mark the allocation as pending. Don't copy it in just yet; further
|
||
// meshes loaded this frame may result in its final allocation location
|
||
// changing.
|
||
if let Some(Slab::General(general_slab)) = self.slabs.get_mut(&mesh_allocation.slab_id) {
|
||
general_slab
|
||
.pending_allocations
|
||
.insert(*mesh_id, mesh_allocation.slab_allocation);
|
||
};
|
||
|
||
self.record_allocation(mesh_id, mesh_allocation.slab_id, layout.class);
|
||
}
|
||
|
||
/// Allocates an object into its own dedicated slab.
|
||
fn allocate_large(&mut self, mesh_id: &AssetId<Mesh>, layout: ElementLayout) {
|
||
let new_slab_id = self.next_slab_id;
|
||
self.next_slab_id.0 = NonMaxU32::new(self.next_slab_id.0.get() + 1).unwrap_or_default();
|
||
|
||
self.record_allocation(mesh_id, new_slab_id, layout.class);
|
||
|
||
self.slabs.insert(
|
||
new_slab_id,
|
||
Slab::LargeObject(LargeObjectSlab {
|
||
buffer: None,
|
||
element_layout: layout,
|
||
}),
|
||
);
|
||
}
|
||
|
||
/// Reallocates a slab that needs to be resized, or allocates a new slab.
|
||
///
|
||
/// This performs the actual growth operation that
|
||
/// [`GeneralSlab::grow_if_necessary`] scheduled. We do the growth in two
|
||
/// phases so that, if a slab grows multiple times in the same frame, only
|
||
/// one new buffer is reallocated, rather than reallocating the buffer
|
||
/// multiple times.
|
||
fn reallocate_slab(
|
||
&mut self,
|
||
render_device: &RenderDevice,
|
||
render_queue: &RenderQueue,
|
||
slab_id: SlabId,
|
||
slab_to_grow: SlabToReallocate,
|
||
) {
|
||
let Some(Slab::General(slab)) = self.slabs.get_mut(&slab_id) else {
|
||
error!("Couldn't find slab {} to grow", slab_id);
|
||
return;
|
||
};
|
||
|
||
let old_buffer = slab.buffer.take();
|
||
|
||
let mut buffer_usages = BufferUsages::COPY_SRC | BufferUsages::COPY_DST;
|
||
match slab.element_layout.class {
|
||
ElementClass::Vertex => buffer_usages |= BufferUsages::VERTEX,
|
||
ElementClass::Index => buffer_usages |= BufferUsages::INDEX,
|
||
};
|
||
|
||
// Create the buffer.
|
||
let new_buffer = render_device.create_buffer(&BufferDescriptor {
|
||
label: Some(&format!(
|
||
"general mesh slab {} ({}buffer)",
|
||
slab_id,
|
||
buffer_usages_to_str(buffer_usages)
|
||
)),
|
||
size: slab.current_slot_capacity as u64 * slab.element_layout.slot_size(),
|
||
usage: buffer_usages,
|
||
mapped_at_creation: false,
|
||
});
|
||
|
||
slab.buffer = Some(new_buffer.clone());
|
||
|
||
let Some(old_buffer) = old_buffer else { return };
|
||
|
||
// In order to do buffer copies, we need a command encoder.
|
||
let mut encoder = render_device.create_command_encoder(&CommandEncoderDescriptor {
|
||
label: Some("slab resize encoder"),
|
||
});
|
||
|
||
// Copy the data from the old buffer into the new one.
|
||
encoder.copy_buffer_to_buffer(
|
||
&old_buffer,
|
||
0,
|
||
&new_buffer,
|
||
0,
|
||
slab_to_grow.old_slot_capacity as u64 * slab.element_layout.slot_size(),
|
||
);
|
||
|
||
let command_buffer = encoder.finish();
|
||
render_queue.submit([command_buffer]);
|
||
}
|
||
|
||
/// Records the location of the given newly-allocated mesh data in the
|
||
/// [`Self::mesh_id_to_vertex_slab`] or [`Self::mesh_id_to_index_slab`]
|
||
/// tables as appropriate.
|
||
fn record_allocation(
|
||
&mut self,
|
||
mesh_id: &AssetId<Mesh>,
|
||
slab_id: SlabId,
|
||
element_class: ElementClass,
|
||
) {
|
||
match element_class {
|
||
ElementClass::Vertex => {
|
||
self.mesh_id_to_vertex_slab.insert(*mesh_id, slab_id);
|
||
}
|
||
ElementClass::Index => {
|
||
self.mesh_id_to_index_slab.insert(*mesh_id, slab_id);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
impl GeneralSlab {
|
||
/// Creates a new growable slab big enough to hold a single element of
|
||
/// `data_slot_count` size with the given `layout`.
|
||
fn new(
|
||
new_slab_id: SlabId,
|
||
mesh_allocation: &mut Option<MeshAllocation>,
|
||
settings: &MeshAllocatorSettings,
|
||
layout: ElementLayout,
|
||
data_slot_count: u32,
|
||
) -> GeneralSlab {
|
||
let initial_slab_slot_capacity = (settings.min_slab_size.div_ceil(layout.slot_size())
|
||
as u32)
|
||
.max(offset_allocator::ext::min_allocator_size(data_slot_count));
|
||
let max_slab_slot_capacity = (settings.max_slab_size.div_ceil(layout.slot_size()) as u32)
|
||
.max(offset_allocator::ext::min_allocator_size(data_slot_count));
|
||
|
||
let mut new_slab = GeneralSlab {
|
||
allocator: Allocator::new(max_slab_slot_capacity),
|
||
buffer: None,
|
||
resident_allocations: HashMap::default(),
|
||
pending_allocations: HashMap::default(),
|
||
element_layout: layout,
|
||
current_slot_capacity: initial_slab_slot_capacity,
|
||
};
|
||
|
||
// This should never fail.
|
||
if let Some(allocation) = new_slab.allocator.allocate(data_slot_count) {
|
||
*mesh_allocation = Some(MeshAllocation {
|
||
slab_id: new_slab_id,
|
||
slab_allocation: SlabAllocation {
|
||
slot_count: data_slot_count,
|
||
allocation,
|
||
},
|
||
});
|
||
}
|
||
|
||
new_slab
|
||
}
|
||
|
||
/// Checks to see if the size of this slab is at least `new_size_in_slots`
|
||
/// and grows the slab if it isn't.
|
||
///
|
||
/// The returned [`SlabGrowthResult`] describes whether the slab needed to
|
||
/// grow and whether, if so, it was successful in doing so.
|
||
fn grow_if_necessary(
|
||
&mut self,
|
||
new_size_in_slots: u32,
|
||
settings: &MeshAllocatorSettings,
|
||
) -> SlabGrowthResult {
|
||
// Is the slab big enough already?
|
||
let initial_slot_capacity = self.current_slot_capacity;
|
||
if self.current_slot_capacity >= new_size_in_slots {
|
||
return SlabGrowthResult::NoGrowthNeeded;
|
||
}
|
||
|
||
// Try to grow in increments of `MeshAllocatorSettings::growth_factor`
|
||
// until we're big enough.
|
||
while self.current_slot_capacity < new_size_in_slots {
|
||
let new_slab_slot_capacity =
|
||
((self.current_slot_capacity as f64 * settings.growth_factor).ceil() as u32)
|
||
.min((settings.max_slab_size / self.element_layout.slot_size()) as u32);
|
||
if new_slab_slot_capacity == self.current_slot_capacity {
|
||
// The slab is full.
|
||
return SlabGrowthResult::CantGrow;
|
||
}
|
||
|
||
self.current_slot_capacity = new_slab_slot_capacity;
|
||
}
|
||
|
||
// Tell our caller what we did.
|
||
SlabGrowthResult::NeededGrowth(SlabToReallocate {
|
||
old_slot_capacity: initial_slot_capacity,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl ElementLayout {
|
||
/// Creates an [`ElementLayout`] for mesh data of the given class (vertex or
|
||
/// index) with the given byte size.
|
||
fn new(class: ElementClass, size: u64) -> ElementLayout {
|
||
const {
|
||
assert!(4 == COPY_BUFFER_ALIGNMENT);
|
||
}
|
||
// this is equivalent to `4 / gcd(4,size)` but lets us not implement gcd.
|
||
// ping @atlv if above assert ever fails (likely never)
|
||
let elements_per_slot = [1, 4, 2, 4][size as usize & 3];
|
||
ElementLayout {
|
||
class,
|
||
size,
|
||
// Make sure that slot boundaries begin and end on
|
||
// `COPY_BUFFER_ALIGNMENT`-byte (4-byte) boundaries.
|
||
elements_per_slot,
|
||
}
|
||
}
|
||
|
||
fn slot_size(&self) -> u64 {
|
||
self.size * self.elements_per_slot as u64
|
||
}
|
||
|
||
/// Creates the appropriate [`ElementLayout`] for the given mesh's vertex
|
||
/// data.
|
||
fn vertex(
|
||
mesh_vertex_buffer_layouts: &mut MeshVertexBufferLayouts,
|
||
mesh: &Mesh,
|
||
) -> ElementLayout {
|
||
let mesh_vertex_buffer_layout =
|
||
mesh.get_mesh_vertex_buffer_layout(mesh_vertex_buffer_layouts);
|
||
ElementLayout::new(
|
||
ElementClass::Vertex,
|
||
mesh_vertex_buffer_layout.0.layout().array_stride,
|
||
)
|
||
}
|
||
|
||
/// Creates the appropriate [`ElementLayout`] for the given mesh's index
|
||
/// data.
|
||
fn index(mesh: &Mesh) -> Option<ElementLayout> {
|
||
let size = match mesh.indices()? {
|
||
Indices::U16(_) => 2,
|
||
Indices::U32(_) => 4,
|
||
};
|
||
Some(ElementLayout::new(ElementClass::Index, size))
|
||
}
|
||
}
|
||
|
||
impl GeneralSlab {
|
||
/// Returns true if this slab is empty.
|
||
fn is_empty(&self) -> bool {
|
||
self.resident_allocations.is_empty() && self.pending_allocations.is_empty()
|
||
}
|
||
}
|
||
|
||
/// Returns a string describing the given buffer usages.
|
||
fn buffer_usages_to_str(buffer_usages: BufferUsages) -> &'static str {
|
||
if buffer_usages.contains(BufferUsages::VERTEX) {
|
||
"vertex "
|
||
} else if buffer_usages.contains(BufferUsages::INDEX) {
|
||
"index "
|
||
} else {
|
||
""
|
||
}
|
||
}
|