wip: add: Abstraction for bindless resources

src/resource/bindless.rs

@@ -0,0 +1,243 @@
+use std::sync::{Arc, Mutex};
+
+use anyhow::Result;
+
+use ash::vk;
+
+use crate::{pipeline::set_layout::DescriptorSetLayoutCreateInfo, util::cache::Resource};
+
+/// The maximum number of resources in a pool
+pub static MAX_BINDLESS_COUNT: u32 = 4096;
+
+/// A resource that can be used bindlessly.
+pub trait BindlessResource {
+ /// Get the [vk::DescriptorType] for this resource.
+ fn descriptor_type() -> vk::DescriptorType;
+ /// Get the [vk::DescriptorSetLayoutBinding] of this resource.
+ 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(),
+ }
+ }
+ /// Get the [vk::DescriptorImageInfo] of this resource.
+ fn descriptor_info(&self) -> vk::DescriptorImageInfo;
+
+ /// Add this resource to `pool` and get an handle to it.
+ 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!()
+ }
+}
+
+/// Resource for a combined image sampler.
+#[derive(Debug)]
+pub struct CombinedImageSampler {
+ sampler: crate::sampler::Sampler,
+ image_view: crate::image::ImageView,
+ image_layout: Option<vk::ImageLayout>,
+}
+
+impl CombinedImageSampler {
+ /// Create a new combined image sampler resource from a sampler and an ImageView.
+ /// The default image layout [vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL] will be used.
+ pub fn new(sampler: crate::sampler::Sampler, image_view: crate::image::ImageView) -> Self {
+ Self {
+ sampler,
+ image_view,
+ image_layout: None,
+ }
+ }
+
+ /// Specify an image layout for this combined image sampler
+ 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),
+ }
+ }
+}
+
+/// An handle to a resource contained in a [BindlessPool]
+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> {
+ /// Get the index of this resource in the pool's descripor.
+ /// This can be sent to a shader through any mean to access it.
+ 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)
+ })
+ }
+ }
+}
+
+/// A bindless pool can hold a number of resources.
+/// It will keep a descriptor set up to date so that any resource contained in it.
+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)
+ }
+
+ /// Allocate a single item from the pool.
+ 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() }
+ }
+ })
+ }
+
+ /// Allocate a number of items from the pool.
+ 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> {
+ /// Create a new bindless pool
+ 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)
+ }
+
+ /// Create a new bindless pool that will allocate a descriptor set from `pool`.
+ /// The provided pool must be large enough for a [MAX_BINDLESS_COUNT] elements descriptor array
+ pub 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)) })
+ }
+}

src/resource/mod.rs

@@ -6,3 +6,4 @@ pub mod pool;
 pub mod query_pool;
 pub mod raytracing;
 pub mod sampler;
+pub mod bindless;