vyre_driver_wgpu/runtime/
indirect.rs

1//! Indirect dispatch path (C-B4).
2//!
3//! Submits `ComputePass::dispatch_workgroups_indirect(buffer, 0)`
4//! on a GPU-resident `[u32; 3]` workgroup-count buffer, so the
5//! upstream kernel can decide the downstream dispatch shape
6//! without a round-trip to host.
7//!
8//! The `core.indirect_dispatch` op (registered in
9//! `vyre-core/src/dialect/core_indirect.rs`) names this operation
10//! in vyre IR; this module is the wgpu-side implementation.
11
12use std::sync::Arc;
13use vyre_driver::BackendError;
14
15use crate::buffer::GpuBufferHandle;
16
17/// Minimum size of an indirect workgroup-count buffer: three
18/// little-endian `u32`s, or 12 bytes.
19pub const INDIRECT_ARGS_BYTES: u64 = 12;
20
21/// Handle + offset identifying where in GPU memory the `[u32; 3]`
22/// workgroup count lives.
23pub struct IndirectArgs {
24    /// GPU buffer containing the workgroup count at the given byte
25    /// offset.
26    pub buffer: Arc<wgpu::Buffer>,
27    /// Byte offset within `buffer`. Must be 4-byte aligned per
28    /// wgpu's contract.
29    pub offset: u64,
30}
31
32impl IndirectArgs {
33    /// Build an `IndirectArgs` from a `GpuBufferHandle` + offset.
34    ///
35    /// # Errors
36    ///
37    /// Returns a `BackendError` when:
38    ///
39    /// * `offset` is not 4-byte aligned (wgpu rejects unaligned
40    ///   indirect dispatches).
41    /// * `offset + INDIRECT_ARGS_BYTES` would exceed the buffer's
42    ///   byte length.
43    /// * The underlying buffer does not carry
44    ///   `wgpu::BufferUsages::INDIRECT`.
45    pub fn from_handle(handle: &GpuBufferHandle, offset: u64) -> Result<Self, BackendError> {
46        if offset & 0b11 != 0 {
47            return Err(BackendError::new(format!(
48                "indirect dispatch offset {offset} is not 4-byte aligned. Fix: align to a u32 boundary."
49            )));
50        }
51        if offset
52            .checked_add(INDIRECT_ARGS_BYTES)
53            .map(|end| end > handle.byte_len())
54            .unwrap_or(true)
55        {
56            return Err(BackendError::new(format!(
57                "indirect dispatch would read past buffer end (offset={offset}, args={INDIRECT_ARGS_BYTES}, buffer byte_len={}). Fix: grow the buffer or lower the offset.",
58                handle.byte_len()
59            )));
60        }
61        if !handle.usage().contains(wgpu::BufferUsages::INDIRECT) {
62            return Err(BackendError::new(
63                "indirect dispatch requires buffer with `wgpu::BufferUsages::INDIRECT`. Fix: allocate the workgroup-count buffer with INDIRECT usage.",
64            ));
65        }
66        Ok(Self {
67            buffer: handle.buffer_arc(),
68            offset,
69        })
70    }
71}
72
73/// Record an indirect dispatch into an existing compute pass.
74///
75/// The caller sets the pipeline + bind group before calling this;
76/// we just submit the `dispatch_workgroups_indirect`.
77pub fn dispatch_indirect<'a>(pass: &mut wgpu::ComputePass<'a>, args: &'a IndirectArgs) {
78    pass.dispatch_workgroups_indirect(&args.buffer, args.offset);
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn args_bytes_is_twelve() {
87        assert_eq!(INDIRECT_ARGS_BYTES, 12);
88    }
89
90    // Note: tests that actually construct IndirectArgs require a
91    // real wgpu::Buffer and hence a GPU. The full dispatch path is
92    // exercised from vyre-wgpu integration tests (`tests/indirect_dispatch.rs`).
93}
vyre_driver_wgpu/runtime/indirect.rs

vyre_driver_wgpu/runtime/
indirect.rs
