vyre_driver_wgpu/runtime/
readback_ring.rs

1//! Async readback ring (Innovation I.5).
2//!
3//! Blocking readback submits a copy + device.poll(Wait) that stalls
4//! the submit queue. Under high dispatch rate this ruins latency and
5//! throughput  -  the GPU goes idle while the CPU waits.
6//!
7//! The readback ring threads N staging buffers. Dispatch \`i\` writes
8//! to \`ring[i % N]\`; the copy submits immediately and readback
9//! happens asynchronously via \`map_async\`. Dispatch \`i+1\` runs in
10//! parallel with readback \`i\`'s copy.
11
12use crossbeam_channel::Receiver;
13use dashmap::mapref::entry::Entry;
14use dashmap::DashMap;
15use rustc_hash::FxHasher;
16use std::hash::BuildHasherDefault;
17use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
18use std::sync::Arc;
19use vyre_driver::accounting::{atomic_max_u64, rebasing_atomic_next_u64};
20use vyre_driver::backend::BackendError;
21
22use crate::staging_reserve::reserve_backend_vec;
23
24const MIN_RING_SIZE: usize = 2;
25const MAX_RING_SIZE: usize = 256;
26const DEFAULT_RING_SLOTS: usize = 256;
27const RING_CAPACITY_GRANULARITY: u64 = 4096;
28const SLOT_FREE: u8 = 0;
29const SLOT_PENDING: u8 = 1;
30const SLOT_READY: u8 = 2;
31const SLOT_ERROR: u8 = 3;
32
33/// Result type produced by one `map_async` callback.
34pub type MapResult = Result<(), wgpu::BufferAsyncError>;
35
36/// Statistics collected by the ring at runtime.
37#[derive(Debug, Default)]
38pub struct RingStats {
39    /// Total dispatches queued.
40    pub dispatches: AtomicU64,
41    /// Readbacks that blocked waiting on map_async.
42    pub readback_stalls: AtomicU64,
43    /// Max outstanding (in-flight) copies.
44    pub peak_inflight: AtomicU64,
45}
46
47impl RingStats {
48    /// Record one dispatch; returns the monotonic dispatch index.
49    pub fn record_dispatch(&self) -> u64 {
50        rebasing_atomic_next_u64(
51            &self.dispatches,
52            0,
53            Ordering::Relaxed,
54            Ordering::Relaxed,
55            Ordering::Relaxed,
56            |_, _| {
57                tracing::error!(
58                    "readback ring dispatch counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
59                );
60            },
61        )
62    }
63
64    /// Record a stall.
65    pub fn record_stall(&self) {
66        rebasing_atomic_next_u64(
67            &self.readback_stalls,
68            0,
69            Ordering::Relaxed,
70            Ordering::Relaxed,
71            Ordering::Relaxed,
72            |_, _| {
73                tracing::error!(
74                    "readback ring stall counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
75                );
76            },
77        );
78    }
79
80    /// Update the peak-in-flight watermark.
81    pub fn update_peak(&self, current: u64) {
82        atomic_max_u64(&self.peak_inflight, current, Ordering::AcqRel);
83    }
84}
85
86/// Lifecycle state for one ring slot.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum SlotState {
89    /// Slot is available for new writes.
90    Free,
91    /// Copy has been submitted, data will be ready after fence.
92    Pending,
93    /// Map has completed and data is visible to the host.
94    Ready,
95    /// Mapping failed and the slot must be collected as an error.
96    Error,
97}
98
99/// GPU-aware ring slot.
100pub struct GpuSlot {
101    /// Underlying wgpu buffer.
102    pub buffer: wgpu::Buffer,
103    /// Atomic lifecycle state (0: Free, 1: Pending, 2: Ready).
104    pub state: Arc<std::sync::atomic::AtomicU8>,
105    byte_len: AtomicU64,
106    mapped_len: AtomicU64,
107    capacity: u64,
108}
109
110/// Submitted copy ticket for one readback-ring slot.
111pub struct ReadbackTicket {
112    idx: usize,
113    byte_len: u64,
114    mapped_len: u64,
115}
116
117/// Size-classed collection of readback rings for direct dispatch.
118pub struct ReadbackRingSet {
119    rings: DashMap<u64, Arc<ReadbackRing>, BuildHasherDefault<FxHasher>>,
120    slots_per_ring: usize,
121}
122
123impl Default for ReadbackRingSet {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129impl ReadbackRingSet {
130    /// Construct an empty ring set using the default slot count.
131    #[must_use]
132    pub fn new() -> Self {
133        Self {
134            rings: DashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
135            slots_per_ring: readback_ring_slots_from_env(),
136        }
137    }
138
139    /// Construct an empty ring set from a raw slot-count setting.
140    ///
141    /// Passing `None` uses the production default. This keeps test and embedded
142    /// callers off process-global environment mutation while preserving the same
143    /// parser and clamping semantics as [`Self::new`].
144    #[must_use]
145    pub fn with_requested_slots(raw_slots: Option<&str>) -> Self {
146        Self {
147            rings: DashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
148            slots_per_ring: readback_ring_slots_from_raw(raw_slots),
149        }
150    }
151
152    /// Return the ring whose staging slots can hold `byte_len`.
153    ///
154    /// # Errors
155    ///
156    /// Returns a backend error if the requested byte length overflows wgpu copy
157    /// alignment.
158    pub fn ring_for(
159        &self,
160        device: &wgpu::Device,
161        byte_len: u64,
162    ) -> Result<Arc<ReadbackRing>, BackendError> {
163        let capacity = Self::capacity_class_for(byte_len)?;
164        self.ring_for_capacity(device, capacity)
165    }
166
167    /// Return a ring for an already-normalized capacity class.
168    #[inline]
169    pub(crate) fn ring_for_capacity(
170        &self,
171        device: &wgpu::Device,
172        capacity: u64,
173    ) -> Result<Arc<ReadbackRing>, BackendError> {
174        Ok(match self.rings.entry(capacity) {
175            Entry::Occupied(entry) => Arc::clone(entry.get()),
176            Entry::Vacant(entry) => {
177                let ring = Arc::new(ReadbackRing::new(device, self.slots_per_ring, capacity)?);
178                entry.insert(Arc::clone(&ring));
179                ring
180            }
181        })
182    }
183
184    /// Convert an arbitrary byte length to the ring capacity class used for
185    /// ring sizing.
186    #[inline]
187    pub(crate) fn capacity_class(byte_len: u64) -> Result<u64, BackendError> {
188        Self::capacity_class_for(byte_len)
189    }
190
191    /// Convert an arbitrary byte length to the ring capacity class used for
192    /// ring sizing.
193    #[inline]
194    pub(crate) fn capacity_class_for(byte_len: u64) -> Result<u64, BackendError> {
195        ring_capacity_class(byte_len)
196    }
197
198    /// Return an existing size-classed ring without taking exclusive access.
199    ///
200    /// # Errors
201    ///
202    /// Returns a backend error if the requested byte length overflows wgpu copy
203    /// alignment.
204    pub fn existing_ring_for(
205        &self,
206        byte_len: u64,
207    ) -> Result<Option<Arc<ReadbackRing>>, BackendError> {
208        let capacity = Self::capacity_class(byte_len)?;
209        Ok(self.existing_ring_for_capacity(capacity))
210    }
211
212    /// Return an existing size-classed ring without taking exclusive access.
213    #[inline]
214    pub(crate) fn existing_ring_for_capacity(&self, capacity: u64) -> Option<Arc<ReadbackRing>> {
215        self.rings
216            .get(&capacity)
217            .map(|ring| Arc::clone(ring.value()))
218    }
219
220    /// Number of slots configured for each runtime ring instance.
221    #[must_use]
222    pub fn slots_per_ring(&self) -> usize {
223        self.slots_per_ring
224    }
225}
226
227/// Async readback ring buffer with GPU-resident staging buffers.
228pub struct ReadbackRing {
229    slots: Vec<GpuSlot>,
230    stats: Arc<RingStats>,
231    next_idx: AtomicU64,
232}
233
234impl ReadbackRing {
235    /// Construct a ring with N staging buffers.
236    #[must_use]
237    pub fn new(device: &wgpu::Device, size: usize, buffer_size: u64) -> Result<Self, BackendError> {
238        let size = size.clamp(MIN_RING_SIZE, MAX_RING_SIZE);
239        let capacity = staging_capacity(buffer_size)?;
240        let mut slots = Vec::new();
241        reserve_backend_vec(&mut slots, size, "readback ring slot table")?;
242        for i in 0..size {
243            let buffer = device.create_buffer(&wgpu::BufferDescriptor {
244                label: Some(&format!("vyre readback ring slot {i}")),
245                size: capacity,
246                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
247                mapped_at_creation: false,
248            });
249            slots.push(GpuSlot {
250                buffer,
251                state: Arc::new(std::sync::atomic::AtomicU8::new(SLOT_FREE)),
252                byte_len: AtomicU64::new(0),
253                mapped_len: AtomicU64::new(0),
254                capacity,
255            });
256        }
257        Ok(Self {
258            slots,
259            stats: Arc::new(RingStats::default()),
260            next_idx: AtomicU64::new(0),
261        })
262    }
263
264    /// Ensure slot `idx` is reusable for a fresh readback: either already
265    /// `SLOT_FREE`, or `SLOT_PENDING` that completes to `SLOT_FREE` after one
266    /// device poll. Any other terminal state is a caller contract violation and
267    /// is reported as a distinct, fail-closed error.
268    ///
269    /// VYRE-WGPU-002: earlier code conflated `SLOT_READY` / `SLOT_ERROR` with a
270    /// single misleading wrap-overflow message. We now name each
271    /// state (but we DO NOT silently recycle an uncollected `SLOT_READY` slot).
272    /// Recycling would unmap and discard a completed-but-uncollected readback
273    /// a silent recall loss (Law 10). The caller MUST collect every readback
274    /// before the ring wraps back to its slot; if it has not, we fail closed so
275    /// the data loss is impossible to miss.
276    fn ensure_slot_reusable(
277        &self,
278        idx: usize,
279        slot: &GpuSlot,
280        device: &wgpu::Device,
281    ) -> Result<(), BackendError> {
282        let mut state = slot.state.load(Ordering::Acquire);
283        if state == SLOT_PENDING {
284            self.stats.record_stall();
285            crate::runtime::device::poll_device_once(device)?;
286            state = slot.state.load(Ordering::Acquire);
287        }
288        match state {
289            SLOT_FREE => Ok(()),
290            SLOT_READY => Err(BackendError::new(format!(
291                "readback ring slot {idx} holds an uncollected completed readback (SLOT_READY). Fix: collect every ReadbackTicket via collect_slot_into before the ring wraps back to this slot (recycling it would silently drop the prior result (a recall loss))."
292            ))),
293            SLOT_ERROR => Err(BackendError::new(format!(
294                "readback ring slot {idx} is in SLOT_ERROR (prior map_async failed) and was not collected before reuse. Fix: collect error slots via collect_slot_into before submitting new readbacks to the same slot."
295            ))),
296            SLOT_PENDING => Err(BackendError::new(format!(
297                "readback ring slot {idx} is still SLOT_PENDING after a device poll, the prior readback has not completed. Fix: increase ring depth (more slots) or collect outstanding readbacks before submitting more."
298            ))),
299            other => Err(BackendError::new(format!(
300                "readback ring slot {idx} has unexpected state {other}. Fix: do not modify readback ring slot state outside the ring API."
301            ))),
302        }
303    }
304
305    /// Record a readback copy into the next available ring slot.
306    ///
307    /// The caller must submit the encoder and then arm the returned ticket with
308    /// [`Self::arm_ticket`]. This path lets the main dispatch encoder copy into
309    /// preallocated ring slots instead of allocating a fresh staging buffer per
310    /// output.
311    ///
312    /// # Errors
313    ///
314    /// Returns [`BackendError`] if the byte range cannot be represented, the
315    /// ring slot is not reusable (an uncollected `SLOT_READY`/`SLOT_ERROR` slot
316    /// or a still-pending slot after a device poll), or the requested readback
317    /// exceeds slot capacity.
318    pub fn record_copy(
319        &self,
320        device: &wgpu::Device,
321        encoder: &mut wgpu::CommandEncoder,
322        src_buffer: &wgpu::Buffer,
323        src_offset: u64,
324        byte_len: u64,
325    ) -> Result<ReadbackTicket, BackendError> {
326        let idx = self.next_slot_index()?;
327        let slot = &self.slots[idx];
328        let mapped_len = aligned_copy_len(byte_len)?;
329        if mapped_len > slot.capacity {
330            return Err(BackendError::new(format!(
331                "readback request of {byte_len} bytes ({} bytes after wgpu copy alignment) exceeds ring slot capacity {} bytes. Fix: construct ReadbackRing with a buffer_size at least as large as the largest readback.",
332                mapped_len, slot.capacity
333            )));
334        }
335
336        self.ensure_slot_reusable(idx, slot, device)?;
337
338        slot.byte_len.store(byte_len, Ordering::Release);
339        slot.mapped_len.store(mapped_len, Ordering::Release);
340        slot.state.store(SLOT_PENDING, Ordering::Release);
341        if mapped_len != 0 {
342            encoder.copy_buffer_to_buffer(src_buffer, src_offset, &slot.buffer, 0, mapped_len);
343        } else {
344            slot.state.store(SLOT_READY, Ordering::Release);
345        }
346        self.stats.record_dispatch();
347        Ok(ReadbackTicket {
348            idx,
349            byte_len,
350            mapped_len,
351        })
352    }
353
354    /// Arm a submitted ticket by registering its `map_async` callback.
355    ///
356    /// # Errors
357    ///
358    /// Returns [`BackendError`] when `ticket` does not reference a live slot.
359    pub fn arm_ticket(
360        &self,
361        ticket: &ReadbackTicket,
362    ) -> Result<(Receiver<MapResult>, Arc<AtomicBool>), BackendError> {
363        let Some(slot) = self.slots.get(ticket.idx) else {
364            return Err(BackendError::new(format!(
365                "readback ring ticket slot {} is out of bounds for {} slots. Fix: keep tickets paired with their originating ring.",
366                ticket.idx,
367                self.slots.len()
368            )));
369        };
370        let (sender, receiver) = crossbeam_channel::bounded(1);
371        let ready = Arc::new(AtomicBool::new(false));
372        if ticket.mapped_len == 0 {
373            if let Err(error) = sender.send(Ok(())) {
374                tracing::error!(
375                    ?error,
376                    "readback ring zero-length callback result was lost because the receiver dropped"
377                );
378            }
379            ready.store(true, Ordering::Release);
380            return Ok((receiver, ready));
381        }
382
383        let state = Arc::clone(&slot.state);
384        let ready_cb = Arc::clone(&ready);
385        slot.buffer
386            .slice(0..ticket.mapped_len)
387            .map_async(wgpu::MapMode::Read, move |result| {
388                match &result {
389                    Ok(()) => state.store(SLOT_READY, Ordering::Release),
390                    Err(error) => {
391                        tracing::error!(
392                            "readback ring map_async failed: {error:?}. Fix: inspect device health and readback buffer usage."
393                        );
394                        state.store(SLOT_ERROR, Ordering::Release);
395                    }
396                }
397                if let Err(error) = sender.send(result) {
398                    tracing::error!(
399                        ?error,
400                        "readback ring callback result was lost because the receiver dropped"
401                    );
402                }
403                ready_cb.store(true, Ordering::Release);
404            });
405        Ok((receiver, ready))
406    }
407
408    /// Expose a ready ticket's mapped bytes to `visitor`, then free the slot.
409    ///
410    /// # Errors
411    ///
412    /// Returns [`BackendError`] when the ticket is stale, the slot is not ready,
413    /// or mapped length metadata is inconsistent.
414    pub fn with_mapped_ticket<R>(
415        &self,
416        ticket: &ReadbackTicket,
417        visitor: impl FnOnce(&[u8]) -> Result<R, BackendError>,
418    ) -> Result<R, BackendError> {
419        let Some(slot) = self.slots.get(ticket.idx) else {
420            return Err(BackendError::new(format!(
421                "readback ring ticket slot {} is out of bounds for {} slots. Fix: keep tickets paired with their originating ring.",
422                ticket.idx,
423                self.slots.len()
424            )));
425        };
426        match slot.state.load(Ordering::Acquire) {
427            SLOT_READY => {}
428            SLOT_ERROR => {
429                slot.byte_len.store(0, Ordering::Release);
430                slot.mapped_len.store(0, Ordering::Release);
431                slot.state.store(SLOT_FREE, Ordering::Release);
432                return Err(BackendError::new(
433                    "readback ring map_async failed. Fix: inspect GPU device health and ensure the slot buffer has MAP_READ usage.",
434                ));
435            }
436            _ => {
437                return Err(BackendError::new(
438                    "readback ring ticket was collected before its map callback completed. Fix: poll the device or wait for the submitted GPU work before collection.",
439                ));
440            }
441        }
442
443        let len = usize::try_from(ticket.byte_len).map_err(|source| {
444            BackendError::new(format!(
445                "readback ring byte length {} cannot fit usize: {source}. Fix: split the readback before collecting it.",
446                ticket.byte_len
447            ))
448        })?;
449        if ticket.mapped_len == 0 {
450            slot.byte_len.store(0, Ordering::Release);
451            slot.mapped_len.store(0, Ordering::Release);
452            slot.state.store(SLOT_FREE, Ordering::Release);
453            return visitor(&[]);
454        }
455        let view = slot.buffer.slice(0..ticket.mapped_len).get_mapped_range();
456        if len > view.len() {
457            let mapped_len = view.len();
458            drop(view);
459            slot.buffer.unmap();
460            slot.byte_len.store(0, Ordering::Release);
461            slot.mapped_len.store(0, Ordering::Release);
462            slot.state.store(SLOT_FREE, Ordering::Release);
463            return Err(BackendError::new(format!(
464                "readback ring mapped length {mapped_len} is shorter than requested length {len}. Fix: keep ticket and slot byte lengths synchronized."
465            )));
466        }
467        let result = visitor(&view[..len]);
468        drop(view);
469        slot.buffer.unmap();
470        slot.byte_len.store(0, Ordering::Release);
471        slot.mapped_len.store(0, Ordering::Release);
472        slot.state.store(SLOT_FREE, Ordering::Release);
473        result
474    }
475
476    /// Submit a copy from `src_buffer` at `src_offset` and mark the slot pending.
477    ///
478    /// `src_offset` is the byte offset within `src_buffer` to copy from. Pass
479    /// `0` to read from the start of the buffer. This mirrors the `src_offset`
480    /// parameter accepted by `record_copy`; callers that need a sub-range of
481    /// the source buffer must supply a non-zero offset here rather than
482    /// wrapping a slice (the wgpu copy API requires aligned buffer offsets).
483    ///
484    /// # Errors
485    /// Returns `BackendError` if encoder or queue submission fails.
486    pub fn submit_readback(
487        &self,
488        device: &wgpu::Device,
489        queue: &wgpu::Queue,
490        src_buffer: &wgpu::Buffer,
491        src_offset: u64,
492        byte_len: u64,
493    ) -> Result<usize, BackendError> {
494        let idx = self.next_slot_index()?;
495        let slot = &self.slots[idx];
496        let mapped_len = aligned_copy_len(byte_len)?;
497        if mapped_len > slot.capacity {
498            return Err(BackendError::new(format!(
499                "readback request of {byte_len} bytes ({} bytes after wgpu copy alignment) exceeds ring slot capacity {} bytes. Fix: construct ReadbackRing with a buffer_size at least as large as the largest readback.",
500                mapped_len, slot.capacity
501            )));
502        }
503
504        self.ensure_slot_reusable(idx, slot, device)?;
505
506        let state_clone = Arc::clone(&slot.state);
507        slot.byte_len.store(byte_len, Ordering::Release);
508        slot.mapped_len.store(mapped_len, Ordering::Release);
509        state_clone.store(SLOT_PENDING, Ordering::Release);
510
511        if mapped_len == 0 {
512            state_clone.store(SLOT_READY, Ordering::Release);
513            self.stats.record_dispatch();
514            return Ok(idx);
515        }
516
517        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
518            label: Some("vyre readback ring copy"),
519        });
520        encoder.copy_buffer_to_buffer(src_buffer, src_offset, &slot.buffer, 0, mapped_len);
521        queue.submit(std::iter::once(encoder.finish()));
522
523        slot.buffer
524            .slice(0..mapped_len)
525            .map_async(wgpu::MapMode::Read, move |result| {
526                match result {
527                    Ok(()) => state_clone.store(SLOT_READY, Ordering::Release),
528                    Err(error) => {
529                        tracing::error!(
530                            "readback ring map_async failed: {error:?}. Fix: inspect device health and readback buffer usage."
531                        );
532                        state_clone.store(SLOT_ERROR, Ordering::Release);
533                    }
534                }
535            });
536
537        self.stats.record_dispatch();
538
539        Ok(idx)
540    }
541
542    /// Try to collect data from a specific slot.
543    ///
544    /// # Errors
545    ///
546    /// Returns [`BackendError`] when `idx` is out of bounds or `map_async`
547    /// failed for the slot.
548    pub fn collect_slot(
549        &self,
550        device: &wgpu::Device,
551        idx: usize,
552    ) -> Result<Option<Vec<u8>>, BackendError> {
553        let mut data = Vec::new();
554        if self.collect_slot_into(device, idx, &mut data)?.is_some() {
555            Ok(Some(data))
556        } else {
557            Ok(None)
558        }
559    }
560
561    /// Try to collect data from a specific slot into a caller-owned buffer.
562    ///
563    /// Reusing `out` avoids an allocation on every ready readback. The buffer is
564    /// cleared before bytes are appended.
565    ///
566    /// # Errors
567    ///
568    /// Returns [`BackendError`] when `idx` is out of bounds or `map_async`
569    /// failed for the slot.
570    pub fn collect_slot_into(
571        &self,
572        device: &wgpu::Device,
573        idx: usize,
574        out: &mut Vec<u8>,
575    ) -> Result<Option<usize>, BackendError> {
576        let Some(slot) = self.slots.get(idx) else {
577            return Err(BackendError::new(format!(
578                "readback ring slot index {idx} is out of bounds for {} slots. Fix: collect only indices returned by submit_readback.",
579                self.slots.len()
580            )));
581        };
582        match slot.state.load(Ordering::Acquire) {
583            SLOT_READY => {
584                let len = self.copy_ready_slot_into(idx, out)?;
585                Ok(Some(len))
586            }
587            SLOT_ERROR => {
588                slot.byte_len.store(0, Ordering::Release);
589                slot.mapped_len.store(0, Ordering::Release);
590                slot.state.store(SLOT_FREE, Ordering::Release);
591                Err(BackendError::new(
592                    "readback ring map_async failed. Fix: inspect GPU device health and ensure the slot buffer has MAP_READ usage.",
593                ))
594            }
595            _ => {
596                crate::runtime::device::poll_device_once(device)?;
597                Ok(None)
598            }
599        }
600    }
601
602    fn copy_ready_slot_into(&self, idx: usize, out: &mut Vec<u8>) -> Result<usize, BackendError> {
603        let slot = &self.slots[idx];
604        let byte_len = slot.byte_len.load(Ordering::Acquire);
605        let mapped_len = slot.mapped_len.load(Ordering::Acquire);
606        let len = usize::try_from(byte_len).map_err(|source| {
607            BackendError::new(format!(
608                "readback ring byte length {byte_len} cannot fit usize: {source}. Fix: split the readback before collecting it."
609            ))
610        })?;
611        if mapped_len != 0 {
612            let view = slot.buffer.slice(0..mapped_len).get_mapped_range();
613            let bytes = &view[..len];
614            if out.len() == len {
615                out.copy_from_slice(bytes);
616            } else {
617                if len > out.capacity() {
618                    let additional = len - out.capacity();
619                    out.try_reserve_exact(additional).map_err(|source| {
620                        BackendError::new(format!(
621                            "readback ring collection could not reserve {len} output bytes exactly: {source}. Fix: lower max_output_bytes or collect readback in smaller shards."
622                        ))
623                    })?;
624                }
625                out.clear();
626                out.extend_from_slice(bytes);
627            }
628            drop(view);
629            slot.buffer.unmap();
630        } else {
631            out.clear();
632        }
633        slot.byte_len.store(0, Ordering::Release);
634        slot.mapped_len.store(0, Ordering::Release);
635        slot.state.store(SLOT_FREE, Ordering::Release);
636        Ok(len)
637    }
638
639    #[inline]
640    fn next_slot_index(&self) -> Result<usize, BackendError> {
641        let slot_len = u64::try_from(self.slots.len()).map_err(|source| {
642            BackendError::new(format!(
643                "readback ring slot count {} cannot fit u64: {source}. Fix: reduce readback ring slot count.",
644                self.slots.len()
645            ))
646        })?;
647        if slot_len == 0 {
648            return Err(BackendError::new(
649                "readback ring has zero slots. Fix: construct rings with at least two slots.",
650            ));
651        }
652        let next = rebasing_atomic_next_u64(
653            &self.next_idx,
654            0,
655            Ordering::Relaxed,
656            Ordering::Relaxed,
657            Ordering::Relaxed,
658            |_, _| {
659                tracing::error!(
660                    "readback ring slot counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
661                );
662            },
663        );
664        usize::try_from(next % slot_len).map_err(|source| {
665            BackendError::new(format!(
666                "readback ring slot index cannot fit usize: {source}. Fix: reduce readback ring slot count."
667            ))
668        })
669    }
670}
671
672#[inline]
673fn staging_capacity(byte_len: u64) -> Result<u64, BackendError> {
674    aligned_copy_len(byte_len).map_err(|error| {
675        tracing::warn!(
676            "readback ring staging capacity overflowed for {byte_len} bytes: {error}. Fix: shard the readback buffer before constructing the ring."
677        );
678        error
679    }).map(|len| len.max(4))
680}
681
682#[inline]
683fn ring_capacity_class(byte_len: u64) -> Result<u64, BackendError> {
684    let aligned = aligned_copy_len(byte_len)?.max(4);
685    aligned
686        .checked_add(RING_CAPACITY_GRANULARITY - 1)
687        .map(|len| len & !(RING_CAPACITY_GRANULARITY - 1))
688        .ok_or_else(|| {
689            BackendError::new(
690                "readback ring capacity class overflows u64. Fix: split the readback before submitting it to the ring.",
691            )
692        })
693}
694
695#[inline]
696fn aligned_copy_len(byte_len: u64) -> Result<u64, BackendError> {
697    crate::numeric::WGPU_NUMERIC.align_up_u64(byte_len, 4, 0, "readback byte length")
698}
699
700fn readback_ring_slots_from_env() -> usize {
701    let raw = std::env::var("VYRE_WGPU_READBACK_RING_SLOTS").ok();
702    readback_ring_slots_from_raw(raw.as_deref())
703}
704
705fn readback_ring_slots_from_raw(raw: Option<&str>) -> usize {
706    let Some(raw) = raw else {
707        return DEFAULT_RING_SLOTS;
708    };
709    let slots = match raw.parse::<usize>() {
710        Ok(0) => {
711            tracing::warn!(
712                "VYRE_WGPU_READBACK_RING_SLOTS=0 is invalid for GPU readback rings; defaulting to {MIN_RING_SIZE}. Fix: set it to a positive integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
713            );
714            MIN_RING_SIZE
715        }
716        Ok(value) if value > MAX_RING_SIZE => {
717            tracing::warn!(
718                "VYRE_WGPU_READBACK_RING_SLOTS={value} exceeds the safe cap of {MAX_RING_SIZE}; clamping.
719                Fix: set it to an integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
720            );
721            MAX_RING_SIZE
722        }
723        Ok(value) => value,
724        Err(error) => {
725            tracing::warn!(
726                "VYRE_WGPU_READBACK_RING_SLOTS={raw:?} is invalid ({error:?}); defaulting to {DEFAULT_RING_SLOTS}. Fix: set it to a positive integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
727            );
728            DEFAULT_RING_SLOTS
729        }
730    };
731    slots.clamp(MIN_RING_SIZE, MAX_RING_SIZE)
732}
733
734#[cfg(test)]
735mod tests {
736    use super::*;
737
738    #[test]
739    fn capacity_class_classifies_by_alignment_and_granularity() {
740        assert_eq!(
741            ReadbackRingSet::capacity_class_for(16).unwrap(),
742            4096,
743            "16-byte requests must promote to 4096-byte slot class"
744        );
745        assert_eq!(
746            ReadbackRingSet::capacity_class_for(1).unwrap(),
747            4096,
748            "1-byte requests must promote to minimum aligned 4096-byte class"
749        );
750        assert_eq!(
751            ReadbackRingSet::capacity_class_for(4097).unwrap(),
752            8192,
753            "4KB boundary crossings must promote to the next class"
754        );
755    }
756
757    #[test]
758    fn existing_ring_for_and_capacity_variant_agree_on_lookup_key() {
759        let ring_set = ReadbackRingSet::new();
760        let from_raw = ring_set
761            .existing_ring_for(16)
762            .expect("Fix: lookup with raw byte length should not fail");
763        let from_class = ring_set.existing_ring_for_capacity(4096);
764        assert!(
765            from_raw.is_none() && from_class.is_none(),
766            "raw and capacity-based lookups should agree on an empty set"
767        );
768    }
769}
vyre_driver_wgpu/runtime/readback_ring.rs

vyre_driver_wgpu/runtime/
readback_ring.rs
