vk_graph/driver/
compute.rs

1//! Compute pipeline types.
2
3use {
4    super::{
5        DriverError,
6        device::Device,
7        shader::{DescriptorBindingMap, PipelineDescriptorInfo, Shader},
8    },
9    crate::lazy_str,
10    ash::vk::{self, Handle as _},
11    derive_builder::Builder,
12    log::{trace, warn},
13    std::{
14        ffi::CString,
15        fmt::{Debug, Formatter},
16        hash::{Hash, Hasher},
17        slice,
18        sync::Arc,
19        thread::panicking,
20    },
21};
22
23/// Smart pointer handle of a compute pipeline object.
24///
25/// Also contains information about the object.
26///
27/// See [`VkPipeline`](https://registry.khronos.org/vulkan/specs/latest/man/html/VkPipeline.html).
28#[derive(Clone)]
29pub struct ComputePipeline {
30    pub(crate) inner: Arc<ComputePipelineInner>,
31}
32
33impl ComputePipeline {
34    /// Creates a new compute pipeline on the given device.
35    ///
36    /// `shader` may be a pre-built [`Shader`] or any input that can be converted into one.
37    /// Invalid shader data is returned as [`DriverError::InvalidData`] through the `Result`
38    /// instead of panicking.
39    ///
40    /// See [`VkComputePipelineCreateInfo`](https://registry.khronos.org/vulkan/specs/latest/man/html/VkComputePipelineCreateInfo.html).
41    ///
42    /// # Examples
43    ///
44    /// Basic usage:
45    ///
46    /// ```no_run
47    /// # use std::sync::Arc;
48    /// # use ash::vk;
49    /// # use vk_graph::driver::DriverError;
50    /// # use vk_graph::driver::device::{Device, DeviceInfo};
51    /// # use vk_graph::driver::compute::{ComputePipeline, ComputePipelineInfo};
52    /// # use vk_graph::driver::shader::{Shader};
53    /// # fn main() -> Result<(), DriverError> {
54    /// # let device = Device::create(DeviceInfo::default())?;
55    /// # let my_shader_code = [0u8; 1];
56    /// // my_shader_code is raw SPIR-V code as bytes
57    /// let shader = Shader::new_compute(my_shader_code.as_slice());
58    /// let pipeline = ComputePipeline::create(&device, ComputePipelineInfo::default(), shader)?;
59    ///
60    /// assert_ne!(pipeline.handle(), vk::Pipeline::null());
61    /// # Ok(()) }
62    /// ```
63    #[profiling::function]
64    pub fn create<S>(
65        device: &Device,
66        info: impl Into<ComputePipelineInfo>,
67        shader: S,
68    ) -> Result<Self, DriverError>
69    where
70        S: TryInto<Shader>,
71        S::Error: Into<DriverError>,
72    {
73        trace!("create");
74
75        let info = info.into();
76        let shader = shader.try_into().map_err(Into::into)?;
77
78        // Use SPIR-V reflection to get the types and counts of all descriptors
79        let mut descriptor_bindings = shader.descriptor_bindings();
80        for (descriptor_info, _) in descriptor_bindings.values_mut() {
81            if descriptor_info.binding_count() == 0 {
82                descriptor_info.set_binding_count(info.bindless_descriptor_count);
83            }
84        }
85
86        let descriptor_info = PipelineDescriptorInfo::create(device, &descriptor_bindings)?;
87        let descriptor_set_layouts = descriptor_info
88            .layouts
89            .values()
90            .map(|descriptor_set_layout| descriptor_set_layout.handle)
91            .collect::<Box<_>>();
92
93        unsafe {
94            let shader_module = device
95                .create_shader_module(
96                    &vk::ShaderModuleCreateInfo::default().code(shader.spirv.words()),
97                    None,
98                )
99                .map_err(|err| {
100                    warn!("unable to create compute shader module: {err}");
101
102                    DriverError::Unsupported
103                })?;
104            let entry_name = CString::new(shader.entry_name.as_bytes()).map_err(|err| {
105                warn!("invalid compute shader entry name: {err}");
106
107                DriverError::InvalidData
108            })?;
109            let mut stage_create_info = vk::PipelineShaderStageCreateInfo::default()
110                .module(shader_module)
111                .stage(shader.stage)
112                .name(&entry_name);
113            let specialization_info = shader.specialization.as_ref().map(Into::into);
114
115            if let Some(specialization_info) = &specialization_info {
116                stage_create_info = stage_create_info.specialization_info(specialization_info);
117            }
118
119            let mut layout_info =
120                vk::PipelineLayoutCreateInfo::default().set_layouts(&descriptor_set_layouts);
121
122            let push_constants = shader.push_constant_range();
123            if let Some(push_constants) = &push_constants {
124                layout_info = layout_info.push_constant_ranges(slice::from_ref(push_constants));
125            }
126
127            let layout = device
128                .create_pipeline_layout(&layout_info, None)
129                .map_err(|err| {
130                    warn!("unable to create compute pipeline layout: {err}");
131
132                    device.destroy_shader_module(shader_module, None);
133
134                    DriverError::Unsupported
135                })?;
136            let create_info = vk::ComputePipelineCreateInfo::default()
137                .stage(stage_create_info)
138                .layout(layout);
139            let handle = device
140                .create_compute_pipelines(
141                    Device::pipeline_cache(device),
142                    slice::from_ref(&create_info),
143                    None,
144                )
145                .map_err(|(_, err)| {
146                    warn!("unable to create compute pipeline: {err}");
147
148                    device.destroy_shader_module(shader_module, None);
149
150                    DriverError::Unsupported
151                })?
152                .into_iter()
153                .find(|handle| !handle.is_null())
154                .ok_or_else(|| {
155                    warn!("missing pipeline handle");
156
157                    DriverError::Unsupported
158                })?;
159
160            device.destroy_shader_module(shader_module, None);
161
162            Ok(ComputePipeline {
163                inner: Arc::new(ComputePipelineInner {
164                    descriptor_bindings,
165                    descriptor_info,
166                    device: device.clone(),
167                    handle,
168                    info,
169                    layout,
170                    push_constants,
171                }),
172            })
173        }
174    }
175
176    /// The device which owns this compute pipeline.
177    pub fn device(&self) -> &Device {
178        &self.inner.device
179    }
180
181    /// The native Vulkan pipeline handle of this compute pipeline.
182    pub fn handle(&self) -> vk::Pipeline {
183        self.inner.handle
184    }
185
186    /// Gets the information used to create this object.
187    pub fn info(&self) -> ComputePipelineInfo {
188        self.inner.info
189    }
190
191    /// Sets the debugging name assigned to this pipeline.
192    pub fn set_debug_name(&self, name: impl AsRef<str>) {
193        Device::try_set_debug_utils_object_name(&self.inner.device, self.inner.handle, &name);
194        Device::try_set_private_data_object_name(
195            &self.inner.device,
196            vk::ObjectType::PIPELINE,
197            self.inner.handle,
198            &name,
199        );
200
201        Device::try_set_debug_utils_object_name(
202            &self.inner.device,
203            self.inner.layout,
204            lazy_str!("{} (layout)", name.as_ref()),
205        );
206
207        for (set_idx, layout) in &self.inner.descriptor_info.layouts {
208            layout.set_debug_name(lazy_str!("{} (DS{set_idx})", name.as_ref()));
209        }
210    }
211
212    /// Sets the debugging name assigned to this pipeline.
213    pub fn with_debug_name(self, name: impl AsRef<str>) -> Self {
214        self.set_debug_name(name);
215
216        self
217    }
218}
219
220impl Debug for ComputePipeline {
221    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
222        let mut res = f.debug_struct(stringify!(ComputePipeline));
223
224        if let Some(debug_name) = &Device::private_data_object_name(
225            &self.inner.device,
226            vk::ObjectType::PIPELINE,
227            self.inner.handle,
228        ) {
229            res.field("debug_name", debug_name);
230        }
231
232        res.field("handle", &self.inner.handle)
233            .finish_non_exhaustive()
234    }
235}
236
237impl Eq for ComputePipeline {}
238
239impl Hash for ComputePipeline {
240    fn hash<H: Hasher>(&self, state: &mut H) {
241        Arc::as_ptr(&self.inner).hash(state);
242    }
243}
244
245impl PartialEq for ComputePipeline {
246    fn eq(&self, other: &Self) -> bool {
247        Arc::ptr_eq(&self.inner, &other.inner)
248    }
249}
250
251/// Information used to create a [`ComputePipeline`] instance.
252///
253/// See [`VkComputePipelineCreateInfo`](https://registry.khronos.org/vulkan/specs/latest/man/html/VkComputePipelineCreateInfo.html).
254#[derive(Builder, Clone, Copy, Debug, Eq, Hash, PartialEq)]
255#[builder(
256    build_fn(private, name = "fallible_build"),
257    derive(Clone, Copy, Debug),
258    pattern = "owned"
259)]
260pub struct ComputePipelineInfo {
261    /// The number of descriptors to allocate for a given binding when using bindless (unbounded)
262    /// syntax.
263    ///
264    /// The default is `8192`.
265    ///
266    /// # Examples
267    ///
268    /// Basic usage (GLSL):
269    ///
270    /// ```
271    /// # vk_shader_macros::glsl!(r#"
272    /// #version 460 core
273    /// #extension GL_EXT_nonuniform_qualifier : require
274    /// #pragma shader_stage(compute)
275    ///
276    /// layout(set = 0, binding = 0, rgba8) writeonly uniform image2D my_binding[];
277    ///
278    /// void main()
279    /// {
280    ///     // my_binding will have space for 8,192 images by default
281    /// }
282    /// # "#);
283    /// ```
284    #[builder(default = "8192")]
285    pub bindless_descriptor_count: u32,
286}
287
288impl ComputePipelineInfo {
289    /// Creates a default `ComputePipelineInfoBuilder`.
290    pub fn builder() -> ComputePipelineInfoBuilder {
291        Default::default()
292    }
293
294    /// Converts a `ComputePipelineInfo` into a `ComputePipelineInfoBuilder`.
295    pub fn into_builder(self) -> ComputePipelineInfoBuilder {
296        ComputePipelineInfoBuilder {
297            bindless_descriptor_count: Some(self.bindless_descriptor_count),
298        }
299    }
300}
301
302impl Default for ComputePipelineInfo {
303    fn default() -> Self {
304        Self {
305            bindless_descriptor_count: 8192,
306        }
307    }
308}
309
310impl From<ComputePipelineInfoBuilder> for ComputePipelineInfo {
311    fn from(info: ComputePipelineInfoBuilder) -> Self {
312        info.build()
313    }
314}
315
316impl ComputePipelineInfoBuilder {
317    /// Builds a new `ComputePipelineInfo`.
318    #[inline(always)]
319    pub fn build(self) -> ComputePipelineInfo {
320        self.fallible_build()
321            .expect("invalid compute pipeline info")
322    }
323}
324
325#[derive(Debug)]
326pub(crate) struct ComputePipelineInner {
327    pub descriptor_bindings: DescriptorBindingMap,
328    pub descriptor_info: PipelineDescriptorInfo,
329    pub device: Device,
330    pub handle: vk::Pipeline,
331    pub info: ComputePipelineInfo,
332    pub layout: vk::PipelineLayout,
333    pub push_constants: Option<vk::PushConstantRange>,
334}
335
336impl Drop for ComputePipelineInner {
337    #[profiling::function]
338    fn drop(&mut self) {
339        if panicking() {
340            return;
341        }
342
343        Device::try_clear_private_data_object_name(
344            &self.device,
345            vk::ObjectType::PIPELINE,
346            self.handle,
347        );
348
349        unsafe {
350            self.device.destroy_pipeline(self.handle, None);
351            self.device.destroy_pipeline_layout(self.layout, None);
352        }
353    }
354}
355
356#[cfg(test)]
357mod test {
358    use super::*;
359
360    type Info = ComputePipelineInfo;
361    type Builder = ComputePipelineInfoBuilder;
362
363    #[test]
364    pub fn compute_pipeline_info() {
365        let info = Info::default();
366        let builder = info.into_builder().build();
367
368        assert_eq!(info, builder);
369    }
370
371    #[test]
372    pub fn compute_pipeline_info_builder() {
373        let info = Info::default();
374        let builder = Builder::default().build();
375
376        assert_eq!(info, builder);
377    }
378}
vk_graph/driver/compute.rs

vk_graph/driver/
compute.rs
