Skip to content

Commit

Permalink
Improve pool memory management (#70)
Browse files Browse the repository at this point in the history
* Clean up pool code
* Add basic pool memory management functions
* Add FifoPool and configurable pool bucket sizes
  • Loading branch information
attackgoat authored Jan 30, 2024
1 parent 3de19d6 commit 43293d4
Show file tree
Hide file tree
Showing 7 changed files with 808 additions and 305 deletions.
2 changes: 1 addition & 1 deletion examples/msaa.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ fn main() -> anyhow::Result<()> {
let mesh_msaa_pipeline = create_mesh_pipeline(&event_loop.device, sample_count)?;
let mesh_noaa_pipeline = create_mesh_pipeline(&event_loop.device, SampleCount::X1)?;
let cube_mesh = load_cube_mesh(&event_loop.device)?;
let mut pool = HashPool::new(&event_loop.device);
let mut pool = FifoPool::new(&event_loop.device);

let mut angle = 0f32;

Expand Down
2 changes: 1 addition & 1 deletion examples/vsm_omni.rs
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ fn main() -> anyhow::Result<()> {
}?;

// A pool will be used for per-frame resources
let mut pool = LazyPool::new(&event_loop.device);
let mut pool = FifoPool::new(&event_loop.device);

let mut elapsed = 0.0;
event_loop.run(|frame| {
Expand Down
5 changes: 4 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,10 @@ pub mod prelude {
input::{
update_input, update_keyboard, update_mouse, KeyBuf, KeyMap, MouseBuf, MouseButton,
},
pool::{hash::HashPool, lazy::LazyPool, Lease, Pool},
pool::{
fifo::FifoPool, hash::HashPool, lazy::LazyPool, Lease, Pool, PoolInfo,
PoolInfoBuilder,
},
},
ash::vk,
log::{debug, error, info, logger, trace, warn}, // Everyone wants a log
Expand Down
289 changes: 289 additions & 0 deletions src/pool/fifo.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,289 @@
//! Pool which leases from a single bucket per resource type.
use {
super::{can_lease_command_buffer, Cache, Lease, Pool, PoolInfo},
crate::driver::{
accel_struct::{AccelerationStructure, AccelerationStructureInfo},
buffer::{Buffer, BufferInfo},
device::Device,
image::{Image, ImageInfo},
CommandBuffer, CommandBufferInfo, DescriptorPool, DescriptorPoolInfo, DriverError,
RenderPass, RenderPassInfo,
},
log::debug,
std::{collections::HashMap, sync::Arc},
};

/// A memory-efficient resource allocator.
///
/// The information for each lease request is compared against the stored resources for
/// compatibility. If no acceptable resources are stored for the information provided a new resource
/// is created and returned.
///
/// # Details
///
/// * Acceleration structures may be larger than requested
/// * Buffers may be larger than requested or have additional usage flags
/// * Images may have additional usage flags
///
/// # Bucket Strategy
///
/// All resources are stored in a single bucket per resource type, regardless of their individual
/// attributes.
///
/// In practice this means that for a [`PoolInfo::image_capacity`] of `4`, a maximum of `4` images
/// will be stored. Requests to lease an image or other resource will first look for a compatible
/// resource in the bucket and create a new resource as needed.
///
/// # Memory Management
///
/// The single-bucket strategy means that there will always be a reasonable and predictable number
/// of stored resources, however you may call [`FifoPool::clear`] or the other memory management
/// functions at any time to discard stored resources.
#[derive(Debug)]
pub struct FifoPool {
accel_struct_cache: Cache<AccelerationStructure>,
buffer_cache: Cache<Buffer>,
command_buffer_cache: HashMap<u32, Cache<CommandBuffer>>,
descriptor_pool_cache: Cache<DescriptorPool>,
device: Arc<Device>,
image_cache: Cache<Image>,
info: PoolInfo,
render_pass_cache: HashMap<RenderPassInfo, Cache<RenderPass>>,
}

impl FifoPool {
/// Constructs a new `FifoPool`.
pub fn new(device: &Arc<Device>) -> Self {
Self::with_capacity(device, PoolInfo::default())
}

/// Constructs a new `FifoPool` with the given capacity information.
pub fn with_capacity(device: &Arc<Device>, info: impl Into<PoolInfo>) -> Self {
let info: PoolInfo = info.into();
let device = Arc::clone(device);

Self {
accel_struct_cache: PoolInfo::explicit_cache(info.accel_struct_capacity),
buffer_cache: PoolInfo::explicit_cache(info.buffer_capacity),
command_buffer_cache: Default::default(),
descriptor_pool_cache: PoolInfo::default_cache(),
device,
image_cache: PoolInfo::explicit_cache(info.image_capacity),
info,
render_pass_cache: Default::default(),
}
}

/// Clears the pool, removing all resources.
pub fn clear(&mut self) {
self.clear_accel_structs();
self.clear_buffers();
self.clear_images();
}

/// Clears the pool of acceleration structure resources.
pub fn clear_accel_structs(&mut self) {
self.accel_struct_cache = PoolInfo::explicit_cache(self.info.accel_struct_capacity);
}

/// Clears the pool of buffer resources.
pub fn clear_buffers(&mut self) {
self.buffer_cache = PoolInfo::explicit_cache(self.info.buffer_capacity);
}

/// Clears the pool of image resources.
pub fn clear_images(&mut self) {
self.image_cache = PoolInfo::explicit_cache(self.info.image_capacity);
}
}

impl Pool<AccelerationStructureInfo, AccelerationStructure> for FifoPool {
#[profiling::function]
fn lease(
&mut self,
info: AccelerationStructureInfo,
) -> Result<Lease<AccelerationStructure>, DriverError> {
let cache_ref = Arc::downgrade(&self.accel_struct_cache);
let mut cache = self.accel_struct_cache.lock();

{
profiling::scope!("Check cache");

// Look for a compatible acceleration structure (big enough and same type)
for idx in 0..cache.len() {
let item = &cache[idx];
if item.info.size >= info.size && item.info.ty == info.ty {
let item = cache.remove(idx).unwrap();

return Ok(Lease::new(cache_ref, item));
}
}
}

debug!("Creating new {}", stringify!(AccelerationStructure));

let item = AccelerationStructure::create(&self.device, info)?;

Ok(Lease::new(cache_ref, item))
}
}

impl Pool<BufferInfo, Buffer> for FifoPool {
#[profiling::function]
fn lease(&mut self, info: BufferInfo) -> Result<Lease<Buffer>, DriverError> {
let cache_ref = Arc::downgrade(&self.buffer_cache);
let mut cache = self.buffer_cache.lock();

{
profiling::scope!("Check cache");

// Look for a compatible buffer (compatible alignment, same mapping mode, big enough and
// superset of usage flags)
for idx in 0..cache.len() {
let item = &cache[idx];
if item.info.alignment >= info.alignment
&& item.info.can_map == info.can_map
&& item.info.size >= info.size
&& item.info.usage.contains(info.usage)
{
let item = cache.remove(idx).unwrap();

return Ok(Lease::new(cache_ref, item));
}
}
}

debug!("Creating new {}", stringify!(Buffer));

let item = Buffer::create(&self.device, info)?;

Ok(Lease::new(cache_ref, item))
}
}

impl Pool<CommandBufferInfo, CommandBuffer> for FifoPool {
#[profiling::function]
fn lease(&mut self, info: CommandBufferInfo) -> Result<Lease<CommandBuffer>, DriverError> {
let cache = self
.command_buffer_cache
.entry(info.queue_family_index)
.or_insert_with(PoolInfo::default_cache);
let mut item = cache
.lock()
.pop_front()
.filter(can_lease_command_buffer)
.map(Ok)
.unwrap_or_else(|| {
debug!("Creating new {}", stringify!(CommandBuffer));

CommandBuffer::create(&self.device, info)
})?;

// Drop anything we were holding from the last submission
CommandBuffer::drop_fenced(&mut item);

Ok(Lease::new(Arc::downgrade(cache), item))
}
}

impl Pool<DescriptorPoolInfo, DescriptorPool> for FifoPool {
#[profiling::function]
fn lease(&mut self, info: DescriptorPoolInfo) -> Result<Lease<DescriptorPool>, DriverError> {
let cache_ref = Arc::downgrade(&self.descriptor_pool_cache);
let mut cache = self.descriptor_pool_cache.lock();

{
profiling::scope!("Check cache");

// Look for a compatible descriptor pool (has enough sets and descriptors)
for idx in 0..cache.len() {
let item = &cache[idx];
if item.info.max_sets >= info.max_sets
&& item.info.acceleration_structure_count >= info.acceleration_structure_count
&& item.info.combined_image_sampler_count >= info.combined_image_sampler_count
&& item.info.input_attachment_count >= info.input_attachment_count
&& item.info.sampled_image_count >= info.sampled_image_count
&& item.info.storage_buffer_count >= info.storage_buffer_count
&& item.info.storage_buffer_dynamic_count >= info.storage_buffer_dynamic_count
&& item.info.storage_image_count >= info.storage_image_count
&& item.info.storage_texel_buffer_count >= info.storage_texel_buffer_count
&& item.info.uniform_buffer_count >= info.uniform_buffer_count
&& item.info.uniform_buffer_dynamic_count >= info.uniform_buffer_dynamic_count
&& item.info.uniform_texel_buffer_count >= info.uniform_texel_buffer_count
{
let item = cache.remove(idx).unwrap();

return Ok(Lease::new(cache_ref, item));
}
}
}

debug!("Creating new {}", stringify!(DescriptorPool));

let item = DescriptorPool::create(&self.device, info)?;

Ok(Lease::new(cache_ref, item))
}
}

impl Pool<ImageInfo, Image> for FifoPool {
#[profiling::function]
fn lease(&mut self, info: ImageInfo) -> Result<Lease<Image>, DriverError> {
let cache_ref = Arc::downgrade(&self.image_cache);
let mut cache = self.image_cache.lock();

{
profiling::scope!("Check cache");

// Look for a compatible image (same properties, superset of creation flags and usage
// flags)
for idx in 0..cache.len() {
let item = &cache[idx];
if item.info.array_elements == info.array_elements
&& item.info.depth == info.depth
&& item.info.fmt == info.fmt
&& item.info.height == info.height
&& item.info.linear_tiling == info.linear_tiling
&& item.info.mip_level_count == info.mip_level_count
&& item.info.sample_count == info.sample_count
&& item.info.ty == info.ty
&& item.info.width == info.width
&& item.info.flags.contains(info.flags)
&& item.info.usage.contains(info.usage)
{
let item = cache.remove(idx).unwrap();

return Ok(Lease::new(cache_ref, item));
}
}
}

debug!("Creating new {}", stringify!(Image));

let item = Image::create(&self.device, info)?;

Ok(Lease::new(cache_ref, item))
}
}

impl Pool<RenderPassInfo, RenderPass> for FifoPool {
#[profiling::function]
fn lease(&mut self, info: RenderPassInfo) -> Result<Lease<RenderPass>, DriverError> {
let cache = if let Some(cache) = self.render_pass_cache.get(&info) {
cache
} else {
// We tried to get the cache first in order to avoid this clone
self.render_pass_cache
.entry(info.clone())
.or_insert_with(PoolInfo::default_cache)
};
let item = cache.lock().pop_front().map(Ok).unwrap_or_else(|| {
debug!("Creating new {}", stringify!(RenderPass));

RenderPass::create(&self.device, info)
})?;

Ok(Lease::new(Arc::downgrade(cache), item))
}
}
Loading

0 comments on commit 43293d4

Please sign in to comment.