-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
wip: add: Abstraction for bindless resources
- Loading branch information
Showing
2 changed files
with
223 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,222 @@ | ||
use std::sync::{Arc, Mutex}; | ||
|
||
use anyhow::Result; | ||
|
||
use ash::vk; | ||
|
||
use crate::{pipeline::set_layout::DescriptorSetLayoutCreateInfo, util::cache::Resource}; | ||
|
||
static MAX_BINDLESS_COUNT: u32 = 4096; | ||
|
||
pub trait BindlessResource { | ||
fn descriptor_type() -> vk::DescriptorType; | ||
fn resource_binding(binding: u32) -> vk::DescriptorSetLayoutBinding { | ||
vk::DescriptorSetLayoutBinding { | ||
binding, | ||
descriptor_type: Self::descriptor_type(), | ||
descriptor_count: MAX_BINDLESS_COUNT, | ||
stage_flags: vk::ShaderStageFlags::ALL, | ||
p_immutable_samplers: std::ptr::null(), | ||
} | ||
} | ||
fn descriptor_info(&self) -> vk::DescriptorImageInfo; | ||
|
||
fn into_bindless(self, pool: &BindlessPool<Self>) -> BindlessHandle<Self> | ||
where | ||
Self: Sized | ||
{ | ||
pool.alloc(self) | ||
} | ||
} | ||
|
||
impl BindlessResource for crate::image::Image { | ||
fn descriptor_type() -> vk::DescriptorType { | ||
vk::DescriptorType::STORAGE_IMAGE | ||
} | ||
|
||
fn descriptor_info(&self) -> vk::DescriptorImageInfo { | ||
todo!() | ||
} | ||
} | ||
|
||
#[derive(Debug)] | ||
pub struct CombinedImageSampler { | ||
sampler: crate::sampler::Sampler, | ||
image_view: crate::image::ImageView, | ||
image_layout: Option<vk::ImageLayout>, | ||
} | ||
|
||
impl CombinedImageSampler { | ||
pub fn new(sampler: crate::sampler::Sampler, image_view: crate::image::ImageView) -> Self { | ||
Self { | ||
sampler, | ||
image_view, | ||
image_layout: None, | ||
} | ||
} | ||
|
||
pub fn with_layout(self, image_layout: Option<vk::ImageLayout>) -> Self { | ||
Self { image_layout, .. self } | ||
} | ||
} | ||
|
||
impl BindlessResource for CombinedImageSampler { | ||
fn descriptor_type() -> vk::DescriptorType { | ||
vk::DescriptorType::COMBINED_IMAGE_SAMPLER | ||
} | ||
|
||
fn descriptor_info(&self) -> vk::DescriptorImageInfo { | ||
vk::DescriptorImageInfo { | ||
sampler: unsafe { self.sampler.handle() }, | ||
image_view: unsafe { self.image_view.handle() }, | ||
image_layout: self.image_layout.unwrap_or(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL), | ||
} | ||
} | ||
} | ||
|
||
pub struct BindlessHandle<R: BindlessResource> { | ||
key: u32, | ||
pool: BindlessPool<R>, | ||
} | ||
|
||
impl<R: BindlessResource> Drop for BindlessHandle<R> { | ||
fn drop(&mut self) { | ||
self.pool.with(|p| p.take(self.key)); | ||
} | ||
} | ||
|
||
impl<R: BindlessResource> BindlessHandle<R> { | ||
pub fn index(&self) -> u32 { | ||
self.key | ||
} | ||
} | ||
|
||
struct BindlessPoolInner<R> { | ||
items: Vec<Option<R>>, | ||
free: Vec<u32>, | ||
descriptor_set: crate::DescriptorSet, | ||
} | ||
|
||
impl<R: BindlessResource> BindlessPoolInner<R> { | ||
fn update_descriptor_set<'a>(&'a mut self, r: impl Iterator<Item = (u32, &'a R)>) { | ||
let vk_writes = r | ||
.map(|(i, r)| { | ||
vk::WriteDescriptorSet { | ||
s_type: vk::StructureType::WRITE_DESCRIPTOR_SET, | ||
p_next: std::ptr::null(), | ||
dst_set: self.descriptor_set.handle, | ||
dst_binding: 0, | ||
dst_array_element: i, | ||
descriptor_count: 1, | ||
descriptor_type: R::descriptor_type(), | ||
p_image_info: &r.descriptor_info() as *const _, | ||
p_buffer_info: std::ptr::null(), | ||
p_texel_buffer_view: std::ptr::null(), | ||
} | ||
}) | ||
.collect::<Vec<_>>(); | ||
|
||
unsafe { | ||
self.descriptor_set.device.update_descriptor_sets(vk_writes.as_slice(), &[]); | ||
} | ||
} | ||
|
||
|
||
fn take(&mut self, key: u32) -> Option<R> { | ||
if key <= self.items.len() as _ { | ||
None | ||
} else { | ||
self.items[key as usize].take().and_then(|ob| { | ||
self.free.push(key); | ||
Some(ob) | ||
}) | ||
} | ||
} | ||
} | ||
|
||
pub struct BindlessPool<R> { | ||
inner: Arc<Mutex<BindlessPoolInner<R>>>, | ||
} | ||
|
||
impl<P: BindlessResource> BindlessPool<P> { | ||
fn with<F: FnOnce(&mut BindlessPoolInner<P>) -> R, R>(&self, f: F) -> R { | ||
let mut inner = self.inner.lock().unwrap(); | ||
f(&mut inner) | ||
} | ||
|
||
pub fn alloc(&self, item: P) -> BindlessHandle<P> { | ||
self.with(|p| { | ||
let key = p.free | ||
.pop() | ||
.unwrap_or_else(|| { | ||
p.items.push(None); | ||
p.items.len() as u32 - 1 | ||
}); | ||
p.update_descriptor_set(std::iter::once((key, &item))); | ||
p.items[key as usize] = Some(item); | ||
BindlessHandle { | ||
key, | ||
pool: Self { inner: self.inner.clone() } | ||
} | ||
}) | ||
} | ||
|
||
pub fn alloc_items(&self, items: &[P]) -> impl Iterator<Item = BindlessHandle<P>> { | ||
self.with(|p| { | ||
let mut keys = p.free.iter().cloned().rev().take(items.len()).collect::<Vec<_>>(); | ||
if keys.len() < items.len() { | ||
let old_len = p.items.len() as u32; | ||
p.items.resize_with(items.len(), || None); | ||
keys.extend(old_len..p.items.len() as u32); | ||
} | ||
p.update_descriptor_set(keys.iter().cloned().zip(items)); | ||
let pool_inner = self.inner.clone(); | ||
keys | ||
.into_iter() | ||
.map(move |key| { | ||
BindlessHandle { | ||
key, | ||
pool: Self { inner: pool_inner.clone() } | ||
} | ||
}) | ||
}) | ||
} | ||
} | ||
|
||
impl<P: BindlessResource> BindlessPool<P> { | ||
pub fn new(device: crate::Device) -> Result<Self> { | ||
let pool_size = vk::DescriptorPoolSize { | ||
ty: P::descriptor_type(), | ||
descriptor_count: MAX_BINDLESS_COUNT, | ||
}; | ||
let pool_create_info = vk::DescriptorPoolCreateInfo { | ||
s_type: vk::StructureType::DESCRIPTOR_POOL_CREATE_INFO, | ||
p_next: std::ptr::null(), | ||
flags: vk::DescriptorPoolCreateFlags::UPDATE_AFTER_BIND, | ||
max_sets: 1, | ||
pool_size_count: 1, | ||
p_pool_sizes: &pool_size as *const _, | ||
}; | ||
let pool = unsafe { device.create_descriptor_pool(&pool_create_info, None)? }; | ||
Self::new_with_pool(device, pool) | ||
} | ||
fn new_with_pool(device: crate::Device, pool: vk::DescriptorPool) -> Result<Self> { | ||
let dsl_info = DescriptorSetLayoutCreateInfo { | ||
bindings: vec![ | ||
P::resource_binding(0) | ||
], | ||
persistent: true, | ||
flags: vec![vk::DescriptorBindingFlags::UPDATE_AFTER_BIND, vk::DescriptorBindingFlags::PARTIALLY_BOUND], | ||
}; | ||
let dsl = crate::pipeline::set_layout::DescriptorSetLayout::create(device.clone(), &dsl_info, ())?; | ||
let descriptor_set = unsafe { crate::DescriptorSet::new_uninitialized(device, dsl.handle(), pool)? }; | ||
|
||
let inner = BindlessPoolInner { | ||
items: vec![], | ||
free: vec![], | ||
descriptor_set, | ||
}; | ||
|
||
Ok(Self { inner: Arc::new(Mutex::new(inner)) }) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -6,3 +6,4 @@ pub mod pool; | |
pub mod query_pool; | ||
pub mod raytracing; | ||
pub mod sampler; | ||
pub mod bindless; |