vk_graph/cmd/
ray_tracing.rs

1use {
2    super::{PipelineCommand, cmd_ref::CommandRef},
3    crate::driver::{device::Device, ray_tracing::RayTracingPipeline},
4    ash::vk,
5    std::ops::Deref,
6};
7
8impl PipelineCommand<'_, RayTracingPipeline> {
9    /// Begin recording ray tracing pipeline work for this graph command.
10    pub fn record_cmd(
11        mut self,
12        func: impl FnOnce(RayTracingCommandRef<'_>) + Send + 'static,
13    ) -> Self {
14        self.record_cmd_mut(func);
15        self
16    }
17
18    /// Mutable-borrow form of [`Self::record_cmd`].
19    pub fn record_cmd_mut(&mut self, func: impl FnOnce(RayTracingCommandRef<'_>) + Send + 'static) {
20        let pipeline = self
21            .cmd
22            .cmd()
23            .expect_last_pipeline()
24            .expect_ray_tracing()
25            .clone();
26
27        self.cmd.push_exec(move |cmd| {
28            func(RayTracingCommandRef { cmd, pipeline });
29        });
30    }
31
32    pub(crate) fn record_stream_mut(
33        &mut self,
34        func: impl for<'r> Fn(RayTracingCommandRef<'r>) + Send + Sync + 'static,
35    ) {
36        let pipeline = self
37            .cmd
38            .cmd()
39            .expect_last_pipeline()
40            .expect_ray_tracing()
41            .clone();
42
43        self.cmd.push_reusable_exec(move |cmd| {
44            func(RayTracingCommandRef {
45                cmd,
46                pipeline: pipeline.clone(),
47            });
48        });
49    }
50}
51
52/// Recording interface for ray tracing commands.
53///
54/// This structure provides a strongly-typed set of methods which allow ray tracing shader code to
55/// be executed. An instance is provided to the closure argument of [`PipelineCommand::record_cmd`]
56/// which may be accessed by binding a [`RayTracingPipeline`] to a command.
57///
58/// # Examples
59///
60/// Basic usage:
61///
62/// ```no_run
63/// # use ash::vk;
64/// # use vk_graph::driver::DriverError;
65/// # use vk_graph::driver::device::{Device, DeviceInfo};
66/// # use vk_graph::driver::ray_tracing::{
67/// #     RayTracingPipeline,
68/// #     RayTracingPipelineInfo,
69/// #     RayTracingShaderGroup,
70/// # };
71/// # use vk_graph::driver::shader::Shader;
72/// # use vk_graph::Graph;
73/// # fn main() -> Result<(), DriverError> {
74/// # let device = Device::create(DeviceInfo::default())?;
75/// # let info = RayTracingPipelineInfo::default();
76/// # let my_miss_code = [0u8; 1];
77/// # let my_ray_tracing_pipeline = RayTracingPipeline::create(&device, info,
78/// #     [Shader::new_miss(my_miss_code.as_slice())],
79/// #     [RayTracingShaderGroup::new_general(0)],
80/// # )?;
81/// # let mut my_graph = Graph::default();
82/// my_graph.begin_cmd()
83///         .debug_name("my ray tracing command")
84///         .bind_pipeline(&my_ray_tracing_pipeline)
85///         .record_cmd(move |cmd| {
86///             // During this closure we have access to the ray tracing functions!
87///         });
88/// # Ok(()) }
89/// ```
90pub struct RayTracingCommandRef<'a> {
91    cmd: CommandRef<'a>,
92    pipeline: RayTracingPipeline,
93}
94
95impl RayTracingCommandRef<'_> {
96    /// Updates push constants.
97    ///
98    /// Push constants represent a high-speed path to modify constant data in pipelines that is
99    /// expected to outperform memory-backed resource updates.
100    ///
101    /// Push constant values can be updated incrementally, causing shader stages to read the new
102    /// data for push constants modified by this command, while still reading the previous data for
103    /// push constants not modified by this command.
104    ///
105    /// # Device limitations
106    ///
107    /// See [`VkPhysicalDeviceLimits::maxPushConstantsSize`](https://registry.khronos.org/vulkan/specs/latest/man/html/VkPhysicalDeviceLimits.html)
108    /// for the limit of the current device. You may also check [gpuinfo.org] for a listing of
109    /// reported limits on other devices.
110    ///
111    /// # Examples
112    ///
113    /// Basic usage:
114    ///
115    /// ```
116    /// # vk_shader_macros::glsl!(target: vulkan1_2, r#"
117    /// #version 460
118    /// #pragma shader_stage(closest)
119    ///
120    /// layout(push_constant) uniform PushConstants {
121    ///     layout(offset = 0) uint some_val;
122    /// } push_constants;
123    ///
124    /// void main() {
125    ///     uint value = push_constants.some_val;
126    /// }
127    /// # "#);
128    /// ```
129    ///
130    /// ```no_run
131    /// # use ash::vk;
132    /// # use vk_graph::driver::DriverError;
133    /// # use vk_graph::driver::device::{Device, DeviceInfo};
134    /// # use vk_graph::driver::buffer::{Buffer, BufferInfo};
135    /// # use vk_graph::driver::ray_tracing::{
136    /// #     RayTracingPipeline,
137    /// #     RayTracingPipelineInfo,
138    /// #     RayTracingShaderGroup,
139    /// # };
140    /// # use vk_graph::driver::shader::Shader;
141    /// # use vk_graph::Graph;
142    /// # fn main() -> Result<(), DriverError> {
143    /// # let device = Device::create(DeviceInfo::default())?;
144    /// # let shader = [0u8; 1];
145    /// # let info = RayTracingPipelineInfo::default();
146    /// # let my_miss_code = [0u8; 1];
147    /// # let my_ray_tracing_pipeline = RayTracingPipeline::create(&device, info,
148    /// #     [Shader::new_miss(my_miss_code.as_slice())],
149    /// #     [RayTracingShaderGroup::new_general(0)],
150    /// # )?;
151    /// # let rgen_sbt = vk::StridedDeviceAddressRegionKHR {
152    /// #     device_address: 0,
153    /// #     stride: 0,
154    /// #     size: 0,
155    /// # };
156    /// # let hit_sbt = vk::StridedDeviceAddressRegionKHR {
157    /// #     device_address: 0,
158    /// #     stride: 0,
159    /// #     size: 0,
160    /// # };
161    /// # let miss_sbt = vk::StridedDeviceAddressRegionKHR {
162    /// #     device_address: 0,
163    /// #     stride: 0,
164    /// #     size: 0,
165    /// # };
166    /// # let call_sbt = vk::StridedDeviceAddressRegionKHR {
167    /// #     device_address: 0,
168    /// #     stride: 0,
169    /// #     size: 0,
170    /// # };
171    /// # let mut my_graph = Graph::default();
172    /// my_graph.begin_cmd()
173    ///         .debug_name("draw a cornell box")
174    ///         .bind_pipeline(&my_ray_tracing_pipeline)
175    ///         .record_cmd(move |cmd| {
176    ///             cmd.push_constants(0, &[0xcb])
177    ///                    .trace_rays(&rgen_sbt, &hit_sbt, &miss_sbt, &call_sbt, 320, 200, 1);
178    ///         });
179    /// # Ok(()) }
180    /// ```
181    ///
182    /// See [`vkCmdPushConstants`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdPushConstants.html).
183    #[profiling::function]
184    pub fn push_constants(&self, offset: u32, data: &[u8]) -> &Self {
185        self.cmd_push_constants(
186            self.pipeline.inner.layout,
187            &self.pipeline.inner.push_constants,
188            offset,
189            data,
190        );
191
192        self
193    }
194
195    /// Sets the stack size dynamically for a ray tracing pipeline.
196    ///
197    /// The pipeline must have been created with
198    /// [`RayTracingPipelineInfo::dynamic_stack_size`](crate::driver::ray_tracing::RayTracingPipelineInfo::dynamic_stack_size)
199    /// enabled.
200    ///
201    /// See [`vkCmdSetRayTracingPipelineStackSizeKHR`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdSetRayTracingPipelineStackSizeKHR.html).
202    #[profiling::function]
203    pub fn set_stack_size(&self, pipeline_stack_size: u32) -> &Self {
204        let khr_ray_tracing_pipeline = Device::expect_vk_khr_ray_tracing_pipeline(&self.cmd.device);
205
206        #[cfg(feature = "checked")]
207        assert!(
208            self.pipeline.inner.info.dynamic_stack_size,
209            "ray tracing pipeline was not created with dynamic_stack_size enabled"
210        );
211
212        unsafe {
213            /*
214            Checked mode catches missing dynamic_stack_size enablement early. Other Vulkan
215            validation remains the responsibility of the validation layer.
216            */
217            khr_ray_tracing_pipeline
218                .cmd_set_ray_tracing_pipeline_stack_size(self.cmd.handle, pipeline_stack_size);
219        }
220
221        self
222    }
223
224    /*
225    TODO: If the rayTraversalPrimitiveCulling or rayQuery features are enabled, the SkipTrianglesKHR
226    and SkipAABBsKHR ray flags can be specified when tracing a ray. SkipTrianglesKHR and
227    SkipAABBsKHR are mutually exclusive.
228    */
229
230    /// Ray traces using the currently-bound [`RayTracingPipeline`] and the given shader binding
231    /// tables.
232    ///
233    /// Shader binding tables must be constructed according to this [example].
234    ///
235    /// See [`vkCmdTraceRaysKHR`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdTraceRaysKHR.html).
236    ///
237    /// # Examples
238    ///
239    /// Basic usage:
240    ///
241    /// ```no_run
242    /// # use ash::vk;
243    /// # use vk_graph::driver::DriverError;
244    /// # use vk_graph::driver::device::{Device, DeviceInfo};
245    /// # use vk_graph::driver::buffer::{Buffer, BufferInfo};
246    /// # use vk_graph::driver::ray_tracing::{
247    /// #     RayTracingPipeline,
248    /// #     RayTracingPipelineInfo,
249    /// #     RayTracingShaderGroup,
250    /// # };
251    /// # use vk_graph::driver::shader::Shader;
252    /// # use vk_graph::Graph;
253    /// # fn main() -> Result<(), DriverError> {
254    /// # let device = Device::create(DeviceInfo::default())?;
255    /// # let shader = [0u8; 1];
256    /// # let info = RayTracingPipelineInfo::default();
257    /// # let my_miss_code = [0u8; 1];
258    /// # let my_ray_tracing_pipeline = RayTracingPipeline::create(&device, info,
259    /// #     [Shader::new_miss(my_miss_code.as_slice())],
260    /// #     [RayTracingShaderGroup::new_general(0)],
261    /// # )?;
262    /// # let rgen_sbt = vk::StridedDeviceAddressRegionKHR {
263    /// #     device_address: 0,
264    /// #     stride: 0,
265    /// #     size: 0,
266    /// # };
267    /// # let hit_sbt = vk::StridedDeviceAddressRegionKHR {
268    /// #     device_address: 0,
269    /// #     stride: 0,
270    /// #     size: 0,
271    /// # };
272    /// # let miss_sbt = vk::StridedDeviceAddressRegionKHR {
273    /// #     device_address: 0,
274    /// #     stride: 0,
275    /// #     size: 0,
276    /// # };
277    /// # let call_sbt = vk::StridedDeviceAddressRegionKHR {
278    /// #     device_address: 0,
279    /// #     stride: 0,
280    /// #     size: 0,
281    /// # };
282    /// # let mut my_graph = Graph::default();
283    /// my_graph.begin_cmd()
284    ///         .debug_name("draw a cornell box")
285    ///         .bind_pipeline(&my_ray_tracing_pipeline)
286    ///         .record_cmd(move |cmd| {
287    ///             cmd.trace_rays(&rgen_sbt, &hit_sbt, &miss_sbt, &call_sbt, 320, 200, 1);
288    ///         });
289    /// # Ok(()) }
290    /// ```
291    ///
292    /// [example]: https://github.com/attackgoat/vk-graph/blob/master/examples/ray_tracing.rs
293    #[allow(clippy::too_many_arguments)]
294    #[profiling::function]
295    pub fn trace_rays(
296        &self,
297        raygen_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
298        miss_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
299        hit_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
300        callable_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
301        width: u32,
302        height: u32,
303        depth: u32,
304    ) -> &Self {
305        let khr_ray_tracing_pipeline = Device::expect_vk_khr_ray_tracing_pipeline(&self.cmd.device);
306
307        unsafe {
308            khr_ray_tracing_pipeline.cmd_trace_rays(
309                self.cmd.handle,
310                raygen_shader_binding_table,
311                miss_shader_binding_table,
312                hit_shader_binding_table,
313                callable_shader_binding_table,
314                width,
315                height,
316                depth,
317            );
318        }
319
320        self
321    }
322
323    /// Ray traces using the currently-bound [`RayTracingPipeline`] and the given shader binding
324    /// tables.
325    ///
326    /// `indirect_device_address` is a [buffer device address] which is a pointer to a
327    /// [`vk::TraceRaysIndirectCommandKHR`] structure containing the trace ray parameters.
328    ///
329    /// See [`vkCmdTraceRaysIndirectKHR`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdTraceRaysIndirectKHR.html).
330    ///
331    /// [buffer device address]: crate::driver::buffer::Buffer::device_address
332    #[profiling::function]
333    pub fn trace_rays_indirect(
334        &self,
335        raygen_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
336        miss_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
337        hit_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
338        callable_shader_binding_table: &vk::StridedDeviceAddressRegionKHR,
339        indirect_device_address: vk::DeviceAddress,
340    ) -> &Self {
341        let khr_ray_tracing_pipeline = Device::expect_vk_khr_ray_tracing_pipeline(&self.cmd.device);
342
343        unsafe {
344            khr_ray_tracing_pipeline.cmd_trace_rays_indirect(
345                self.cmd.handle,
346                raygen_shader_binding_table,
347                miss_shader_binding_table,
348                hit_shader_binding_table,
349                callable_shader_binding_table,
350                indirect_device_address,
351            )
352        }
353
354        self
355    }
356}
357
358impl<'a> Deref for RayTracingCommandRef<'a> {
359    type Target = CommandRef<'a>;
360
361    fn deref(&self) -> &Self::Target {
362        &self.cmd
363    }
364}
vk_graph/cmd/ray_tracing.rs

vk_graph/cmd/
ray_tracing.rs
