use {
super::{lease_command_buffer, Cache, Lease, Pool, PoolInfo},
crate::driver::{
accel_struct::{AccelerationStructure, AccelerationStructureInfo},
buffer::{Buffer, BufferInfo},
device::Device,
image::{Image, ImageInfo, ImageType, SampleCount},
CommandBuffer, CommandBufferInfo, DescriptorPool, DescriptorPoolInfo, DriverError,
RenderPass, RenderPassInfo,
},
ash::vk,
log::debug,
std::{collections::HashMap, sync::Arc},
};
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
struct ImageKey {
array_elements: u32,
depth: u32,
fmt: vk::Format,
height: u32,
mip_level_count: u32,
sample_count: SampleCount,
tiling: vk::ImageTiling,
ty: ImageType,
width: u32,
}
impl From<ImageInfo> for ImageKey {
fn from(info: ImageInfo) -> Self {
Self {
array_elements: info.array_elements,
depth: info.depth,
fmt: info.fmt,
height: info.height,
mip_level_count: info.mip_level_count,
sample_count: info.sample_count,
tiling: info.tiling,
ty: info.ty,
width: info.width,
}
}
}
#[derive(Debug)]
pub struct LazyPool {
accel_struct_cache: HashMap<vk::AccelerationStructureTypeKHR, Cache<AccelerationStructure>>,
buffer_cache: HashMap<(bool, vk::DeviceSize), Cache<Buffer>>,
command_buffer_cache: HashMap<u32, Cache<CommandBuffer>>,
descriptor_pool_cache: Cache<DescriptorPool>,
device: Arc<Device>,
image_cache: HashMap<ImageKey, Cache<Image>>,
info: PoolInfo,
render_pass_cache: HashMap<RenderPassInfo, Cache<RenderPass>>,
}
impl LazyPool {
pub fn new(device: &Arc<Device>) -> Self {
Self::with_capacity(device, PoolInfo::default())
}
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: Default::default(),
buffer_cache: Default::default(),
command_buffer_cache: Default::default(),
descriptor_pool_cache: PoolInfo::default_cache(),
device,
image_cache: Default::default(),
info,
render_pass_cache: Default::default(),
}
}
pub fn clear(&mut self) {
self.clear_accel_structs();
self.clear_buffers();
self.clear_images();
}
pub fn clear_accel_structs(&mut self) {
self.accel_struct_cache.clear();
}
pub fn clear_accel_structs_by_ty(&mut self, ty: vk::AccelerationStructureTypeKHR) {
self.accel_struct_cache.remove(&ty);
}
pub fn clear_buffers(&mut self) {
self.buffer_cache.clear();
}
pub fn clear_images(&mut self) {
self.image_cache.clear();
}
pub fn clear_images_by_info(&mut self, info: impl Into<ImageInfo>) {
self.image_cache.remove(&info.into().into());
}
pub fn retain_accel_structs<F>(&mut self, mut f: F)
where
F: FnMut(vk::AccelerationStructureTypeKHR) -> bool,
{
self.accel_struct_cache.retain(|&ty, _| f(ty))
}
}
impl Pool<AccelerationStructureInfo, AccelerationStructure> for LazyPool {
#[profiling::function]
fn lease(
&mut self,
info: AccelerationStructureInfo,
) -> Result<Lease<AccelerationStructure>, DriverError> {
let cache = self
.accel_struct_cache
.entry(info.ty)
.or_insert_with(|| PoolInfo::explicit_cache(self.info.accel_struct_capacity));
let cache_ref = Arc::downgrade(cache);
{
profiling::scope!("check cache");
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = cache.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
for idx in 0..cache.len() {
let item = unsafe { cache.get_unchecked(idx) };
if item.info.size >= info.size {
let item = cache.swap_remove(idx);
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 LazyPool {
#[profiling::function]
fn lease(&mut self, info: BufferInfo) -> Result<Lease<Buffer>, DriverError> {
let cache = self
.buffer_cache
.entry((info.mappable, info.alignment))
.or_insert_with(|| PoolInfo::explicit_cache(self.info.buffer_capacity));
let cache_ref = Arc::downgrade(cache);
{
profiling::scope!("check cache");
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = cache.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
for idx in 0..cache.len() {
let item = unsafe { cache.get_unchecked(idx) };
if item.info.size >= info.size && item.info.usage.contains(info.usage) {
let item = cache.swap_remove(idx);
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 LazyPool {
#[profiling::function]
fn lease(&mut self, info: CommandBufferInfo) -> Result<Lease<CommandBuffer>, DriverError> {
let cache_ref = self
.command_buffer_cache
.entry(info.queue_family_index)
.or_insert_with(PoolInfo::default_cache);
let mut item = {
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = cache_ref.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
lease_command_buffer(&mut cache)
}
.map(Ok)
.unwrap_or_else(|| {
debug!("Creating new {}", stringify!(CommandBuffer));
CommandBuffer::create(&self.device, info)
})?;
CommandBuffer::drop_fenced(&mut item);
Ok(Lease::new(Arc::downgrade(cache_ref), item))
}
}
impl Pool<DescriptorPoolInfo, DescriptorPool> for LazyPool {
#[profiling::function]
fn lease(&mut self, info: DescriptorPoolInfo) -> Result<Lease<DescriptorPool>, DriverError> {
let cache_ref = Arc::downgrade(&self.descriptor_pool_cache);
{
profiling::scope!("check cache");
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = self.descriptor_pool_cache.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
for idx in 0..cache.len() {
let item = unsafe { cache.get_unchecked(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.swap_remove(idx);
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 LazyPool {
#[profiling::function]
fn lease(&mut self, info: ImageInfo) -> Result<Lease<Image>, DriverError> {
let cache = self
.image_cache
.entry(info.into())
.or_insert_with(|| PoolInfo::explicit_cache(self.info.image_capacity));
let cache_ref = Arc::downgrade(cache);
{
profiling::scope!("check cache");
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = cache.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
for idx in 0..cache.len() {
let item = unsafe { cache.get_unchecked(idx) };
if item.info.flags.contains(info.flags) && item.info.usage.contains(info.usage) {
let item = cache.swap_remove(idx);
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 LazyPool {
#[profiling::function]
fn lease(&mut self, info: RenderPassInfo) -> Result<Lease<RenderPass>, DriverError> {
let cache_ref = if let Some(cache) = self.render_pass_cache.get(&info) {
cache
} else {
self.render_pass_cache
.entry(info.clone())
.or_insert_with(PoolInfo::default_cache)
};
let item = {
#[cfg_attr(not(feature = "parking_lot"), allow(unused_mut))]
let mut cache = cache_ref.lock();
#[cfg(not(feature = "parking_lot"))]
let mut cache = cache.unwrap();
cache.pop()
}
.map(Ok)
.unwrap_or_else(|| {
debug!("Creating new {}", stringify!(RenderPass));
RenderPass::create(&self.device, info)
})?;
Ok(Lease::new(Arc::downgrade(cache_ref), item))
}
}