Add a StandardDescriptorSetAllocatorCreateInfo (#2366)

* Add `StandardDescriptorSetAllocatorCreateInfo`

* Fix tests and examples

examples/src/bin/async-update.rs

@@ -483,7 +483,8 @@ fn main() {
     };
     let mut framebuffers = window_size_dependent_setup(&images, render_pass.clone(), &mut viewport);
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
 
     // A byproduct of always using the same set of uniform buffers is that we can also create one
     // descriptor set for each, reusing them in the same way as the buffers.

examples/src/bin/basic-compute-shader.rs

@@ -160,7 +160,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/deferred/frame/system.rs

@@ -203,6 +203,7 @@ impl FrameSystem {
 
         let descriptor_set_allocator = Arc::new(StandardDescriptorSetAllocator::new(
             gfx_queue.device().clone(),
+            Default::default(),
         ));
 
         // Initialize the three lighting systems. Note that we need to pass to them the subpass

examples/src/bin/dynamic-buffers.rs

@@ -154,7 +154,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/dynamic-local-size.rs

@@ -211,7 +211,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/gl-interop.rs

@@ -261,7 +261,8 @@ mod linux {
             }
         });
 
-        let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+        let descriptor_set_allocator =
+            StandardDescriptorSetAllocator::new(device.clone(), Default::default());
         let command_buffer_allocator =
             StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/image-self-copy-blit/main.rs

@@ -210,7 +210,8 @@ fn main() {
     )
     .unwrap();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
     let mut uploads = AutoCommandBufferBuilder::primary(

examples/src/bin/image/main.rs

@@ -210,7 +210,8 @@ fn main() {
     )
     .unwrap();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
     let mut uploads = AutoCommandBufferBuilder::primary(

examples/src/bin/immutable-sampler/main.rs

@@ -216,7 +216,8 @@ fn main() {
     )
     .unwrap();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
     let mut uploads = AutoCommandBufferBuilder::primary(

examples/src/bin/indirect.rs

@@ -370,7 +370,8 @@ fn main() {
     let mut recreate_swapchain = false;
     let mut previous_frame_end = Some(sync::now(device.clone()).boxed());
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/interactive_fractal/app.rs

@@ -70,6 +70,7 @@ impl FractalApp {
         ));
         let descriptor_set_allocator = Arc::new(StandardDescriptorSetAllocator::new(
             gfx_queue.device().clone(),
+            Default::default(),
         ));
 
         FractalApp {

examples/src/bin/multi_window_game_of_life/app.rs

@@ -117,6 +117,7 @@ impl Default for App {
         ));
         let descriptor_set_allocator = Arc::new(StandardDescriptorSetAllocator::new(
             context.device().clone(),
+            Default::default(),
         ));
 
         App {

examples/src/bin/push-constants.rs

@@ -142,7 +142,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/runtime_array/main.rs

@@ -270,7 +270,8 @@ fn main() {
     )
     .unwrap();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
     let mut uploads = AutoCommandBufferBuilder::primary(

examples/src/bin/self-copy-buffer.rs

@@ -134,7 +134,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/shader-include/main.rs

@@ -142,7 +142,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/shader-types-sharing.rs

@@ -235,7 +235,8 @@ fn main() {
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
 
     // Prepare test array `[0, 1, 2, 3....]`.
     let data_buffer = Buffer::from_iter(

examples/src/bin/simple-particles.rs

@@ -327,7 +327,8 @@ fn main() {
     }
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/specialization-constants.rs

@@ -143,7 +143,8 @@ fn main() {
     };
 
     let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/teapot/main.rs

@@ -258,7 +258,8 @@ fn main() {
     let mut previous_frame_end = Some(sync::now(device.clone()).boxed());
     let rotation_start = Instant::now();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
 

examples/src/bin/texture_array/main.rs

@@ -212,7 +212,8 @@ fn main() {
     )
     .unwrap();
 
-    let descriptor_set_allocator = StandardDescriptorSetAllocator::new(device.clone());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
     let command_buffer_allocator =
         StandardCommandBufferAllocator::new(device.clone(), Default::default());
     let mut uploads = AutoCommandBufferBuilder::primary(

vulkano/src/command_buffer/auto/mod.rs

@@ -773,7 +773,8 @@ mod tests {
             )
             .unwrap();
 
-            let ds_allocator = StandardDescriptorSetAllocator::new(device.clone());
+            let ds_allocator =
+                StandardDescriptorSetAllocator::new(device.clone(), Default::default());
 
             let set = PersistentDescriptorSet::new(
                 &ds_allocator,

vulkano/src/descriptor_set/allocator.rs

@@ -35,8 +35,6 @@ use thread_local::ThreadLocal;
 
 const MAX_POOLS: usize = 32;
 
-const MAX_SETS: usize = 256;
-
 /// Types that manage the memory of descriptor sets.
 ///
 /// # Safety
@@ -96,6 +94,7 @@ pub trait DescriptorSetAlloc: Send + Sync {
 pub struct StandardDescriptorSetAllocator {
     device: InstanceOwnedDebugWrapper<Arc<Device>>,
     pools: ThreadLocal<UnsafeCell<SortedMap<NonZeroU64, Entry>>>,
+    create_info: StandardDescriptorSetAllocatorCreateInfo,
 }
 
 #[derive(Debug)]
@@ -112,10 +111,14 @@ unsafe impl Send for Entry {}
 impl StandardDescriptorSetAllocator {
     /// Creates a new `StandardDescriptorSetAllocator`.
     #[inline]
-    pub fn new(device: Arc<Device>) -> StandardDescriptorSetAllocator {
+    pub fn new(
+        device: Arc<Device>,
+        create_info: StandardDescriptorSetAllocatorCreateInfo,
+    ) -> StandardDescriptorSetAllocator {
         StandardDescriptorSetAllocator {
             device: InstanceOwnedDebugWrapper(device),
             pools: ThreadLocal::new(),
+            create_info,
         }
     }
 
@@ -154,15 +157,15 @@ unsafe impl DescriptorSetAllocator for StandardDescriptorSetAllocator {
 
         let entry = unsafe { &mut *pools.get() }.get_or_try_insert(layout.id(), || {
             if max_count == 0 {
-                FixedEntry::new(layout.clone()).map(Entry::Fixed)
+                FixedEntry::new(layout.clone(), &self.create_info).map(Entry::Fixed)
             } else {
-                VariableEntry::new(layout.clone()).map(Entry::Variable)
+                VariableEntry::new(layout.clone(), &self.create_info).map(Entry::Variable)
             }
         })?;
 
         match entry {
             Entry::Fixed(entry) => entry.allocate(),
-            Entry::Variable(entry) => entry.allocate(variable_descriptor_count),
+            Entry::Variable(entry) => entry.allocate(variable_descriptor_count, &self.create_info),
         }
     }
 }
@@ -199,10 +202,13 @@ struct FixedEntry {
 }
 
 impl FixedEntry {
-    fn new(layout: Arc<DescriptorSetLayout>) -> Result<Self, Validated<VulkanError>> {
+    fn new(
+        layout: Arc<DescriptorSetLayout>,
+        create_info: &StandardDescriptorSetAllocatorCreateInfo,
+    ) -> Result<Self, Validated<VulkanError>> {
         Ok(FixedEntry {
-            pool: FixedPool::new(&layout, MAX_SETS)?,
-            set_count: MAX_SETS,
+            pool: FixedPool::new(&layout, create_info.set_count)?,
+            set_count: create_info.set_count,
             layout,
         })
     }
@@ -304,11 +310,14 @@ struct VariableEntry {
 }
 
 impl VariableEntry {
-    fn new(layout: Arc<DescriptorSetLayout>) -> Result<Self, Validated<VulkanError>> {
+    fn new(
+        layout: Arc<DescriptorSetLayout>,
+        create_info: &StandardDescriptorSetAllocatorCreateInfo,
+    ) -> Result<Self, Validated<VulkanError>> {
         let reserve = Arc::new(ArrayQueue::new(MAX_POOLS));
 
         Ok(VariableEntry {
-            pool: VariablePool::new(&layout, reserve.clone())?,
+            pool: VariablePool::new(&layout, reserve.clone(), create_info.set_count)?,
             reserve,
             layout,
             allocations: 0,
@@ -318,15 +327,16 @@ impl VariableEntry {
     fn allocate(
         &mut self,
         variable_descriptor_count: u32,
+        create_info: &StandardDescriptorSetAllocatorCreateInfo,
     ) -> Result<StandardDescriptorSetAlloc, Validated<VulkanError>> {
-        if self.allocations >= MAX_SETS {
+        if self.allocations >= create_info.set_count {
             self.pool = if let Some(inner) = self.reserve.pop() {
                 Arc::new(VariablePool {
                     inner: ManuallyDrop::new(inner),
                     reserve: self.reserve.clone(),
                 })
             } else {
-                VariablePool::new(&self.layout, self.reserve.clone())?
+                VariablePool::new(&self.layout, self.reserve.clone(), create_info.set_count)?
             };
             self.allocations = 0;
         }
@@ -380,17 +390,18 @@ impl VariablePool {
     fn new(
         layout: &Arc<DescriptorSetLayout>,
         reserve: Arc<ArrayQueue<DescriptorPool>>,
+        set_count: usize,
     ) -> Result<Arc<Self>, VulkanError> {
         DescriptorPool::new(
             layout.device().clone(),
             DescriptorPoolCreateInfo {
-                max_sets: MAX_SETS as u32,
+                max_sets: set_count as u32,
                 pool_sizes: layout
                     .descriptor_counts()
                     .iter()
                     .map(|(&ty, &count)| {
                         assert!(ty != DescriptorType::InlineUniformBlock);
-                        (ty, count * MAX_SETS as u32)
+                        (ty, count * set_count as u32)
                     })
                     .collect(),
                 ..Default::default()
@@ -423,6 +434,40 @@ impl Drop for VariablePool {
     }
 }
 
+/// Parameters to create a new `StandardDescriptorSetAllocator`.
+#[derive(Clone, Debug)]
+pub struct StandardDescriptorSetAllocatorCreateInfo {
+    /// How many descriptor sets should be allocated per pool.
+    ///
+    /// Each time a thread allocates using some descriptor set layout, and either no pools were
+    /// initialized yet or all pools are full, a new pool is allocated for that thread and
+    /// descriptor set layout combination. This option tells the allocator how many descriptor sets
+    /// should be allocated for that pool. For fixed-size descriptor set layouts, it always
+    /// allocates exactly this many descriptor sets at once for the pool, as that is more
+    /// performant than allocating them one-by-one. For descriptor set layouts with a variable
+    /// descriptor count, it allocates a pool capable of holding exactly this many descriptor sets,
+    /// but doesn't allocate any descriptor sets since the variable count isn't known. What this
+    /// means is that you should make sure that this isn't too large, so that you don't end up
+    /// wasting too much memory. You also don't want this to be too low, because that on the other
+    /// hand would mean that the pool would have to be reset more often, or that more pools would
+    /// need to be created, depending on the lifetime of the descriptor sets.
+    ///
+    /// The default value is `256`.
+    pub set_count: usize,
+
+    pub _ne: crate::NonExhaustive,
+}
+
+impl Default for StandardDescriptorSetAllocatorCreateInfo {
+    #[inline]
+    fn default() -> Self {
+        StandardDescriptorSetAllocatorCreateInfo {
+            set_count: 256,
+            _ne: crate::NonExhaustive(()),
+        }
+    }
+}
+
 /// A descriptor set allocated from a [`StandardDescriptorSetAllocator`].
 #[derive(Debug)]
 pub struct StandardDescriptorSetAlloc {
@@ -554,7 +599,7 @@ mod tests {
         )
         .unwrap();
 
-        let allocator = StandardDescriptorSetAllocator::new(device);
+        let allocator = StandardDescriptorSetAllocator::new(device, Default::default());
 
         let pool1 =
             if let AllocParent::Fixed(pool) = &allocator.allocate(&layout, 0).unwrap().parent {

vulkano/src/pipeline/compute.rs

@@ -549,7 +549,7 @@ mod tests {
         )
         .unwrap();
 
-        let ds_allocator = StandardDescriptorSetAllocator::new(device.clone());
+        let ds_allocator = StandardDescriptorSetAllocator::new(device.clone(), Default::default());
         let set = PersistentDescriptorSet::new(
             &ds_allocator,
             pipeline.layout().set_layouts().get(0).unwrap().clone(),
@@ -680,7 +680,7 @@ mod tests {
         )
         .unwrap();
 
-        let ds_allocator = StandardDescriptorSetAllocator::new(device.clone());
+        let ds_allocator = StandardDescriptorSetAllocator::new(device.clone(), Default::default());
         let set = PersistentDescriptorSet::new(
             &ds_allocator,
             pipeline.layout().set_layouts().get(0).unwrap().clone(),