Bindless Descriptor fixes and optimizations (#2515)

* more SmallVec use in DescriptorPool, to improve variable descriptor count allocation

* descriptor updates with no writes or copies return early

* add DescriptorSet::invalidate()

* make InvalidateDescriptorSet pub

* allow descriptor bindings with update_after_bind or partially_bound to be unbound on recording a draw/dispatch command

---------

Co-authored-by: Firestar99 <4696087-firestar99@users.noreply.gitlab.com>

vulkano/src/command_buffer/commands/pipeline.rs

@@ -10,8 +10,8 @@ use crate::{
         DrawMeshTasksIndirectCommand, RecordingCommandBuffer, ResourceInCommand, SubpassContents,
     },
     descriptor_set::{
-        layout::DescriptorType, DescriptorBindingResources, DescriptorBufferInfo,
-        DescriptorImageViewInfo,
+        layout::{DescriptorBindingFlags, DescriptorType},
+        DescriptorBindingResources, DescriptorBufferInfo, DescriptorImageViewInfo,
     },
     device::{DeviceOwned, QueueFlags},
     format::{FormatFeatures, NumericType},
@@ -2146,14 +2146,19 @@ impl RecordingCommandBuffer {
                 Ok(())
             };
 
+            let flags_skip_binding_validation =
+                DescriptorBindingFlags::UPDATE_AFTER_BIND | DescriptorBindingFlags::PARTIALLY_BOUND;
+            let requires_binding_validation =
+                (layout_binding.binding_flags & flags_skip_binding_validation).is_empty();
+            if requires_binding_validation {
                 let set_resources = descriptor_set_state
                     .descriptor_sets
                     .get(&set_num)
                     .ok_or_else(|| {
                         Box::new(ValidationError {
                             problem: format!(
-                            "the currently bound pipeline accesses descriptor set {set_num}, but \
-                        no descriptor set was previously bound"
+                                "the currently bound pipeline accesses descriptor set {set_num}, \
+                                 but no descriptor set was previously bound"
                             )
                             .into(),
                             // vuids?
@@ -2227,8 +2232,8 @@ impl RecordingCommandBuffer {
                     }
                     // Spec:
                     // Descriptor bindings with descriptor type of
-                // VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK can be undefined when
-                // the descriptor set is consumed; though values in that block will be undefined.
+                    // VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK can be undefined when the descriptor
+                    // set is consumed; though values in that block will be undefined.
                     //
                     // TODO: We *may* still want to validate this?
                     DescriptorBindingResources::InlineUniformBlock => (),
@@ -2244,6 +2249,7 @@ impl RecordingCommandBuffer {
                     }
                 }
             }
+        }
 
         Ok(())
     }

vulkano/src/descriptor_set/mod.rs

@@ -75,8 +75,8 @@ use self::{
 pub use self::{
     collection::DescriptorSetsCollection,
     update::{
-        CopyDescriptorSet, DescriptorBufferInfo, DescriptorImageViewInfo, WriteDescriptorSet,
-        WriteDescriptorSetElements,
+        CopyDescriptorSet, DescriptorBufferInfo, DescriptorImageViewInfo, InvalidateDescriptorSet,
+        WriteDescriptorSet, WriteDescriptorSetElements,
     },
 };
 use crate::{
@@ -191,6 +191,9 @@ impl DescriptorSet {
     ) -> Result<(), Box<ValidationError>> {
         let descriptor_writes: SmallVec<[_; 8]> = descriptor_writes.into_iter().collect();
         let descriptor_copies: SmallVec<[_; 8]> = descriptor_copies.into_iter().collect();
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return Ok(());
+        }
 
         self.inner
             .validate_update(&descriptor_writes, &descriptor_copies)?;
@@ -215,6 +218,9 @@ impl DescriptorSet {
     ) {
         let descriptor_writes: SmallVec<[_; 8]> = descriptor_writes.into_iter().collect();
         let descriptor_copies: SmallVec<[_; 8]> = descriptor_copies.into_iter().collect();
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return;
+        }
 
         Self::update_inner(
             &self.inner,
@@ -236,6 +242,9 @@ impl DescriptorSet {
     ) -> Result<(), Box<ValidationError>> {
         let descriptor_writes: SmallVec<[_; 8]> = descriptor_writes.into_iter().collect();
         let descriptor_copies: SmallVec<[_; 8]> = descriptor_copies.into_iter().collect();
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return Ok(());
+        }
 
         self.inner
             .validate_update(&descriptor_writes, &descriptor_copies)?;
@@ -258,6 +267,9 @@ impl DescriptorSet {
     ) {
         let descriptor_writes: SmallVec<[_; 8]> = descriptor_writes.into_iter().collect();
         let descriptor_copies: SmallVec<[_; 8]> = descriptor_copies.into_iter().collect();
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return;
+        }
 
         Self::update_inner(
             &self.inner,
@@ -283,6 +295,42 @@ impl DescriptorSet {
             resources.copy(copy);
         }
     }
+
+    /// Invalidates descriptors within a descriptor set. Doesn't actually call into vulkan and only
+    /// invalidates the descriptors inside vulkano's resource tracking. Invalidated descriptors are
+    /// equivalent to uninitialized descriptors, in that binding a descriptor set to a particular
+    /// pipeline requires all shader-accessible descriptors to be valid.
+    ///
+    /// The intended use-case is an update-after-bind or bindless system, where entries in an
+    /// arrayed binding have to be invalidated so that the backing resource will be freed, and
+    /// not stay forever referenced until overridden by some update.
+    pub fn invalidate(
+        &self,
+        descriptor_invalidates: &[InvalidateDescriptorSet],
+    ) -> Result<(), Box<ValidationError>> {
+        self.validate_invalidate(descriptor_invalidates)?;
+        self.invalidate_unchecked(descriptor_invalidates);
+        Ok(())
+    }
+
+    pub fn invalidate_unchecked(&self, descriptor_invalidates: &[InvalidateDescriptorSet]) {
+        let mut resources = self.resources.write();
+        for invalidate in descriptor_invalidates {
+            resources.invalidate(invalidate);
+        }
+    }
+
+    pub(super) fn validate_invalidate(
+        &self,
+        descriptor_invalidates: &[InvalidateDescriptorSet],
+    ) -> Result<(), Box<ValidationError>> {
+        for (index, write) in descriptor_invalidates.iter().enumerate() {
+            write
+                .validate(self.layout(), self.variable_descriptor_count())
+                .map_err(|err| err.add_context(format!("descriptor_writes[{}]", index)))?;
+        }
+        Ok(())
+    }
 }
 
 unsafe impl VulkanObject for DescriptorSet {
@@ -432,6 +480,14 @@ impl DescriptorSetResources {
                 copy.descriptor_count,
             );
     }
+
+    #[inline]
+    pub(crate) fn invalidate(&mut self, invalidate: &InvalidateDescriptorSet) {
+        self.binding_resources
+            .get_mut(&invalidate.binding)
+            .expect("descriptor write has invalid binding number")
+            .invalidate(invalidate)
+    }
 }
 
 /// The resources that are bound to a single descriptor set binding.
@@ -621,6 +677,47 @@ impl DescriptorBindingResources {
             },
         }
     }
+
+    pub(crate) fn invalidate(&mut self, invalidate: &InvalidateDescriptorSet) {
+        fn invalidate_resources<T: Clone>(
+            resources: &mut [Option<T>],
+            invalidate: &InvalidateDescriptorSet,
+        ) {
+            let first = invalidate.first_array_element as usize;
+            resources
+                .get_mut(first..first + invalidate.descriptor_count as usize)
+                .expect("descriptor write for binding out of bounds")
+                .iter_mut()
+                .for_each(|resource| {
+                    *resource = None;
+                });
+        }
+
+        match self {
+            DescriptorBindingResources::None(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::Buffer(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::BufferView(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::ImageView(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::ImageViewSampler(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::Sampler(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+            DescriptorBindingResources::InlineUniformBlock => (),
+            DescriptorBindingResources::AccelerationStructure(resources) => {
+                invalidate_resources(resources, invalidate)
+            }
+        }
+    }
 }
 
 #[derive(Clone)]

vulkano/src/descriptor_set/pool.rs

@@ -250,9 +250,10 @@ impl DescriptorPool {
         let allocate_infos = allocate_infos.into_iter();
 
         let (lower_size_bound, _) = allocate_infos.size_hint();
-        let mut layouts_vk = Vec::with_capacity(lower_size_bound);
-        let mut variable_descriptor_counts = Vec::with_capacity(lower_size_bound);
-        let mut layouts = Vec::with_capacity(lower_size_bound);
+        let mut layouts_vk: SmallVec<[_; 1]> = SmallVec::with_capacity(lower_size_bound);
+        let mut variable_descriptor_counts: SmallVec<[_; 1]> =
+            SmallVec::with_capacity(lower_size_bound);
+        let mut layouts: SmallVec<[_; 1]> = SmallVec::with_capacity(lower_size_bound);
 
         for info in allocate_infos {
             let DescriptorSetAllocateInfo {
@@ -266,7 +267,7 @@ impl DescriptorPool {
             layouts.push(layout);
         }
 
-        let mut output = vec![];
+        let mut output: SmallVec<[_; 1]> = SmallVec::new();
 
         if !layouts_vk.is_empty() {
             let variable_desc_count_alloc_info = if (self.device.api_version() >= Version::V1_2

vulkano/src/descriptor_set/sys.rs

@@ -86,6 +86,10 @@ impl RawDescriptorSet {
         descriptor_writes: &[WriteDescriptorSet],
         descriptor_copies: &[CopyDescriptorSet],
     ) -> Result<(), Box<ValidationError>> {
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return Ok(());
+        }
+
         self.validate_update(descriptor_writes, descriptor_copies)?;
 
         self.update_unchecked(descriptor_writes, descriptor_copies);
@@ -117,6 +121,10 @@ impl RawDescriptorSet {
         descriptor_writes: &[WriteDescriptorSet],
         descriptor_copies: &[CopyDescriptorSet],
     ) {
+        if descriptor_writes.is_empty() && descriptor_copies.is_empty() {
+            return;
+        }
+
         struct PerDescriptorWrite {
             write_info: DescriptorWriteInfo,
             acceleration_structures: ash::vk::WriteDescriptorSetAccelerationStructureKHR<'static>,

vulkano/src/descriptor_set/update.rs

@@ -1599,6 +1599,8 @@ pub struct CopyDescriptorSet {
     pub dst_binding: u32,
 
     /// The first array element in the destination descriptor set to copy into.
+    ///
+    /// The default value is 0.
     pub dst_first_array_element: u32,
 
     /// The number of descriptors (array elements) to copy.
@@ -1832,3 +1834,93 @@ impl CopyDescriptorSet {
         Ok(())
     }
 }
+
+/// Invalidates descriptors within a descriptor set. Doesn't actually call into vulkan and only
+/// invalidates the descriptors inside vulkano's resource tracking. Invalidated descriptors are
+/// equivalent to uninitialized descriptors, in that binding a descriptor set to a particular
+/// pipeline requires all shader-accessible descriptors to be valid.
+///
+/// The intended use-case is an update-after-bind or bindless system, where entries in an arrayed
+/// binding have to be invalidated so that the backing resource will be freed, and not stay forever
+/// referenced until overridden by some update.
+pub struct InvalidateDescriptorSet {
+    /// The binding number in the descriptor set to invalidate.
+    ///
+    /// The default value is 0.
+    pub binding: u32,
+
+    /// The first array element in the descriptor set to invalidate.
+    ///
+    /// The default value is 0.
+    pub first_array_element: u32,
+
+    /// The number of descriptors (array elements) to invalidate.
+    ///
+    /// The default value is 1.
+    pub descriptor_count: u32,
+
+    pub _ne: crate::NonExhaustive,
+}
+
+impl InvalidateDescriptorSet {
+    pub fn invalidate(binding: u32) -> Self {
+        Self::invalidate_array(binding, 0, 1)
+    }
+
+    pub fn invalidate_array(binding: u32, first_array_element: u32, descriptor_count: u32) -> Self {
+        Self {
+            binding,
+            first_array_element,
+            descriptor_count,
+            _ne: crate::NonExhaustive(()),
+        }
+    }
+
+    pub(crate) fn validate(
+        &self,
+        layout: &DescriptorSetLayout,
+        variable_descriptor_count: u32,
+    ) -> Result<(), Box<ValidationError>> {
+        let &Self {
+            binding,
+            first_array_element,
+            descriptor_count,
+            ..
+        } = self;
+
+        let layout_binding = match layout.bindings().get(&binding) {
+            Some(layout_binding) => layout_binding,
+            None => {
+                return Err(Box::new(ValidationError {
+                    context: "binding".into(),
+                    problem: "does not exist in the descriptor set layout".into(),
+                    vuids: &["VUID-VkWriteDescriptorSet-dstBinding-00315"],
+                    ..Default::default()
+                }));
+            }
+        };
+
+        debug_assert!(descriptor_count != 0);
+        let max_descriptor_count = if layout_binding
+            .binding_flags
+            .intersects(DescriptorBindingFlags::VARIABLE_DESCRIPTOR_COUNT)
+        {
+            variable_descriptor_count
+        } else {
+            layout_binding.descriptor_count
+        };
+
+        // VUID-VkWriteDescriptorSet-dstArrayElement-00321
+        if first_array_element + descriptor_count > max_descriptor_count {
+            return Err(Box::new(ValidationError {
+                problem: "`first_array_element` + the number of provided elements is greater than \
+                    the number of descriptors in the descriptor set binding"
+                    .into(),
+                vuids: &["VUID-VkWriteDescriptorSet-dstArrayElement-00321"],
+                ..Default::default()
+            }));
+        }
+
+        Ok(())
+    }
+}