diff --git a/src/command_buffer/incomplete.rs b/src/command_buffer/incomplete.rs index 2ea410d..afe4e63 100644 --- a/src/command_buffer/incomplete.rs +++ b/src/command_buffer/incomplete.rs @@ -22,6 +22,8 @@ use crate::{ VirtualResource, }; +use super::DescriptorState; + impl<'q, D: ExecutionDomain, A: Allocator> IncompleteCmdBuffer<'q, A> for IncompleteCommandBuffer<'q, D, A> { @@ -123,6 +125,27 @@ impl IncompleteCommandBuffer<'_, D, A> { Ok(()) } + fn set_descriptor_set( + &mut self, + set: u32, + descriptor_set: Arc, + ) -> Result<()> { + if self.current_descriptor_sets.is_none() { + self.current_descriptor_sets = Some(HashMap::new()); + } + + match self.current_descriptor_sets.as_mut().unwrap().entry(set) { + Entry::Occupied(mut entry) => { + *entry.get_mut() = DescriptorState::DescriptorSet(descriptor_set.clone()); + } + Entry::Vacant(entry) => { + entry.insert(DescriptorState::DescriptorSet(descriptor_set.clone())); + } + }; + self.descriptor_state_needs_update = true; + Ok(()) + } + /// Modify the descriptor set state at a given set binding. /// # Errors /// * Fails if the supplied callback fails. @@ -137,12 +160,20 @@ impl IncompleteCommandBuffer<'_, D, A> { match self.current_descriptor_sets.as_mut().unwrap().entry(set) { Entry::Occupied(mut entry) => { - f(entry.get_mut())?; + match entry.get_mut() { + DescriptorState::Builder(ref mut builder) => f(builder)?, + entry => { + // replace descriptor set with builder + let mut builder = DescriptorSetBuilder::new(); + f(&mut builder)?; + *entry = DescriptorState::Builder(builder); + }, + } } Entry::Vacant(entry) => { let mut builder = DescriptorSetBuilder::new(); f(&mut builder)?; - entry.insert(builder); + entry.insert(DescriptorState::Builder(builder)); } }; self.descriptor_state_needs_update = true; @@ -160,13 +191,20 @@ impl IncompleteCommandBuffer<'_, D, A> { } let cache = self.descriptor_cache.clone(); - for (index, builder) in self.current_descriptor_sets.take().unwrap() { - let mut info = builder.build(); - info.layout = *self.current_set_layouts.get(index as usize).unwrap(); - cache.with_descriptor_set(info, |set| { - self.bind_descriptor_set(index, set)?; - Ok(()) - })?; + for (index, state) in self.current_descriptor_sets.take().unwrap() { + match state { + DescriptorState::Builder(builder) => { + let mut info = builder.build(); + info.layout = *self.current_set_layouts.get(index as usize).unwrap(); + cache.with_descriptor_set(info, |set| { + self.bind_descriptor_set(index, set)?; + Ok(()) + })?; + }, + DescriptorState::DescriptorSet(set) => { + self.bind_descriptor_set(index, &set)?; + } + } } // We updated all our descriptor sets, were good now. @@ -299,6 +337,22 @@ impl IncompleteCommandBuffer<'_, D, A> { Ok(self) } + /// Bind a [`BindlessPool`](crate::resouce::bindless::BindlessPool) + pub fn bind_bindless_pool( + mut self, + index: u32, + bindless_pool: &crate::bindless::BindlessPool + ) -> Result + where + R: crate::bindless::BindlessResource + { + bindless_pool.with(|p| { + self.set_descriptor_set(index, p.descriptor_set.clone()) + })?; + + Ok(self) + } + /// Binds a new descriptor with type [`vk::DescriptorType::UNIFORM_BUFFER`]. /// This binding is not actually flushed to the command buffer until the next draw or dispatch call. /// # Errors diff --git a/src/command_buffer/mod.rs b/src/command_buffer/mod.rs index 005d419..f127a48 100644 --- a/src/command_buffer/mod.rs +++ b/src/command_buffer/mod.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use std::marker::PhantomData; -use std::sync::MutexGuard; +use std::sync::{MutexGuard, Arc}; use anyhow::Result; use ash::vk; @@ -57,6 +57,13 @@ pub struct CommandBuffer { _domain: PhantomData, } +#[derive(Derivative)] +#[derivative(Debug)] +enum DescriptorState { + Builder(DescriptorSetBuilder<'static>), + DescriptorSet(Arc), +} + /// This struct represents an incomplete command buffer. /// This is a command buffer that has not been finished yet with [`IncompleteCommandBuffer::finish()`](crate::IncompleteCommandBuffer::finish). /// Calling this method will turn it into an immutable command buffer which can then be submitted @@ -96,7 +103,7 @@ pub struct IncompleteCommandBuffer<'q, D: ExecutionDomain, A: Allocator = Defaul current_bindpoint: vk::PipelineBindPoint, current_rendering_state: Option, current_render_area: vk::Rect2D, - current_descriptor_sets: Option>>, + current_descriptor_sets: Option>, descriptor_state_needs_update: bool, current_sbt_regions: Option<[vk::StridedDeviceAddressRegionKHR; 4]>, // TODO: Only update disturbed descriptor sets diff --git a/src/descriptor/descriptor_set.rs b/src/descriptor/descriptor_set.rs index ef7438f..d8d6ff0 100644 --- a/src/descriptor/descriptor_set.rs +++ b/src/descriptor/descriptor_set.rs @@ -54,6 +54,30 @@ pub struct DescriptorSet { pub(crate) handle: vk::DescriptorSet, } +impl DescriptorSet { + pub(crate) fn new_uninitialized(device: Device, layout: vk::DescriptorSetLayout, pool: vk::DescriptorPool) -> Result { + let info = vk::DescriptorSetAllocateInfo { + s_type: vk::StructureType::DESCRIPTOR_SET_ALLOCATE_INFO, + p_next: std::ptr::null(), + descriptor_pool: pool, + descriptor_set_count: 1, + p_set_layouts: &layout as *const _, + }; + let set = unsafe { device.allocate_descriptor_sets(&info) }? + .first() + .cloned() + .unwrap(); + #[cfg(feature = "log-objects")] + trace!("Created new VkDescriptorSet {set:p}"); + + Ok(DescriptorSet { + device, + pool: pool, + handle: set, + }) + } +} + fn binding_image_info(binding: &DescriptorBinding) -> Vec { binding .descriptors @@ -120,26 +144,13 @@ impl Resource for DescriptorSet { fn create(device: Device, key: &Self::Key, _: Self::ExtraParams<'_>) -> Result where Self: Sized, { - let info = vk::DescriptorSetAllocateInfo { - s_type: vk::StructureType::DESCRIPTOR_SET_ALLOCATE_INFO, - p_next: std::ptr::null(), - descriptor_pool: key.pool, - descriptor_set_count: 1, - p_set_layouts: &key.layout, - }; - let set = unsafe { device.allocate_descriptor_sets(&info) }? - .first() - .cloned() - .unwrap(); - #[cfg(feature = "log-objects")] - trace!("Created new VkDescriptorSet {set:p}"); - + let descriptor_set = DescriptorSet::new_uninitialized(device.clone(), key.layout, key.pool)?; let writes = key .bindings .iter() .map(|binding| { let mut write = WriteDescriptorSet { - set, + set: descriptor_set.handle, binding: binding.binding, array_element: 0, count: binding.descriptors.len() as u32, @@ -226,11 +237,7 @@ impl Resource for DescriptorSet { device.update_descriptor_sets(vk_writes.as_slice(), &[]); } - Ok(DescriptorSet { - device, - pool: key.pool, - handle: set, - }) + Ok(descriptor_set) } } diff --git a/src/resource/bindless.rs b/src/resource/bindless.rs new file mode 100644 index 0000000..8914b19 --- /dev/null +++ b/src/resource/bindless.rs @@ -0,0 +1,244 @@ +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 a handle to it. + fn into_bindless(self, pool: &BindlessPool) -> BindlessHandle + 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: Arc, + image_view: crate::image::ImageView, + image_layout: Option, +} + +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: Arc, 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) -> 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), + } + } +} + +/// A handle to a resource contained in a [BindlessPool] +pub struct BindlessHandle { + key: u32, + pool: BindlessPool, +} + +impl Drop for BindlessHandle { + fn drop(&mut self) { + self.pool.with(|p| p.take(self.key)); + } +} + +impl BindlessHandle { + /// Get the index of this resource in the pool's descripor. + /// This can be sent to a shader through any means to access it. + pub fn index(&self) -> u32 { + self.key + } +} + +pub(crate) struct BindlessPoolInner { + items: Vec>, + free: Vec, + pub(crate) descriptor_set: Arc, +} + +impl BindlessPoolInner { + fn update_descriptor_set<'a>(&'a mut self, r: impl Iterator) { + 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::>(); + + unsafe { + self.descriptor_set.device.update_descriptor_sets(vk_writes.as_slice(), &[]); + } + } + + + fn take(&mut self, key: u32) -> Option { + 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 that gives access to all resources. +pub struct BindlessPool { + inner: Arc>>, +} + +impl BindlessPool

{ + pub(crate) fn with) -> 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

{ + 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> { + self.with(|p| { + let mut keys = p.free.iter().cloned().rev().take(items.len()).collect::>(); + 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() } + } + }) + }) + } + + /// Create a new bindless pool + pub fn new(device: crate::Device) -> Result { + 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 { + 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 descriptor_set = Arc::new(descriptor_set); + + let inner = BindlessPoolInner { + items: vec![], + free: vec![], + descriptor_set, + }; + + Ok(Self { inner: Arc::new(Mutex::new(inner)) }) + } +} diff --git a/src/resource/mod.rs b/src/resource/mod.rs index 0aae812..786dc89 100644 --- a/src/resource/mod.rs +++ b/src/resource/mod.rs @@ -6,3 +6,4 @@ pub mod pool; pub mod query_pool; pub mod raytracing; pub mod sampler; +pub mod bindless;