vulk_ext/vkx/
descriptor.rs

1use super::*;
2
3const DESCRIPTOR_MAX_SIZE: usize = 16;
4
5pub(crate) fn validate_descriptor_sizes(
6    p: &vk::PhysicalDeviceDescriptorBufferPropertiesEXT,
7) -> Result<()> {
8    ensure!(DESCRIPTOR_MAX_SIZE >= p.buffer_capture_replay_descriptor_data_size);
9    ensure!(DESCRIPTOR_MAX_SIZE >= p.image_capture_replay_descriptor_data_size);
10    ensure!(DESCRIPTOR_MAX_SIZE >= p.image_view_capture_replay_descriptor_data_size);
11    ensure!(DESCRIPTOR_MAX_SIZE >= p.sampler_capture_replay_descriptor_data_size);
12    ensure!(DESCRIPTOR_MAX_SIZE >= p.acceleration_structure_capture_replay_descriptor_data_size);
13    ensure!(DESCRIPTOR_MAX_SIZE >= p.sampler_descriptor_size);
14    ensure!(DESCRIPTOR_MAX_SIZE >= p.combined_image_sampler_descriptor_size);
15    ensure!(DESCRIPTOR_MAX_SIZE >= p.sampled_image_descriptor_size);
16    ensure!(DESCRIPTOR_MAX_SIZE >= p.storage_image_descriptor_size);
17    ensure!(DESCRIPTOR_MAX_SIZE >= p.uniform_texel_buffer_descriptor_size);
18    ensure!(DESCRIPTOR_MAX_SIZE >= p.robust_uniform_texel_buffer_descriptor_size);
19    ensure!(DESCRIPTOR_MAX_SIZE >= p.storage_texel_buffer_descriptor_size);
20    ensure!(DESCRIPTOR_MAX_SIZE >= p.robust_storage_texel_buffer_descriptor_size);
21    ensure!(DESCRIPTOR_MAX_SIZE >= p.uniform_buffer_descriptor_size);
22    ensure!(DESCRIPTOR_MAX_SIZE >= p.robust_uniform_buffer_descriptor_size);
23    ensure!(DESCRIPTOR_MAX_SIZE >= p.storage_buffer_descriptor_size);
24    ensure!(DESCRIPTOR_MAX_SIZE >= p.robust_storage_buffer_descriptor_size);
25    ensure!(DESCRIPTOR_MAX_SIZE >= p.input_attachment_descriptor_size);
26    ensure!(DESCRIPTOR_MAX_SIZE >= p.acceleration_structure_descriptor_size);
27    Ok(())
28}
29
30#[derive(Clone, Copy, Debug)]
31pub enum DescriptorCreateInfo {
32    UniformBuffer {
33        address: vk::DeviceAddress,
34        range: vk::DeviceSize,
35    },
36    StorageBuffer {
37        address: vk::DeviceAddress,
38        range: vk::DeviceSize,
39    },
40    SampledImage {
41        image_view: vk::ImageView,
42        image_layout: vk::ImageLayout,
43    },
44    StorageImage {
45        image_view: vk::ImageView,
46        image_layout: vk::ImageLayout,
47    },
48    InputAttachment {
49        image_view: vk::ImageView,
50        image_layout: vk::ImageLayout,
51    },
52    Sampler(vk::Sampler),
53    AccelerationStructure(vk::DeviceAddress),
54}
55
56impl DescriptorCreateInfo {
57    fn size(&self, props: &vk::PhysicalDeviceDescriptorBufferPropertiesEXT) -> usize {
58        match self {
59            DescriptorCreateInfo::UniformBuffer { .. } => props.uniform_buffer_descriptor_size,
60            DescriptorCreateInfo::StorageBuffer { .. } => props.storage_buffer_descriptor_size,
61            DescriptorCreateInfo::SampledImage { .. } => props.sampled_image_descriptor_size,
62            DescriptorCreateInfo::StorageImage { .. } => props.storage_image_descriptor_size,
63            DescriptorCreateInfo::InputAttachment { .. } => props.input_attachment_descriptor_size,
64            DescriptorCreateInfo::Sampler(_) => props.sampler_descriptor_size,
65            DescriptorCreateInfo::AccelerationStructure(_) => {
66                props.acceleration_structure_descriptor_size
67            }
68        }
69    }
70
71    fn ty(&self) -> vk::DescriptorType {
72        match self {
73            DescriptorCreateInfo::UniformBuffer { .. } => vk::DescriptorType::UniformBuffer,
74            DescriptorCreateInfo::StorageBuffer { .. } => vk::DescriptorType::StorageBuffer,
75            DescriptorCreateInfo::SampledImage { .. } => vk::DescriptorType::SampledImage,
76            DescriptorCreateInfo::StorageImage { .. } => vk::DescriptorType::StorageImage,
77            DescriptorCreateInfo::InputAttachment { .. } => vk::DescriptorType::InputAttachment,
78            DescriptorCreateInfo::Sampler(_) => vk::DescriptorType::Sampler,
79            DescriptorCreateInfo::AccelerationStructure(_) => {
80                vk::DescriptorType::AccelerationStructureKHR
81            }
82        }
83    }
84}
85
86type DescriptorData = [u8; DESCRIPTOR_MAX_SIZE];
87
88#[derive(Clone, Copy)]
89pub struct Descriptor {
90    ty: vk::DescriptorType,
91    size: usize,
92    data: DescriptorData,
93}
94
95impl Descriptor {
96    #[must_use]
97    pub unsafe fn create(
98        physical_device: &PhysicalDevice,
99        device: &Device,
100        create_info: DescriptorCreateInfo,
101    ) -> Self {
102        // Descriptor info.
103        let props = physical_device.descriptor_buffer_properties_ext;
104        let size = create_info.size(&props);
105        let ty = create_info.ty();
106
107        // Get descriptor data.
108        let data = match create_info {
109            DescriptorCreateInfo::UniformBuffer { address, range } => Self::get_descriptor_data(
110                device,
111                ty,
112                size,
113                vk::DescriptorDataEXT {
114                    p_uniform_buffer: &vk::DescriptorAddressInfoEXT {
115                        s_type: vk::StructureType::DescriptorAddressInfoEXT,
116                        p_next: null_mut(),
117                        address,
118                        range,
119                        format: vk::Format::Undefined,
120                    },
121                },
122            ),
123            DescriptorCreateInfo::StorageBuffer { address, range } => Self::get_descriptor_data(
124                device,
125                ty,
126                size,
127                vk::DescriptorDataEXT {
128                    p_storage_buffer: &vk::DescriptorAddressInfoEXT {
129                        s_type: vk::StructureType::DescriptorAddressInfoEXT,
130                        p_next: null_mut(),
131                        address,
132                        range,
133                        format: vk::Format::Undefined,
134                    },
135                },
136            ),
137            DescriptorCreateInfo::SampledImage {
138                image_view,
139                image_layout,
140            } => Self::get_descriptor_data(
141                device,
142                ty,
143                size,
144                vk::DescriptorDataEXT {
145                    p_sampled_image: &vk::DescriptorImageInfo {
146                        sampler: vk::Sampler::null(),
147                        image_view,
148                        image_layout,
149                    },
150                },
151            ),
152            DescriptorCreateInfo::StorageImage {
153                image_view,
154                image_layout,
155            } => Self::get_descriptor_data(
156                device,
157                ty,
158                size,
159                vk::DescriptorDataEXT {
160                    p_storage_image: &vk::DescriptorImageInfo {
161                        sampler: vk::Sampler::null(),
162                        image_view,
163                        image_layout,
164                    },
165                },
166            ),
167            DescriptorCreateInfo::InputAttachment {
168                image_view,
169                image_layout,
170            } => Self::get_descriptor_data(
171                device,
172                ty,
173                size,
174                vk::DescriptorDataEXT {
175                    p_input_attachment_image: &vk::DescriptorImageInfo {
176                        sampler: vk::Sampler::null(),
177                        image_view,
178                        image_layout,
179                    },
180                },
181            ),
182            DescriptorCreateInfo::Sampler(sampler) => Self::get_descriptor_data(
183                device,
184                ty,
185                size,
186                vk::DescriptorDataEXT {
187                    p_sampler: &raw const sampler,
188                },
189            ),
190            DescriptorCreateInfo::AccelerationStructure(acceleration_structure) => {
191                Self::get_descriptor_data(
192                    device,
193                    ty,
194                    size,
195                    vk::DescriptorDataEXT {
196                        acceleration_structure,
197                    },
198                )
199            }
200        };
201
202        Self { ty, size, data }
203    }
204
205    unsafe fn get_descriptor_data(
206        device: &Device,
207        ty: vk::DescriptorType,
208        size: usize,
209        data: vk::DescriptorDataEXT,
210    ) -> DescriptorData {
211        let mut descriptor = MaybeUninit::<DescriptorData>::zeroed();
212        device.get_descriptor_ext(
213            &vk::DescriptorGetInfoEXT {
214                s_type: vk::StructureType::DescriptorGetInfoEXT,
215                p_next: null(),
216                ty,
217                data,
218            },
219            size,
220            descriptor.as_mut_ptr().cast(),
221        );
222        descriptor.assume_init()
223    }
224}
225
226impl std::fmt::Debug for Descriptor {
227    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228        let slice = &self.data[0..self.size];
229        f.debug_struct("Descriptor")
230            .field("ty", &self.ty)
231            .field("size", &self.size)
232            .field("data", &slice)
233            .finish()
234    }
235}
236
237#[derive(Debug)]
238pub struct DescriptorBinding<'a> {
239    pub ty: vk::DescriptorType,
240    pub stages: vk::ShaderStageFlags,
241    pub descriptors: &'a [Descriptor],
242}
243
244pub struct DescriptorStorage {
245    buffer: vk::Buffer,
246    allocations: BufferAllocations,
247    pub(super) allocation: BufferAllocation,
248    set_layout: vk::DescriptorSetLayout,
249    pub(super) set_count: u32,
250    pub(super) buffer_indices: Vec<u32>,
251    pub(super) offsets: Vec<vk::DeviceSize>,
252    pub(super) push_constant_range: Option<vk::PushConstantRange>,
253    pub(super) pipeline_layout: vk::PipelineLayout,
254    pub(super) usage: vk::BufferUsageFlags,
255}
256
257impl DescriptorStorage {
258    pub unsafe fn create(
259        physical_device: &PhysicalDevice,
260        device: &Device,
261        bindings: &[DescriptorBinding],
262        push_constant_range: Option<vk::PushConstantRange>,
263    ) -> Result<Self> {
264        // Validation.
265        ensure!(!bindings.is_empty(), "Expected 1 or more bindings");
266        for (binding_index, binding) in bindings.iter().enumerate() {
267            ensure!(
268                !binding.descriptors.is_empty(),
269                "Binding {binding_index} expected 1 or more descriptors",
270            );
271            for descriptor in binding.descriptors {
272                ensure!(
273                    binding.ty == descriptor.ty,
274                    "Binding {} expected descriptor type to be equal to {:?}, got {:?} instead",
275                    binding_index,
276                    binding.ty,
277                    descriptor.ty
278                );
279            }
280        }
281
282        // Descriptor set layout.
283        let set_layout_bindings = bindings
284            .iter()
285            .enumerate()
286            .map(|(binding_index, binding)| vk::DescriptorSetLayoutBinding {
287                binding: binding_index as _,
288                descriptor_type: binding.ty,
289                descriptor_count: binding.descriptors.len() as _,
290                stage_flags: binding.stages,
291                p_immutable_samplers: null(),
292            })
293            .collect::<Vec<_>>();
294        let set_layout =
295            device.create_descriptor_set_layout(&vk::DescriptorSetLayoutCreateInfo {
296                s_type: vk::StructureType::DescriptorSetLayoutCreateInfo,
297                p_next: null(),
298                flags: vk::DescriptorSetLayoutCreateFlagBits::DescriptorBufferEXT.into(),
299                binding_count: set_layout_bindings.len() as _,
300                p_bindings: set_layout_bindings.as_ptr(),
301            })?;
302        let set_count = 1;
303        let buffer_indices = vec![0];
304        let offsets = vec![0];
305        let size = device.get_descriptor_set_layout_size_ext(set_layout);
306
307        // Buffer usage.
308        let usage = vk::BufferUsageFlagBits::ResourceDescriptorBufferEXT
309            | vk::BufferUsageFlagBits::SamplerDescriptorBufferEXT;
310
311        // Buffer.
312        let (buffer, buffer_create_info) = BufferCreator::new(size, usage)
313            .create(device)
314            .context("Creating buffer object")?;
315
316        // Allocate.
317        let allocations = BufferAllocations::allocate(
318            physical_device,
319            device,
320            &[buffer],
321            &[buffer_create_info],
322            vk::MemoryPropertyFlagBits::HostVisible | vk::MemoryPropertyFlagBits::HostCoherent,
323        )?;
324        let allocation = allocations.allocations()[0];
325
326        // Write descriptors.
327        for (binding_index, binding) in bindings.iter().enumerate() {
328            let binding_index = binding_index as u32;
329            let descriptor_offset =
330                device.get_descriptor_set_layout_binding_offset_ext(set_layout, binding_index);
331            let descriptor_offset = descriptor_offset as usize;
332            for (array_index, descriptor) in binding.descriptors.iter().enumerate() {
333                let dst_offset = descriptor_offset + array_index * descriptor.size;
334                std::ptr::copy_nonoverlapping(
335                    descriptor.data.as_ptr(),
336                    allocation.as_mut_ptr::<u8>().add(dst_offset),
337                    descriptor.size,
338                );
339            }
340        }
341
342        // Pipeline layout.
343        let pipeline_layout = {
344            let mut create_info = vk::PipelineLayoutCreateInfo {
345                s_type: vk::StructureType::PipelineLayoutCreateInfo,
346                p_next: null(),
347                flags: vk::PipelineLayoutCreateFlags::empty(),
348                set_layout_count: 1,
349                p_set_layouts: &raw const set_layout,
350                push_constant_range_count: 0,
351                p_push_constant_ranges: null(),
352            };
353            let mut pcr: vk::PushConstantRange = zeroed();
354            if let Some(push_constant_range) = &push_constant_range {
355                pcr.stage_flags = push_constant_range.stage_flags;
356                pcr.size = push_constant_range.size;
357                pcr.offset = push_constant_range.offset;
358                create_info.push_constant_range_count = 1;
359                create_info.p_push_constant_ranges = &raw const pcr;
360            }
361            device.create_pipeline_layout(&raw const create_info)?
362        };
363
364        Ok(Self {
365            buffer,
366            allocations,
367            allocation,
368            set_layout,
369            set_count,
370            buffer_indices,
371            offsets,
372            push_constant_range,
373            pipeline_layout,
374            usage,
375        })
376    }
377
378    pub unsafe fn destroy(self, device: &Device) {
379        device.destroy_pipeline_layout(self.pipeline_layout);
380        device.destroy_descriptor_set_layout(self.set_layout);
381        device.destroy_buffer(self.buffer);
382        self.allocations.free(device);
383    }
384
385    pub unsafe fn bind(&self, device: &Device, cmd: vk::CommandBuffer) {
386        device.cmd_bind_descriptor_buffers_ext(
387            cmd,
388            1,
389            &vk::DescriptorBufferBindingInfoEXT {
390                s_type: vk::StructureType::DescriptorBufferBindingInfoEXT,
391                p_next: null_mut(),
392                address: self.allocation.device_address(),
393                usage: self.usage,
394            },
395        );
396    }
397
398    pub unsafe fn set_offsets(
399        &self,
400        device: &Device,
401        cmd: vk::CommandBuffer,
402        pipeline_bind_point: vk::PipelineBindPoint,
403    ) {
404        device.cmd_set_descriptor_buffer_offsets_ext(
405            cmd,
406            pipeline_bind_point,
407            self.pipeline_layout,
408            0,
409            self.set_count,
410            self.buffer_indices.as_ptr(),
411            self.offsets.as_ptr(),
412        );
413    }
414
415    pub unsafe fn push_constants<T>(
416        &self,
417        device: &Device,
418        cmd: vk::CommandBuffer,
419        data: &T,
420    ) -> Result<()> {
421        let Some(pcr) = self.push_constant_range else {
422            bail!("Missing push constant range");
423        };
424        ensure!(pcr.size as usize == size_of::<T>());
425        device.cmd_push_constants(
426            cmd,
427            self.pipeline_layout,
428            pcr.stage_flags,
429            pcr.offset,
430            pcr.size,
431            std::ptr::from_ref::<T>(data).cast(),
432        );
433        Ok(())
434    }
435
436    #[must_use]
437    pub fn pipeline_layout(&self) -> vk::PipelineLayout {
438        self.pipeline_layout
439    }
440
441    #[must_use]
442    pub fn set_layouts(&self) -> &[vk::DescriptorSetLayout] {
443        std::slice::from_ref(&self.set_layout)
444    }
445
446    #[must_use]
447    pub fn push_constant_ranges(&self) -> &[vk::PushConstantRange] {
448        if let Some(pcr) = &self.push_constant_range {
449            std::slice::from_ref(pcr)
450        } else {
451            &[]
452        }
453    }
454}
vulk_ext/vkx/descriptor.rs

vulk_ext/vkx/
descriptor.rs
