oxigdal_gpu/
indirect_dispatch.rs

1//! Indirect compute dispatch support for GPU-driven workloads.
2//!
3//! Indirect dispatch allows the workgroup counts for a compute shader to be
4//! stored in a GPU-side buffer (written by a preceding GPU pass), eliminating
5//! the need for a CPU readback round-trip before issuing the next dispatch.
6//!
7//! # Key types
8//!
9//! - [`DispatchIndirectArgs`] — layout-stable, 12-byte struct that maps
10//!   directly onto the `VkDispatchIndirectCommand` / D3D12 / Metal equivalent.
11//! - [`IndirectDispatchBuffer`] — a GPU buffer holding one or more dispatch
12//!   argument slots, initialised from the CPU and updateable at any time.
13//!
14//! # Helper functions
15//!
16//! - [`workgroup_count_1d`] / [`workgroup_count_2d`] / [`workgroup_count_3d`]
17//!   — compute the ceiling-division workgroup counts for common dispatch
18//!   dimensions without requiring a GPU context.
19//! - [`args_for_elements`] — short-hand for building a 1-D dispatch args struct.
20//! - [`dispatch_indirect_on_pass`] — convenience wrapper that sets the pipeline,
21//!   bind groups, and issues the indirect dispatch in one call.
22
23use std::sync::Arc;
24use wgpu::{Buffer, BufferAddress, BufferDescriptor, BufferUsages, ComputePipeline};
25
26use crate::context::GpuContext;
27use crate::error::{GpuError, GpuResult};
28
29// ─────────────────────────────────────────────────────────────────────────────
30// DispatchIndirectArgs
31// ─────────────────────────────────────────────────────────────────────────────
32
33/// Layout-stable representation of workgroup dispatch dimensions stored in GPU
34/// memory.
35///
36/// The struct is `repr(C)` so that its in-memory byte layout is well-defined
37/// and can be written directly to a GPU buffer via
38/// [`wgpu::Queue::write_buffer`].  It corresponds to the 12-byte
39/// `VkDispatchIndirectCommand` structure (and its equivalents in D3D12 and
40/// Metal).
41///
42/// # Layout
43///
44/// | Bytes | Field | Meaning                         |
45/// |-------|-------|---------------------------------|
46/// | 0–3   | `x`   | Workgroup count along X axis    |
47/// | 4–7   | `y`   | Workgroup count along Y axis    |
48/// | 8–11  | `z`   | Workgroup count along Z axis    |
49///
50/// All values are stored in little-endian byte order when serialised via
51/// [`DispatchIndirectArgs::as_bytes`].
52#[repr(C)]
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct DispatchIndirectArgs {
55    /// Workgroup count along the X dispatch dimension.
56    pub x: u32,
57    /// Workgroup count along the Y dispatch dimension.
58    pub y: u32,
59    /// Workgroup count along the Z dispatch dimension.
60    pub z: u32,
61}
62
63impl DispatchIndirectArgs {
64    /// Create new dispatch arguments with explicit x / y / z workgroup counts.
65    #[inline]
66    pub fn new(x: u32, y: u32, z: u32) -> Self {
67        Self { x, y, z }
68    }
69
70    /// Create 1-D dispatch arguments for `total` elements with a workgroup
71    /// size of `workgroup_size`.
72    ///
73    /// The Y and Z dimensions are set to 1. If `workgroup_size` is 0 the X
74    /// count is also 0.
75    #[inline]
76    pub fn for_1d(total: u32, workgroup_size: u32) -> Self {
77        Self {
78            x: workgroup_count_1d(total, workgroup_size),
79            y: 1,
80            z: 1,
81        }
82    }
83
84    /// Serialise the struct into a 12-byte little-endian array suitable for
85    /// writing to a GPU buffer.
86    #[inline]
87    pub fn as_bytes(&self) -> [u8; 12] {
88        let mut buf = [0u8; 12];
89        buf[0..4].copy_from_slice(&self.x.to_le_bytes());
90        buf[4..8].copy_from_slice(&self.y.to_le_bytes());
91        buf[8..12].copy_from_slice(&self.z.to_le_bytes());
92        buf
93    }
94}
95
96// ─────────────────────────────────────────────────────────────────────────────
97// IndirectDispatchBuffer
98// ─────────────────────────────────────────────────────────────────────────────
99
100/// A GPU buffer that holds one or more [`DispatchIndirectArgs`] slots for use
101/// with `wgpu::ComputePass::dispatch_workgroups_indirect`.
102///
103/// Each slot occupies exactly 12 bytes (the size of one
104/// [`DispatchIndirectArgs`]).  The buffer is created with
105/// `INDIRECT | COPY_DST | STORAGE` usages so it can both receive CPU-side
106/// writes and be written by a preceding compute shader pass.
107///
108/// # Multi-slot layout
109///
110/// ```text
111/// Offset 0       12     24     36 …
112///         ├──────┼──────┼──────┤
113///         │ [0]  │ [1]  │ [2]  │  …
114///         └──────┴──────┴──────┘
115///           12 B   12 B   12 B
116/// ```
117///
118/// The active slot for a dispatch is selected via [`IndirectDispatchBuffer::offset`]
119/// (defaults to slot 0) and can be overridden by calling
120/// [`IndirectDispatchBuffer::update_at`] to write to a specific slot and then
121/// constructing a view with the desired byte offset — or by creating separate
122/// buffers for each independently-driven dispatch.
123pub struct IndirectDispatchBuffer {
124    buffer: Arc<Buffer>,
125    /// Byte offset into `buffer` at which slot 0 begins.
126    ///
127    /// Always 0 for buffers created via the public constructors; exposed as a
128    /// field to allow future sub-buffer views.
129    offset: BufferAddress,
130    /// Total number of 12-byte dispatch slots in the buffer.
131    capacity: u32,
132}
133
134impl IndirectDispatchBuffer {
135    /// Create a single-slot indirect dispatch buffer initialised with `args`.
136    ///
137    /// The buffer is 12 bytes and supports `INDIRECT | COPY_DST | STORAGE`
138    /// usages.
139    ///
140    /// # Errors
141    ///
142    /// Returns [`GpuError::InvalidKernelParams`] if the device rejects the
143    /// buffer descriptor (should not happen in practice).
144    pub fn new(ctx: &GpuContext, args: DispatchIndirectArgs) -> GpuResult<Self> {
145        let buffer = ctx.device().create_buffer(&BufferDescriptor {
146            label: Some("indirect_dispatch"),
147            size: 12,
148            usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST | BufferUsages::STORAGE,
149            mapped_at_creation: false,
150        });
151        let buf = Arc::new(buffer);
152        ctx.queue().write_buffer(&buf, 0, &args.as_bytes());
153        Ok(Self {
154            buffer: buf,
155            offset: 0,
156            capacity: 1,
157        })
158    }
159
160    /// Create a multi-slot indirect dispatch buffer with `max_dispatches`
161    /// slots, all initialised to zero.
162    ///
163    /// The total buffer size is `12 * max_dispatches` bytes.  Each slot can be
164    /// written individually via [`IndirectDispatchBuffer::update_at`].
165    ///
166    /// # Errors
167    ///
168    /// Returns [`GpuError::InvalidKernelParams`] if `max_dispatches` is 0.
169    pub fn new_with_capacity(ctx: &GpuContext, max_dispatches: u32) -> GpuResult<Self> {
170        if max_dispatches == 0 {
171            return Err(GpuError::InvalidKernelParams {
172                reason: "max_dispatches must be > 0".to_string(),
173            });
174        }
175        let size = 12u64 * max_dispatches as u64;
176        let buffer = ctx.device().create_buffer(&BufferDescriptor {
177            label: Some("indirect_dispatch_multi"),
178            size,
179            usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST | BufferUsages::STORAGE,
180            mapped_at_creation: false,
181        });
182        Ok(Self {
183            buffer: Arc::new(buffer),
184            offset: 0,
185            capacity: max_dispatches,
186        })
187    }
188
189    /// Overwrite slot 0 with new dispatch arguments.
190    ///
191    /// This is a CPU-side write using `queue.write_buffer`; the data is
192    /// uploaded before the next `queue.submit`.
193    pub fn update(&self, ctx: &GpuContext, args: DispatchIndirectArgs) {
194        ctx.queue()
195            .write_buffer(&self.buffer, self.offset, &args.as_bytes());
196    }
197
198    /// Overwrite a specific slot (0-indexed) with new dispatch arguments.
199    ///
200    /// # Errors
201    ///
202    /// Returns [`GpuError::InvalidKernelParams`] if `slot_idx >= capacity`.
203    pub fn update_at(
204        &self,
205        ctx: &GpuContext,
206        slot_idx: u32,
207        args: DispatchIndirectArgs,
208    ) -> GpuResult<()> {
209        if slot_idx >= self.capacity {
210            return Err(GpuError::InvalidKernelParams {
211                reason: format!(
212                    "slot_idx {} is out of range for capacity {}",
213                    slot_idx, self.capacity
214                ),
215            });
216        }
217        let slot_offset = self.offset + 12 * slot_idx as u64;
218        ctx.queue()
219            .write_buffer(&self.buffer, slot_offset, &args.as_bytes());
220        Ok(())
221    }
222
223    /// Return a reference to the underlying [`wgpu::Buffer`].
224    #[inline]
225    pub fn buffer(&self) -> &Buffer {
226        &self.buffer
227    }
228
229    /// Return the byte offset of slot 0 within the buffer.
230    ///
231    /// This is the value passed as the second argument to
232    /// `ComputePass::dispatch_workgroups_indirect`.
233    #[inline]
234    pub fn offset(&self) -> BufferAddress {
235        self.offset
236    }
237
238    /// Return the number of 12-byte dispatch slots in the buffer.
239    #[inline]
240    pub fn capacity(&self) -> u32 {
241        self.capacity
242    }
243}
244
245// ─────────────────────────────────────────────────────────────────────────────
246// dispatch_indirect_on_pass
247// ─────────────────────────────────────────────────────────────────────────────
248
249/// Set the pipeline, bind groups, and issue an indirect dispatch in one call.
250///
251/// This is a thin convenience wrapper around the raw `wgpu::ComputePass` API:
252///
253/// ```text
254/// pass.set_pipeline(pipeline);
255/// for (i, bg) in bind_groups.iter().enumerate() {
256///     pass.set_bind_group(i as u32, *bg, &[]);
257/// }
258/// pass.dispatch_workgroups_indirect(indirect.buffer(), indirect.offset());
259/// ```
260///
261/// # Arguments
262///
263/// * `pass`        — an in-progress `ComputePass`.
264/// * `pipeline`    — the compute pipeline to bind.
265/// * `bind_groups` — slice of bind groups; index 0 → slot 0, etc.
266/// * `indirect`    — the buffer supplying the workgroup counts.
267///
268/// # Panics
269///
270/// Panics with a wgpu validation error (not a Rust panic) if the pipeline,
271/// bind groups, or indirect buffer are invalid for the current device.
272pub fn dispatch_indirect_on_pass<'a>(
273    pass: &mut wgpu::ComputePass<'a>,
274    pipeline: &'a ComputePipeline,
275    bind_groups: &[&'a wgpu::BindGroup],
276    indirect: &'a IndirectDispatchBuffer,
277) {
278    pass.set_pipeline(pipeline);
279    for (i, bg) in bind_groups.iter().enumerate() {
280        pass.set_bind_group(i as u32, *bg, &[]);
281    }
282    pass.dispatch_workgroups_indirect(indirect.buffer(), indirect.offset());
283}
284
285// ─────────────────────────────────────────────────────────────────────────────
286// Pure arithmetic helpers — no GPU context required
287// ─────────────────────────────────────────────────────────────────────────────
288
289/// Compute the ceiling-division workgroup count for a 1-D dispatch.
290///
291/// Returns `ceil(total_elements / workgroup_size)`.  If `workgroup_size` is 0
292/// the result is 0 (to avoid a division-by-zero panic).
293///
294/// # Examples
295///
296/// ```
297/// use oxigdal_gpu::workgroup_count_1d;
298/// assert_eq!(workgroup_count_1d(64, 64), 1);
299/// assert_eq!(workgroup_count_1d(65, 64), 2);
300/// assert_eq!(workgroup_count_1d(0,  64), 0);
301/// ```
302#[inline]
303pub fn workgroup_count_1d(total_elements: u32, workgroup_size: u32) -> u32 {
304    if workgroup_size == 0 {
305        return 0;
306    }
307    total_elements.saturating_add(workgroup_size - 1) / workgroup_size
308}
309
310/// Compute the ceiling-division workgroup counts for a 2-D dispatch.
311///
312/// Returns `(ceil(width / wg_x), ceil(height / wg_y))`.
313///
314/// # Examples
315///
316/// ```
317/// use oxigdal_gpu::workgroup_count_2d;
318/// assert_eq!(workgroup_count_2d(128, 64, 16, 8), (8, 8));
319/// ```
320#[inline]
321pub fn workgroup_count_2d(width: u32, height: u32, wg_x: u32, wg_y: u32) -> (u32, u32) {
322    (
323        workgroup_count_1d(width, wg_x),
324        workgroup_count_1d(height, wg_y),
325    )
326}
327
328/// Compute the ceiling-division workgroup counts for a 3-D dispatch.
329///
330/// Returns `(ceil(d0 / wg0), ceil(d1 / wg1), ceil(d2 / wg2))`.
331///
332/// # Examples
333///
334/// ```
335/// use oxigdal_gpu::workgroup_count_3d;
336/// assert_eq!(workgroup_count_3d((10, 10, 10), (4, 4, 4)), (3, 3, 3));
337/// ```
338#[inline]
339pub fn workgroup_count_3d(dims: (u32, u32, u32), wg_dims: (u32, u32, u32)) -> (u32, u32, u32) {
340    (
341        workgroup_count_1d(dims.0, wg_dims.0),
342        workgroup_count_1d(dims.1, wg_dims.1),
343        workgroup_count_1d(dims.2, wg_dims.2),
344    )
345}
346
347/// Build [`DispatchIndirectArgs`] for a 1-D dispatch over `total` elements
348/// with the given `workgroup_size`.
349///
350/// Y and Z are set to 1.
351///
352/// # Examples
353///
354/// ```
355/// use oxigdal_gpu::{args_for_elements, DispatchIndirectArgs};
356/// let args = args_for_elements(128, 64);
357/// assert_eq!(args, DispatchIndirectArgs { x: 2, y: 1, z: 1 });
358/// ```
359#[inline]
360pub fn args_for_elements(total: u32, workgroup_size: u32) -> DispatchIndirectArgs {
361    DispatchIndirectArgs {
362        x: workgroup_count_1d(total, workgroup_size),
363        y: 1,
364        z: 1,
365    }
366}
oxigdal_gpu/indirect_dispatch.rs

oxigdal_gpu/
indirect_dispatch.rs
