vk_graph/cmd/compute.rs
1use {
2 super::{cmd_ref::CommandRef, pipeline::PipelineCommand},
3 crate::{driver::compute::ComputePipeline, node::AnyBufferNode},
4 ash::vk,
5 std::ops::Deref,
6};
7
8impl PipelineCommand<'_, ComputePipeline> {
9 /// Begin recording compute pipeline work for this graph command.
10 pub fn record_cmd(mut self, func: impl FnOnce(ComputeCommandRef<'_>) + Send + 'static) -> Self {
11 self.record_cmd_mut(func);
12 self
13 }
14
15 /// Mutable-borrow form of [`Self::record_cmd`].
16 pub fn record_cmd_mut(&mut self, func: impl FnOnce(ComputeCommandRef<'_>) + Send + 'static) {
17 let pipeline = self
18 .cmd
19 .cmd()
20 .expect_last_pipeline()
21 .expect_compute()
22 .clone();
23
24 self.cmd.push_exec(move |cmd| {
25 func(ComputeCommandRef { cmd, pipeline });
26 });
27 }
28
29 pub(crate) fn record_stream_mut(
30 &mut self,
31 func: impl for<'r> Fn(ComputeCommandRef<'r>) + Send + Sync + 'static,
32 ) {
33 let pipeline = self
34 .cmd
35 .cmd()
36 .expect_last_pipeline()
37 .expect_compute()
38 .clone();
39
40 self.cmd.push_reusable_exec(move |cmd| {
41 func(ComputeCommandRef {
42 cmd,
43 pipeline: pipeline.clone(),
44 });
45 });
46 }
47}
48
49/// Recording interface for compute commands.
50///
51/// This structure provides a strongly-typed set of methods which allow compute shader code to be
52/// executed. An instance is provided to the closure argument of
53/// [`PipelineCommand::record_cmd`] which may be accessed by binding a [`ComputePipeline`] to a
54/// command.
55///
56/// # Examples
57///
58/// Basic usage:
59///
60/// ```no_run
61/// # use ash::vk;
62/// # use vk_graph::driver::DriverError;
63/// # use vk_graph::driver::device::{Device, DeviceInfo};
64/// # use vk_graph::driver::compute::{ComputePipeline, ComputePipelineInfo};
65/// # use vk_graph::driver::shader::{Shader};
66/// # use vk_graph::Graph;
67/// # fn main() -> Result<(), DriverError> {
68/// # let device = Device::create(DeviceInfo::default())?;
69/// # let info = ComputePipelineInfo::default();
70/// # let shader = Shader::new_compute([0u8; 1].as_slice());
71/// # let my_compute_pipeline = ComputePipeline::create(&device, info, shader)?;
72/// # let mut my_graph = Graph::default();
73/// my_graph
74/// .begin_cmd()
75/// .bind_pipeline(&my_compute_pipeline)
76/// .record_cmd(move |cmd| {
77/// // During this closure we have access to the compute functions!
78/// cmd.dispatch(64, 1, 1);
79/// });
80/// # Ok(()) }
81/// ```
82pub struct ComputeCommandRef<'a> {
83 cmd: CommandRef<'a>,
84 pipeline: ComputePipeline,
85}
86
87impl ComputeCommandRef<'_> {
88 /// [`Self::dispatch`] compute work items.
89 ///
90 /// When the command is executed, a global workgroup consisting of
91 /// `group_count_x × group_count_y × group_count_z` local workgroups is assembled.
92 ///
93 /// # Examples
94 ///
95 /// Basic usage:
96 ///
97 /// ```
98 /// # vk_shader_macros::glsl!(r#"
99 /// #version 450
100 /// #pragma shader_stage(compute)
101 ///
102 /// layout(set = 0, binding = 0, std430) restrict writeonly buffer MyBuffer {
103 /// uint my_buf[];
104 /// };
105 ///
106 /// void main() {
107 /// my_buf[0] = 1;
108 /// }
109 /// # "#);
110 /// ```
111 ///
112 /// ```no_run
113 /// # use ash::vk;
114 /// # use vk_sync::AccessType;
115 /// # use vk_graph::driver::DriverError;
116 /// # use vk_graph::driver::device::{Device, DeviceInfo};
117 /// # use vk_graph::driver::buffer::{Buffer, BufferInfo};
118 /// # use vk_graph::driver::compute::{ComputePipeline, ComputePipelineInfo};
119 /// # use vk_graph::driver::shader::{Shader};
120 /// # use vk_graph::Graph;
121 /// # fn main() -> Result<(), DriverError> {
122 /// # let device = Device::create(DeviceInfo::default())?;
123 /// # let buf_info = BufferInfo::device_mem(8, vk::BufferUsageFlags::STORAGE_BUFFER);
124 /// # let my_buf = Buffer::create(&device, buf_info)?;
125 /// # let info = ComputePipelineInfo::default();
126 /// # let shader = Shader::new_compute([0u8; 1].as_slice());
127 /// # let my_compute_pipeline = ComputePipeline::create(&device, info, shader)?;
128 /// # let mut my_graph = Graph::default();
129 /// # let my_buf_node = my_graph.bind_resource(my_buf);
130 /// my_graph
131 /// .begin_cmd()
132 /// .debug_name("fill my_buf_node with data")
133 /// .bind_pipeline(&my_compute_pipeline)
134 /// .shader_resource_access(0, my_buf_node, AccessType::ComputeShaderWrite)
135 /// .record_cmd(move |cmd| {
136 /// cmd.dispatch(128, 64, 32);
137 /// });
138 /// # Ok(()) }
139 /// ```
140 ///
141 /// See [`vkCmdDispatch`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdDispatch.html).
142 #[profiling::function]
143 pub fn dispatch(&self, group_count_x: u32, group_count_y: u32, group_count_z: u32) -> &Self {
144 unsafe {
145 self.cmd.device.cmd_dispatch(
146 self.cmd.handle,
147 group_count_x,
148 group_count_y,
149 group_count_z,
150 );
151 }
152
153 self
154 }
155
156 /// [`Self::dispatch_base`] compute work items with non-zero base values for the workgroup IDs.
157 ///
158 /// When the command is executed, a global workgroup consisting of
159 /// `group_count_x × group_count_y × group_count_z` local workgroups is assembled, with
160 /// WorkgroupId values ranging from `[base_group*, base_group* + group_count*)` in each
161 /// component.
162 ///
163 /// [`Self::dispatch`] is equivalent to
164 /// `dispatch_base(0, 0, 0, group_count_x, group_count_y, group_count_z)`.
165 ///
166 /// See [`vkCmdDispatchBase`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdDispatchBase.html).
167 #[profiling::function]
168 pub fn dispatch_base(
169 &self,
170 base_group_x: u32,
171 base_group_y: u32,
172 base_group_z: u32,
173 group_count_x: u32,
174 group_count_y: u32,
175 group_count_z: u32,
176 ) -> &Self {
177 unsafe {
178 self.cmd.device.cmd_dispatch_base(
179 self.cmd.handle,
180 base_group_x,
181 base_group_y,
182 base_group_z,
183 group_count_x,
184 group_count_y,
185 group_count_z,
186 );
187 }
188
189 self
190 }
191
192 /// Dispatch compute work items with indirect parameters.
193 ///
194 /// `dispatch_indirect` behaves similarly to [`Self::dispatch`] except that the parameters
195 /// are read by the device from `args_buf` during execution. The parameters of the dispatch are
196 /// encoded in a [`vk::DispatchIndirectCommand`] structure taken from `args_buf` starting at
197 /// `args_offset`.
198 ///
199 /// # Examples
200 ///
201 /// Basic usage:
202 ///
203 /// ```no_run
204 /// # use ash::vk;
205 /// # use bytemuck::{bytes_of, Pod, Zeroable};
206 /// # use vk_sync::AccessType;
207 /// # use vk_graph::driver::DriverError;
208 /// # use vk_graph::driver::device::{Device, DeviceInfo};
209 /// # use vk_graph::driver::buffer::{Buffer, BufferInfo};
210 /// # use vk_graph::driver::compute::{ComputePipeline, ComputePipelineInfo};
211 /// # use vk_graph::driver::shader::{Shader};
212 /// # use vk_graph::Graph;
213 /// # fn main() -> Result<(), DriverError> {
214 /// # let device = Device::create(DeviceInfo::default())?;
215 /// # let buf_info = BufferInfo::device_mem(8, vk::BufferUsageFlags::STORAGE_BUFFER);
216 /// # let my_buf = Buffer::create(&device, buf_info)?;
217 /// # let info = ComputePipelineInfo::default();
218 /// # let shader = Shader::new_compute([0u8; 1].as_slice());
219 /// # let my_compute_pipeline = ComputePipeline::create(&device, info, shader)?;
220 /// # let mut my_graph = Graph::default();
221 /// # let my_buf_node = my_graph.bind_resource(my_buf);
222 /// # #[repr(C)]
223 /// # #[derive(Clone, Copy, Pod, Zeroable)]
224 /// # struct DispatchIndirectCommand { x: u32, y: u32, z: u32, }
225 /// let args = DispatchIndirectCommand {
226 /// x: 1,
227 /// y: 2,
228 /// z: 3,
229 /// };
230 /// let data = bytes_of(&args);
231 /// let usage = vk::BufferUsageFlags::INDIRECT_BUFFER | vk::BufferUsageFlags::STORAGE_BUFFER;
232 /// let args_buf = Buffer::create_from_slice(&device, usage, data)?;
233 /// let args_buf = my_graph.bind_resource(args_buf);
234 ///
235 /// my_graph
236 /// .begin_cmd()
237 /// .debug_name("fill my_buf_node with data")
238 /// .bind_pipeline(&my_compute_pipeline)
239 /// .resource_access(args_buf, AccessType::IndirectBuffer)
240 /// .shader_resource_access(0, my_buf_node, AccessType::ComputeShaderWrite)
241 /// .record_cmd(move |cmd| {
242 /// cmd.dispatch_indirect(args_buf, 0);
243 /// });
244 /// # Ok(()) }
245 /// ```
246 ///
247 /// See [`vkCmdDispatchIndirect`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdDispatchIndirect.html).
248 #[profiling::function]
249 pub fn dispatch_indirect(
250 &self,
251 args_buf: impl Into<AnyBufferNode>,
252 args_offset: vk::DeviceSize,
253 ) -> &Self {
254 let args_buf = args_buf.into();
255 let args_buf = self.resource(args_buf);
256
257 unsafe {
258 self.cmd
259 .device
260 .cmd_dispatch_indirect(self.cmd.handle, args_buf.handle, args_offset);
261 }
262
263 self
264 }
265
266 /// Updates push constants.
267 ///
268 /// Push constants represent a high-speed path to modify constant data in pipelines that is
269 /// expected to outperform memory-backed resource updates.
270 ///
271 /// Push constant values can be updated incrementally, causing shader stages to read the new
272 /// data for push constants modified by this command, while still reading the previous data for
273 /// push constants not modified by this command.
274 ///
275 /// # Device limitations
276 ///
277 /// See [`VkPhysicalDeviceLimits::maxPushConstantsSize`](https://registry.khronos.org/vulkan/specs/latest/man/html/VkPhysicalDeviceLimits.html)
278 /// for the limit of the current device. You may also check [gpuinfo.org] for a listing of
279 /// reported limits on other devices.
280 ///
281 /// # Examples
282 ///
283 /// Basic usage:
284 ///
285 /// ```
286 /// # vk_shader_macros::glsl!(r#"
287 /// #version 450
288 /// #pragma shader_stage(compute)
289 ///
290 /// layout(push_constant) uniform PushConstants {
291 /// layout(offset = 0) uint the_answer;
292 /// } push_constants;
293 ///
294 /// void main()
295 /// {
296 /// uint value = push_constants.the_answer;
297 /// }
298 /// # "#);
299 /// ```
300 ///
301 /// ```no_run
302 /// # use ash::vk;
303 /// # use vk_graph::driver::DriverError;
304 /// # use vk_graph::driver::device::{Device, DeviceInfo};
305 /// # use vk_graph::driver::buffer::{Buffer, BufferInfo};
306 /// # use vk_graph::driver::compute::{ComputePipeline, ComputePipelineInfo};
307 /// # use vk_graph::driver::shader::{Shader};
308 /// # use vk_graph::Graph;
309 /// # fn main() -> Result<(), DriverError> {
310 /// # let device = Device::create(DeviceInfo::default())?;
311 /// # let info = ComputePipelineInfo::default();
312 /// # let shader = Shader::new_compute([0u8; 1].as_slice());
313 /// # let my_compute_pipeline = ComputePipeline::create(&device, info, shader)?;
314 /// # let mut my_graph = Graph::default();
315 /// my_graph
316 /// .begin_cmd()
317 /// .debug_name("compute the ultimate question")
318 /// .bind_pipeline(&my_compute_pipeline)
319 /// .record_cmd(move |cmd| {
320 /// cmd
321 /// .push_constants(0, &[42])
322 /// .dispatch(1, 1, 1);
323 /// });
324 /// # Ok(()) }
325 /// ```
326 ///
327 /// See [`vkCmdPushConstants`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkCmdPushConstants.html).
328 #[profiling::function]
329 pub fn push_constants(&self, offset: u32, data: &[u8]) -> &Self {
330 self.cmd_push_constants(
331 self.pipeline.inner.layout,
332 self.pipeline.inner.push_constants.as_slice(),
333 offset,
334 data,
335 );
336
337 self
338 }
339}
340
341impl<'a> Deref for ComputeCommandRef<'a> {
342 type Target = CommandRef<'a>;
343
344 fn deref(&self) -> &Self::Target {
345 &self.cmd
346 }
347}