lgpu_hal/vulkan/
bindless.rs

1//! Bindless Descriptor Heap implementation for Vulkan 1.2+
2//!
3//! Provides a persistent, global descriptor array where buffers, textures,
4//! and samplers can be indexed directly by a 32-bit integer in shaders.
5
6use crate::HalError;
7use ash::vk;
8use std::sync::Mutex;
9
10pub const DEFAULT_MAX_STORAGE_BUFFERS: u32 = 16384;
11pub const DEFAULT_MAX_TEXTURES: u32 = 16384;
12pub const DEFAULT_MAX_SAMPLERS: u32 = 256;
13pub const DEFAULT_MAX_STORAGE_IMAGES: u32 = 16384;
14pub const PUSH_CONSTANT_SIZE: u32 = 128;
15
16pub struct VulkanBindlessHeap {
17    pub descriptor_set_layout: vk::DescriptorSetLayout,
18    pub descriptor_pool: vk::DescriptorPool,
19    pub descriptor_set: vk::DescriptorSet,
20    pub pipeline_layout: vk::PipelineLayout,
21    pub max_buffers: u32,
22    pub max_textures: u32,
23    pub max_samplers: u32,
24    pub max_storage_images: u32,
25    buffer_slots: Mutex<SlotPool>,
26    texture_slots: Mutex<SlotPool>,
27    sampler_slots: Mutex<SlotPool>,
28    storage_image_slots: Mutex<SlotPool>,
29}
30
31struct SlotPool {
32    free: Vec<u32>,
33    next: u32,
34}
35
36impl VulkanBindlessHeap {
37    pub fn new(
38        device: &ash::Device,
39        max_buffers: u32,
40        max_textures: u32,
41        max_samplers: u32,
42        max_storage_images: u32,
43    ) -> Result<Self, HalError> {
44        // 1. Descriptor Set Layout Bindings
45        // Binding 0: Storage Buffers array
46        // Binding 1: Sampled Images array
47        // Binding 2: Samplers array
48        // Binding 3: Storage Images array (writeable)
49        let bindings = [
50            vk::DescriptorSetLayoutBinding::default()
51                .binding(0)
52                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
53                .descriptor_count(max_buffers)
54                .stage_flags(vk::ShaderStageFlags::ALL),
55            vk::DescriptorSetLayoutBinding::default()
56                .binding(1)
57                .descriptor_type(vk::DescriptorType::SAMPLED_IMAGE)
58                .descriptor_count(max_textures)
59                .stage_flags(vk::ShaderStageFlags::ALL),
60            vk::DescriptorSetLayoutBinding::default()
61                .binding(2)
62                .descriptor_type(vk::DescriptorType::SAMPLER)
63                .descriptor_count(max_samplers)
64                .stage_flags(vk::ShaderStageFlags::ALL),
65            vk::DescriptorSetLayoutBinding::default()
66                .binding(3)
67                .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
68                .descriptor_count(max_storage_images)
69                .stage_flags(vk::ShaderStageFlags::ALL),
70        ];
71
72        let binding_flags = [
73            vk::DescriptorBindingFlags::PARTIALLY_BOUND
74                | vk::DescriptorBindingFlags::UPDATE_AFTER_BIND,
75            vk::DescriptorBindingFlags::PARTIALLY_BOUND
76                | vk::DescriptorBindingFlags::UPDATE_AFTER_BIND,
77            vk::DescriptorBindingFlags::PARTIALLY_BOUND
78                | vk::DescriptorBindingFlags::UPDATE_AFTER_BIND,
79            vk::DescriptorBindingFlags::PARTIALLY_BOUND
80                | vk::DescriptorBindingFlags::UPDATE_AFTER_BIND,
81        ];
82
83        let mut flags_info =
84            vk::DescriptorSetLayoutBindingFlagsCreateInfo::default().binding_flags(&binding_flags);
85
86        let layout_info = vk::DescriptorSetLayoutCreateInfo::default()
87            .bindings(&bindings)
88            .flags(vk::DescriptorSetLayoutCreateFlags::UPDATE_AFTER_BIND_POOL)
89            .push_next(&mut flags_info);
90
91        let descriptor_set_layout = unsafe {
92            device
93                .create_descriptor_set_layout(&layout_info, None)
94                .map_err(|e| {
95                    HalError::PipelineError(format!("Failed to create bindless set layout: {e}"))
96                })?
97        };
98
99        // 2. Global Pipeline Layout (with 128 bytes Push Constants)
100        let push_constant_ranges = [vk::PushConstantRange::default()
101            .stage_flags(vk::ShaderStageFlags::ALL)
102            .offset(0)
103            .size(PUSH_CONSTANT_SIZE)];
104
105        let set_layouts = [descriptor_set_layout];
106        let pipeline_layout_info = vk::PipelineLayoutCreateInfo::default()
107            .set_layouts(&set_layouts)
108            .push_constant_ranges(&push_constant_ranges);
109
110        let pipeline_layout = unsafe {
111            device
112                .create_pipeline_layout(&pipeline_layout_info, None)
113                .map_err(|e| {
114                    HalError::PipelineError(format!(
115                        "Failed to create bindless pipeline layout: {e}"
116                    ))
117                })?
118        };
119
120        // 3. Descriptor Pool
121        let pool_sizes = [
122            vk::DescriptorPoolSize::default()
123                .ty(vk::DescriptorType::STORAGE_BUFFER)
124                .descriptor_count(max_buffers),
125            vk::DescriptorPoolSize::default()
126                .ty(vk::DescriptorType::SAMPLED_IMAGE)
127                .descriptor_count(max_textures),
128            vk::DescriptorPoolSize::default()
129                .ty(vk::DescriptorType::SAMPLER)
130                .descriptor_count(max_samplers),
131            vk::DescriptorPoolSize::default()
132                .ty(vk::DescriptorType::STORAGE_IMAGE)
133                .descriptor_count(max_storage_images),
134        ];
135
136        let pool_info = vk::DescriptorPoolCreateInfo::default()
137            .pool_sizes(&pool_sizes)
138            .max_sets(1)
139            .flags(vk::DescriptorPoolCreateFlags::UPDATE_AFTER_BIND);
140
141        let descriptor_pool = unsafe {
142            device
143                .create_descriptor_pool(&pool_info, None)
144                .map_err(|e| {
145                    HalError::PipelineError(format!(
146                        "Failed to create bindless descriptor pool: {e}"
147                    ))
148                })?
149        };
150
151        // 4. Allocate the single global Descriptor Set
152        let alloc_info = vk::DescriptorSetAllocateInfo::default()
153            .descriptor_pool(descriptor_pool)
154            .set_layouts(&set_layouts);
155
156        let descriptor_sets = unsafe {
157            device.allocate_descriptor_sets(&alloc_info).map_err(|e| {
158                HalError::PipelineError(format!("Failed to allocate bindless descriptor set: {e}"))
159            })?
160        };
161
162        Ok(Self {
163            descriptor_set_layout,
164            descriptor_pool,
165            descriptor_set: descriptor_sets[0],
166            pipeline_layout,
167            max_buffers,
168            max_textures,
169            max_samplers,
170            max_storage_images,
171            buffer_slots: Mutex::new(SlotPool {
172                free: Vec::new(),
173                next: 0,
174            }),
175            texture_slots: Mutex::new(SlotPool {
176                free: Vec::new(),
177                next: 0,
178            }),
179            sampler_slots: Mutex::new(SlotPool {
180                free: Vec::new(),
181                next: 0,
182            }),
183            storage_image_slots: Mutex::new(SlotPool {
184                free: Vec::new(),
185                next: 0,
186            }),
187        })
188    }
189
190    pub fn allocate_buffer_slot(&self) -> Result<u32, HalError> {
191        let mut pool = self.buffer_slots.lock().unwrap();
192        if let Some(slot) = pool.free.pop() {
193            return Ok(slot);
194        }
195        if pool.next >= self.max_buffers {
196            return Err(HalError::OutOfMemory(
197                "Exceeded max bindless buffer capacity".to_string(),
198            ));
199        }
200        let slot = pool.next;
201        pool.next += 1;
202        Ok(slot)
203    }
204
205    pub fn free_buffer_slot(&self, slot: u32) {
206        let mut pool = self.buffer_slots.lock().unwrap();
207        pool.free.push(slot);
208    }
209
210    pub fn write_buffer_descriptor(
211        &self,
212        device: &ash::Device,
213        slot: u32,
214        buffer: vk::Buffer,
215        size: u64,
216    ) {
217        let buffer_info = vk::DescriptorBufferInfo::default()
218            .buffer(buffer)
219            .offset(0)
220            .range(size);
221
222        let write = vk::WriteDescriptorSet::default()
223            .dst_set(self.descriptor_set)
224            .dst_binding(0)
225            .dst_array_element(slot)
226            .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
227            .buffer_info(std::slice::from_ref(&buffer_info));
228
229        unsafe {
230            device.update_descriptor_sets(&[write], &[]);
231        }
232    }
233
234    pub fn allocate_texture_slot(&self) -> Result<u32, HalError> {
235        let mut pool = self.texture_slots.lock().unwrap();
236        if let Some(slot) = pool.free.pop() {
237            return Ok(slot);
238        }
239        if pool.next >= self.max_textures {
240            return Err(HalError::OutOfMemory(
241                "Exceeded max bindless texture capacity".to_string(),
242            ));
243        }
244        let slot = pool.next;
245        pool.next += 1;
246        Ok(slot)
247    }
248
249    pub fn free_texture_slot(&self, slot: u32) {
250        let mut pool = self.texture_slots.lock().unwrap();
251        pool.free.push(slot);
252    }
253
254    pub fn write_texture_descriptor(
255        &self,
256        device: &ash::Device,
257        slot: u32,
258        image_view: vk::ImageView,
259        image_layout: vk::ImageLayout,
260    ) {
261        let image_info = vk::DescriptorImageInfo::default()
262            .image_view(image_view)
263            .image_layout(image_layout);
264
265        let write = vk::WriteDescriptorSet::default()
266            .dst_set(self.descriptor_set)
267            .dst_binding(1)
268            .dst_array_element(slot)
269            .descriptor_type(vk::DescriptorType::SAMPLED_IMAGE)
270            .image_info(std::slice::from_ref(&image_info));
271
272        unsafe {
273            device.update_descriptor_sets(&[write], &[]);
274        }
275    }
276
277    pub fn allocate_sampler_slot(&self) -> Result<u32, HalError> {
278        let mut pool = self.sampler_slots.lock().unwrap();
279        if let Some(slot) = pool.free.pop() {
280            return Ok(slot);
281        }
282        if pool.next >= self.max_samplers {
283            return Err(HalError::OutOfMemory(
284                "Exceeded max bindless sampler capacity".to_string(),
285            ));
286        }
287        let slot = pool.next;
288        pool.next += 1;
289        Ok(slot)
290    }
291
292    pub fn free_sampler_slot(&self, slot: u32) {
293        let mut pool = self.sampler_slots.lock().unwrap();
294        pool.free.push(slot);
295    }
296
297    pub fn write_sampler_descriptor(&self, device: &ash::Device, slot: u32, sampler: vk::Sampler) {
298        let sampler_info = vk::DescriptorImageInfo::default().sampler(sampler);
299
300        let write = vk::WriteDescriptorSet::default()
301            .dst_set(self.descriptor_set)
302            .dst_binding(2)
303            .dst_array_element(slot)
304            .descriptor_type(vk::DescriptorType::SAMPLER)
305            .image_info(std::slice::from_ref(&sampler_info));
306
307        unsafe {
308            device.update_descriptor_sets(&[write], &[]);
309        }
310    }
311
312    pub fn allocate_storage_image_slot(&self) -> Result<u32, HalError> {
313        let mut pool = self.storage_image_slots.lock().unwrap();
314        if let Some(slot) = pool.free.pop() {
315            return Ok(slot);
316        }
317        if pool.next >= self.max_storage_images {
318            return Err(HalError::OutOfMemory(
319                "Exceeded max bindless storage image capacity".to_string(),
320            ));
321        }
322        let slot = pool.next;
323        pool.next += 1;
324        Ok(slot)
325    }
326
327    pub fn free_storage_image_slot(&self, slot: u32) {
328        let mut pool = self.storage_image_slots.lock().unwrap();
329        pool.free.push(slot);
330    }
331
332    pub fn write_storage_image_descriptor(
333        &self,
334        device: &ash::Device,
335        slot: u32,
336        image_view: vk::ImageView,
337        image_layout: vk::ImageLayout,
338    ) {
339        let image_info = vk::DescriptorImageInfo::default()
340            .image_view(image_view)
341            .image_layout(image_layout);
342
343        let write = vk::WriteDescriptorSet::default()
344            .dst_set(self.descriptor_set)
345            .dst_binding(3)
346            .dst_array_element(slot)
347            .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
348            .image_info(std::slice::from_ref(&image_info));
349
350        unsafe {
351            device.update_descriptor_sets(&[write], &[]);
352        }
353    }
354
355    pub fn write_buffer_descriptors_batch(
356        &self,
357        device: &ash::Device,
358        writes: &[(u32, vk::Buffer, u64)],
359    ) {
360        if writes.is_empty() {
361            return;
362        }
363        let buffer_infos: Vec<vk::DescriptorBufferInfo> = writes
364            .iter()
365            .map(|(_, buffer, size)| {
366                vk::DescriptorBufferInfo::default()
367                    .buffer(*buffer)
368                    .offset(0)
369                    .range(*size)
370            })
371            .collect();
372
373        let descriptor_writes: Vec<vk::WriteDescriptorSet> = writes
374            .iter()
375            .zip(buffer_infos.iter())
376            .map(|(&(slot, _, _), buffer_info)| {
377                vk::WriteDescriptorSet::default()
378                    .dst_set(self.descriptor_set)
379                    .dst_binding(0)
380                    .dst_array_element(slot)
381                    .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
382                    .buffer_info(std::slice::from_ref(buffer_info))
383            })
384            .collect();
385
386        unsafe {
387            device.update_descriptor_sets(&descriptor_writes, &[]);
388        }
389    }
390
391    pub fn write_texture_descriptors_batch(
392        &self,
393        device: &ash::Device,
394        writes: &[(u32, vk::ImageView, vk::ImageLayout)],
395    ) {
396        if writes.is_empty() {
397            return;
398        }
399        let image_infos: Vec<vk::DescriptorImageInfo> = writes
400            .iter()
401            .map(|(_, view, layout)| {
402                vk::DescriptorImageInfo::default()
403                    .image_view(*view)
404                    .image_layout(*layout)
405            })
406            .collect();
407
408        let descriptor_writes: Vec<vk::WriteDescriptorSet> = writes
409            .iter()
410            .zip(image_infos.iter())
411            .map(|(&(slot, _, _), image_info)| {
412                vk::WriteDescriptorSet::default()
413                    .dst_set(self.descriptor_set)
414                    .dst_binding(1)
415                    .dst_array_element(slot)
416                    .descriptor_type(vk::DescriptorType::SAMPLED_IMAGE)
417                    .image_info(std::slice::from_ref(image_info))
418            })
419            .collect();
420
421        unsafe {
422            device.update_descriptor_sets(&descriptor_writes, &[]);
423        }
424    }
425
426    pub fn write_storage_image_descriptors_batch(
427        &self,
428        device: &ash::Device,
429        writes: &[(u32, vk::ImageView, vk::ImageLayout)],
430    ) {
431        if writes.is_empty() {
432            return;
433        }
434        let image_infos: Vec<vk::DescriptorImageInfo> = writes
435            .iter()
436            .map(|(_, view, layout)| {
437                vk::DescriptorImageInfo::default()
438                    .image_view(*view)
439                    .image_layout(*layout)
440            })
441            .collect();
442
443        let descriptor_writes: Vec<vk::WriteDescriptorSet> = writes
444            .iter()
445            .zip(image_infos.iter())
446            .map(|(&(slot, _, _), image_info)| {
447                vk::WriteDescriptorSet::default()
448                    .dst_set(self.descriptor_set)
449                    .dst_binding(3)
450                    .dst_array_element(slot)
451                    .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
452                    .image_info(std::slice::from_ref(image_info))
453            })
454            .collect();
455
456        unsafe {
457            device.update_descriptor_sets(&descriptor_writes, &[]);
458        }
459    }
460
461    pub fn destroy(&self, device: &ash::Device) {
462        unsafe {
463            device.destroy_descriptor_pool(self.descriptor_pool, None);
464            device.destroy_descriptor_set_layout(self.descriptor_set_layout, None);
465            device.destroy_pipeline_layout(self.pipeline_layout, None);
466        }
467    }
468}
lgpu_hal/vulkan/bindless.rs

lgpu_hal/vulkan/
bindless.rs
