vyre_driver_wgpu/runtime/
prerecorded.rs

1//! Pre-recorded persistent dispatch command buffers.
2
3use std::sync::{Arc, Mutex};
4
5use smallvec::SmallVec;
6use vyre_driver::BackendError;
7
8use crate::allocation::padded_wgpu_u64;
9use crate::buffer::GpuBufferHandle;
10use crate::pipeline::binding::{clear_outputs_for_bound, validate_handle};
11use crate::pipeline::{BufferBindingInfo, WgpuPipeline};
12
13/// GPU work recorded ahead of submission for encoder-free dispatch handoff.
14///
15/// `wgpu::CommandBuffer` is single-submit. This type prevents the raw wgpu
16/// panic by consuming the stored command buffer on the first replay and
17/// returning a structured error on repeated replay attempts.
18pub struct PrerecordedDispatch {
19    /// Pre-recorded command buffer.
20    pub cb: Mutex<Option<wgpu::CommandBuffer>>,
21    /// Bind groups captured by the command buffer.
22    pub bind_groups: Vec<Arc<wgpu::BindGroup>>,
23    /// Buffer handles kept alive for the lifetime of the recorded commands.
24    pub handles: Vec<GpuBufferHandle>,
25    /// Output handles recorded for terminal readback by tests and callers.
26    pub output_handles: Vec<GpuBufferHandle>,
27    /// Device used to record this dispatch.
28    pub device: wgpu::Device,
29    /// Queue paired with `device`.
30    pub queue: wgpu::Queue,
31}
32
33impl PrerecordedDispatch {
34    /// Submit the pre-recorded command buffer to `queue`.
35    ///
36    /// # Errors
37    ///
38    /// Returns a backend error when this command buffer was already submitted.
39    pub fn replay(&self, queue: &wgpu::Queue) -> Result<wgpu::SubmissionIndex, BackendError> {
40        let command_buffer = self
41            .cb
42            .lock()
43            .map_err(|source| {
44                BackendError::new(format!(
45                    "pre-recorded dispatch mutex poisoned: {source}. Fix: drop this dispatch and record a fresh command buffer."
46                ))
47            })?
48            .take()
49            .ok_or_else(|| {
50                BackendError::new(
51                    "pre-recorded wgpu command buffer was already submitted. Fix: record a new PrerecordedDispatch for each replay slot; wgpu command buffers are single-submit.",
52                )
53            })?;
54        Ok(queue.submit(std::iter::once(command_buffer)))
55    }
56
57    /// Read one recorded output buffer into a byte vector.
58    ///
59    /// # Errors
60    ///
61    /// Returns a backend error when the output index is invalid or mapping
62    /// fails.
63    pub fn read_output(&self, index: usize) -> Result<Vec<u8>, BackendError> {
64        let output = self.output_handles.get(index).ok_or_else(|| {
65            BackendError::new(format!(
66                "pre-recorded output index {index} is out of bounds for {} outputs. Fix: request an output produced by this dispatch.",
67                self.output_handles.len()
68            ))
69        })?;
70        let byte_capacity = usize::try_from(output.byte_len()).map_err(|error| {
71            BackendError::new(format!(
72                "pre-recorded output byte length {} does not fit usize on this host: {error}. Fix: shard the GPU output before readback.",
73                output.byte_len()
74            ))
75        })?;
76        let mut bytes = Vec::new();
77        vyre_driver::allocation::try_reserve_vec_to_capacity(&mut bytes, byte_capacity).map_err(
78            |source| {
79                BackendError::new(format!(
80                    "pre-recorded output readback could not reserve {byte_capacity} byte(s): {source}. Fix: shard the GPU output before readback."
81                ))
82            },
83        )?;
84        output.readback(&self.device, &self.queue, &mut bytes)?;
85        Ok(bytes)
86    }
87
88    /// Read one recorded output buffer into caller-owned storage.
89    ///
90    /// Clears `out`, then reuses its allocation.
91    ///
92    /// # Errors
93    ///
94    /// Returns a backend error when the output index is invalid or mapping
95    /// fails.
96    pub fn read_output_into(&self, index: usize, out: &mut Vec<u8>) -> Result<(), BackendError> {
97        let output = self.output_handles.get(index).ok_or_else(|| {
98            BackendError::new(format!(
99                "pre-recorded output index {index} is out of bounds for {} outputs. Fix: request an output produced by this dispatch.",
100                self.output_handles.len()
101            ))
102        })?;
103        output.readback(&self.device, &self.queue, out)
104    }
105}
106
107impl WgpuPipeline {
108    /// Record a persistent dispatch once so later submission bypasses encoder
109    /// construction, output clears, bind-group lookup, and compute-pass setup.
110    ///
111    /// # Errors
112    ///
113    /// Returns a backend error when the handles do not match the compiled
114    /// program's binding contract or command recording fails.
115    pub fn prerecord_persistent_dispatch(
116        &self,
117        inputs: &[GpuBufferHandle],
118        outputs: &[GpuBufferHandle],
119        params: Option<&GpuBufferHandle>,
120        workgroups: [u32; 3],
121    ) -> Result<PrerecordedDispatch, BackendError> {
122        let (device, queue) = &*self.device_queue;
123        let bound = bind_handles(&self.buffer_bindings, inputs, outputs, params)?;
124        let mut grouped_bound: Vec<SmallVec<[(&BufferBindingInfo, &GpuBufferHandle); 16]>> =
125            Vec::new();
126        vyre_driver::allocation::try_reserve_vec_to_capacity(
127            &mut grouped_bound,
128            self.bind_group_layouts.len(),
129        )
130        .map_err(|source| {
131            BackendError::new(format!(
132                "pre-recorded bind-group staging could not reserve {} group slot(s): {source}. Fix: split bind-group resources before recording.",
133                self.bind_group_layouts.len()
134            ))
135        })?;
136        grouped_bound.resize_with(self.bind_group_layouts.len(), SmallVec::new);
137        for (info, handle) in &bound {
138            let group = usize::try_from(info.group).map_err(|source| {
139                BackendError::new(format!(
140                    "pre-recorded bind group {} cannot fit usize: {source}. Fix: keep group indices representable on this host.",
141                    info.group
142                ))
143            })?;
144            let Some(slot) = grouped_bound.get_mut(group) else {
145                return Err(BackendError::new(format!(
146                    "pre-recorded binding {} (`{}`) targets group {}, but the pipeline only has {} bind-group layouts. Fix: keep reflection metadata synchronized with bind-group layouts.",
147                    info.binding,
148                    info.name,
149                    info.group,
150                    self.bind_group_layouts.len()
151                )));
152            };
153            slot.push((*info, *handle));
154        }
155        let mut bind_groups = Vec::new();
156        vyre_driver::allocation::try_reserve_vec_to_capacity(
157            &mut bind_groups,
158            self.bind_group_layouts.len(),
159        )
160        .map_err(|source| {
161            BackendError::new(format!(
162                "pre-recorded bind-group cache result could not reserve {} group slot(s): {source}. Fix: split bind-group resources before recording.",
163                self.bind_group_layouts.len()
164            ))
165        })?;
166        for (group_index, layout) in self.bind_group_layouts.iter().enumerate() {
167            let group_bound = &grouped_bound[group_index];
168            let handle_id_capacity = group_bound.len().checked_mul(2).ok_or_else(|| {
169                BackendError::new(
170                    "pre-recorded bind group handle-id count overflowed usize. Fix: split bind-group resources before recording.",
171                )
172            })?;
173            let mut handle_ids: SmallVec<[u64; 16]> = SmallVec::new();
174            vyre_foundation::allocation::try_reserve_smallvec_to_capacity(
175                &mut handle_ids,
176                handle_id_capacity,
177            )
178            .map_err(|source| {
179                BackendError::new(format!(
180                    "pre-recorded bind-group handle-id cache key could not reserve {handle_id_capacity} word slot(s): {source}. Fix: split bind-group resources before recording."
181                ))
182            })?;
183            let mut checked_bound: SmallVec<[(&BufferBindingInfo, &GpuBufferHandle, u64); 16]> =
184                SmallVec::new();
185            vyre_foundation::allocation::try_reserve_smallvec_to_capacity(
186                &mut checked_bound,
187                group_bound.len(),
188            )
189            .map_err(|source| {
190                BackendError::new(format!(
191                    "pre-recorded bind-group checked binding staging could not reserve {} binding slot(s): {source}. Fix: split bind-group resources before recording.",
192                    group_bound.len()
193                ))
194            })?;
195            for (_, handle) in group_bound {
196                handle_ids.push(handle.allocation_identity());
197                let bind_size = padded_wgpu_u64(
198                    handle.byte_len(),
199                    "pre-recorded bind-group cache key byte length",
200                    "split the pre-recorded dispatch buffer",
201                )?;
202                handle_ids.push(bind_size);
203            }
204            for (info, handle) in group_bound {
205                checked_bound.push((
206                    info,
207                    handle,
208                    padded_wgpu_u64(
209                        handle.byte_len(),
210                        "pre-recorded bind-group binding size",
211                        "split the pre-recorded dispatch buffer",
212                    )?,
213                ));
214            }
215            let layout_id = Arc::as_ptr(layout).addr();
216            let bg = self
217                .bind_group_cache
218                .get_or_create_by_ids(layout_id, handle_ids, || {
219                    let mut entries = SmallVec::<[wgpu::BindGroupEntry<'_>; 16]>::with_capacity(
220                        group_bound.len(),
221                    );
222                    entries.extend(checked_bound.iter().map(|(info, handle, bind_size)| {
223                        wgpu::BindGroupEntry {
224                            binding: info.binding,
225                            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
226                                buffer: handle.buffer(),
227                                offset: 0,
228                                size: wgpu::BufferSize::new(*bind_size),
229                            }),
230                        }
231                    }));
232                    device.create_bind_group(&wgpu::BindGroupDescriptor {
233                        label: Some("vyre pre-recorded persistent bind group"),
234                        layout,
235                        entries: &entries,
236                    })
237                });
238            bind_groups.push(bg);
239        }
240
241        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
242            label: Some("vyre pre-recorded persistent dispatch"),
243        });
244        clear_outputs_for_bound("pre-recorded", &mut encoder, &bound, |binding| {
245            self.output_binding(binding).cloned()
246        })?;
247        {
248            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
249                label: Some("vyre pre-recorded persistent compute"),
250                timestamp_writes: None,
251            });
252            pass.set_pipeline(&self.pipeline);
253            for (i, bg) in bind_groups.iter().enumerate() {
254                let bind_group_index = u32::try_from(i).map_err(|_| {
255                    BackendError::new(
256                        "pre-recorded bind group index exceeds u32::MAX. Fix: reduce bind group fanout before recording.",
257                    )
258                })?;
259                pass.set_bind_group(bind_group_index, bg.as_ref(), &[]);
260            }
261            if let Some(indirect) = &self.indirect {
262                let indirect_handle = bound
263                    .iter()
264                    .find(|(info, _)| info.name.as_ref() == indirect.count_buffer.as_str())
265                    .map(|(_, handle)| *handle)
266                    .ok_or_else(|| {
267                        BackendError::new(format!(
268                            "indirect dispatch count buffer `{}` not bound in pre-recorded dispatch. Fix: supply the declared buffer handle.",
269                            indirect.count_buffer
270                        ))
271                    })?;
272                pass.dispatch_workgroups_indirect(indirect_handle.buffer(), indirect.count_offset);
273            } else {
274                pass.dispatch_workgroups(workgroups[0], workgroups[1], workgroups[2]);
275            }
276        }
277
278        let mut handles = Vec::new();
279        vyre_driver::allocation::try_reserve_vec_to_capacity(&mut handles, bound.len()).map_err(
280            |source| {
281                BackendError::new(format!(
282                    "pre-recorded dispatch handle retention could not reserve {} handle slot(s): {source}. Fix: split bound resources before recording.",
283                    bound.len()
284                ))
285            },
286        )?;
287        handles.extend(bound.iter().map(|(_, handle)| (*handle).clone()));
288        let mut output_handles = Vec::new();
289        vyre_driver::allocation::try_reserve_vec_to_capacity(&mut output_handles, outputs.len())
290            .map_err(|source| {
291                BackendError::new(format!(
292                    "pre-recorded output handle retention could not reserve {} output slot(s): {source}. Fix: split output resources before recording.",
293                    outputs.len()
294                ))
295            })?;
296        output_handles.extend(outputs.iter().cloned());
297        Ok(PrerecordedDispatch {
298            cb: Mutex::new(Some(encoder.finish())),
299            bind_groups,
300            handles,
301            output_handles,
302            device: device.clone(),
303            queue: queue.clone(),
304        })
305    }
306
307    /// Upload borrowed host inputs, allocate output handles, and pre-record
308    /// one persistent dispatch using this pipeline's device.
309    ///
310    /// # Errors
311    ///
312    /// Returns a backend error when upload, output allocation, or command
313    /// recording fails.
314    pub fn prerecord_borrowed_dispatch(
315        &self,
316        inputs: &[&[u8]],
317        workgroups: [u32; 3],
318    ) -> Result<PrerecordedDispatch, BackendError> {
319        let (input_handles, output_handles) = self.legacy_handles_from_inputs(inputs)?;
320        self.prerecord_persistent_dispatch(&input_handles, &output_handles, None, workgroups)
321    }
322}
323
324fn bind_handles<'a>(
325    bindings: &'a [BufferBindingInfo],
326    inputs: &'a [GpuBufferHandle],
327    outputs: &'a [GpuBufferHandle],
328    params: Option<&'a GpuBufferHandle>,
329) -> Result<SmallVec<[(&'a BufferBindingInfo, &'a GpuBufferHandle); 8]>, BackendError> {
330    let mut input_index = 0usize;
331    let mut output_index = 0usize;
332    let mut params_used = false;
333    let mut bound = SmallVec::new();
334    vyre_foundation::allocation::try_reserve_smallvec_to_capacity(&mut bound, bindings.len())
335        .map_err(|source| {
336            BackendError::new(format!(
337                "pre-recorded binding resolution could not reserve {} binding slot(s): {source}. Fix: split bound resources before recording.",
338                bindings.len()
339            ))
340        })?;
341    for info in bindings {
342        if info.kind == vyre_foundation::ir::MemoryKind::Shared {
343            continue;
344        }
345        let handle = if info.is_output {
346            let handle = outputs.get(output_index).ok_or_else(|| {
347                BackendError::new(format!(
348                    "pre-recorded dispatch missing output handle for binding {} (`{}`). Fix: pass one output handle per output BufferDecl.",
349                    info.binding, info.name
350                ))
351            })?;
352            output_index += 1;
353            handle
354        } else if matches!(
355            info.kind,
356            vyre_foundation::ir::MemoryKind::Uniform | vyre_foundation::ir::MemoryKind::Push
357        ) && params.is_some()
358            && !params_used
359        {
360            params_used = true;
361            if let Some(handle) = params {
362                handle
363            } else {
364                return Err(BackendError::new(
365                    "pre-recorded dispatch parameter handle disappeared after validation. Fix: retry recording with a stable params handle.",
366                ));
367            }
368        } else {
369            let handle = inputs.get(input_index).ok_or_else(|| {
370                BackendError::new(format!(
371                    "pre-recorded dispatch missing input handle for binding {} (`{}`). Fix: pass non-output handles in BufferDecl order.",
372                    info.binding, info.name
373                ))
374            })?;
375            input_index += 1;
376            handle
377        };
378        validate_handle("pre-recorded", info, handle)?;
379        bound.push((info, handle));
380    }
381    if input_index != inputs.len() {
382        return Err(BackendError::new(format!(
383            "pre-recorded dispatch received {} input handles but consumed {input_index}. Fix: pass handles matching non-output BufferDecl order.",
384            inputs.len()
385        )));
386    }
387    if output_index != outputs.len() {
388        return Err(BackendError::new(format!(
389            "pre-recorded dispatch received {} output handles but consumed {output_index}. Fix: pass handles matching output BufferDecl order.",
390            outputs.len()
391        )));
392    }
393    Ok(bound)
394}
vyre_driver_wgpu/runtime/prerecorded.rs

vyre_driver_wgpu/runtime/
prerecorded.rs
