vyre_driver_wgpu/buffer/
handle.rs

1//! Public persistent GPU buffer handle.
2
3use std::cmp::{Ordering as CmpOrdering, Reverse};
4use std::collections::BinaryHeap;
5use std::hash::{Hash, Hasher};
6use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
7use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
8use std::time::Instant;
9
10use dashmap::DashMap;
11use rustc_hash::{FxHashMap, FxHasher};
12use smallvec::SmallVec;
13use vyre_driver::BackendError;
14
15use super::pool::PoolReturn;
16
17static NEXT_BUFFER_ID: AtomicU64 = AtomicU64::new(1);
18static RESIDENT_BUFFERS: OnceLock<DashMap<u64, Weak<GpuBufferInner>>> = OnceLock::new();
19const STAGING_BUFFER_POOL_CLASS_CAP: usize = 16;
20
21fn resident_buffers() -> &'static DashMap<u64, Weak<GpuBufferInner>> {
22    RESIDENT_BUFFERS.get_or_init(DashMap::new)
23}
24
25fn pointer_identity_key<T>(ptr: *const T) -> u64 {
26    let mut hasher = FxHasher::default();
27    ptr.addr().hash(&mut hasher);
28    hasher.finish()
29}
30
31/// Cheaply cloneable handle for a GPU-resident buffer.
32///
33/// The handle records the byte length originally requested by the caller,
34/// the backing allocation length, the logical element count, and the actual
35/// usage flags used to create the underlying `wgpu::Buffer`.
36#[derive(Clone)]
37pub struct GpuBufferHandle {
38    inner: Arc<GpuBufferInner>,
39}
40
41struct GpuBufferInner {
42    id: u64,
43    buffer: Arc<wgpu::Buffer>,
44    byte_len: u64,
45    allocation_len: u64,
46    element_count: usize,
47    usage: wgpu::BufferUsages,
48    pool_return: Option<PoolReturn>,
49}
50
51/// Snapshot of [`StagingBufferPool`] counters.
52#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
53pub struct StagingBufferPoolStats {
54    /// Number of fresh GPU buffer allocations.
55    pub allocations: usize,
56    /// Number of times a free buffer was reused.
57    pub hits: usize,
58}
59
60/// Device-local staging buffer pool keyed by `(size, usage)`.
61///
62/// Hot dispatch paths (e.g. [`GpuBufferHandle::readback_until`]) acquire
63/// readback staging buffers from this pool instead of creating a fresh
64/// `wgpu::Buffer` on every call. Each `(size, usage)` class is capped at
65/// [`STAGING_BUFFER_POOL_CLASS_CAP`] entries; evictions drop the
66/// least-recently-used buffer.
67#[derive(Clone, Default)]
68pub struct StagingBufferPool {
69    inner: Arc<Mutex<StagingBufferPoolInner>>,
70}
71
72impl std::fmt::Debug for StagingBufferPool {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("StagingBufferPool").finish_non_exhaustive()
75    }
76}
77
78#[derive(Default)]
79struct StagingBufferPoolInner {
80    free: FxHashMap<(u64, u32), SmallVec<[wgpu::Buffer; STAGING_BUFFER_POOL_CLASS_CAP]>>,
81    allocations: usize,
82    hits: usize,
83}
84
85impl StagingBufferPool {
86    fn lock_inner(&self) -> MutexGuard<'_, StagingBufferPoolInner> {
87        self.inner.lock().unwrap_or_else(|error| {
88            tracing::error!(
89                "Vyre WGPU staging buffer pool lock was poisoned: {error}. Fix: discard the pool after a panic; continuing with recovered state."
90            );
91            error.into_inner()
92        })
93    }
94
95    /// Create an empty staging buffer pool.
96    #[must_use]
97    #[inline]
98    pub fn new() -> Self {
99        Self::default()
100    }
101
102    /// Return allocation and hit counters.
103    #[must_use]
104    pub fn stats(&self) -> StagingBufferPoolStats {
105        let inner = self.lock_inner();
106        StagingBufferPoolStats {
107            allocations: inner.allocations,
108            hits: inner.hits,
109        }
110    }
111
112    /// Acquire a staging buffer with exactly `size` bytes and `usage`.
113    ///
114    /// Reuses a free buffer when one is available; otherwise creates a fresh
115    /// GPU allocation and increments the allocation counter.
116    pub fn acquire(
117        &self,
118        device: &wgpu::Device,
119        size: u64,
120        usage: wgpu::BufferUsages,
121    ) -> wgpu::Buffer {
122        let key = (size, usage.bits());
123        let mut inner = self.lock_inner();
124        if let Some(buffers) = inner.free.get_mut(&key) {
125            if let Some(buffer) = buffers.pop() {
126                inner.hits += 1;
127                return buffer;
128            }
129        }
130        inner.allocations += 1;
131        drop(inner);
132        device.create_buffer(&wgpu::BufferDescriptor {
133            label: Some("vyre staging readback"),
134            size,
135            usage,
136            mapped_at_creation: false,
137        })
138    }
139
140    /// Release a staging buffer back to the pool.
141    ///
142    /// The buffer is pushed to the MRU position of its `(size, usage)` class.
143    /// If the class already holds 16 buffers, the LRU entry is dropped.
144    pub fn release(&self, buffer: wgpu::Buffer, size: u64, usage: wgpu::BufferUsages) {
145        let key = (size, usage.bits());
146        let mut inner = self.lock_inner();
147        let buffers = inner.free.entry(key).or_insert_with(SmallVec::new);
148        if buffers.len() == STAGING_BUFFER_POOL_CLASS_CAP {
149            buffers.remove(0);
150        }
151        buffers.push(buffer);
152    }
153}
154
155impl GpuBufferHandle {
156    /// Upload `bytes` into a new GPU buffer.
157    ///
158    /// The created buffer always includes `COPY_DST` so the upload is legal.
159    ///
160    /// # Errors
161    ///
162    /// Returns a backend error when the requested allocation length cannot fit
163    /// `u64`.
164    pub fn upload(
165        device: &wgpu::Device,
166        queue: &wgpu::Queue,
167        bytes: &[u8],
168        usage: wgpu::BufferUsages,
169    ) -> Result<Self, BackendError> {
170        let allocation_len = aligned_len(bytes.len())?;
171        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
172            label: Some("vyre persistent upload"),
173            size: allocation_len,
174            usage: usage | wgpu::BufferUsages::COPY_DST,
175            mapped_at_creation: false,
176        });
177        write_padded(queue, &buffer, bytes, allocation_len)?;
178        let logical_len = u64::try_from(bytes.len()).map_err(|source| {
179            BackendError::new(format!(
180                "GPU upload logical byte length cannot fit u64: {source}. Fix: split the dispatch input."
181            ))
182        })?;
183        Ok(Self::from_parts(
184            Arc::new(buffer),
185            logical_len,
186            allocation_len,
187            bytes.len(),
188            usage | wgpu::BufferUsages::COPY_DST,
189            None,
190        ))
191    }
192
193    /// Allocate a GPU-resident buffer without uploading host contents.
194    ///
195    /// # Errors
196    ///
197    /// Returns a backend error when `len` cannot be represented as a valid
198    /// wgpu buffer size.
199    pub fn alloc(
200        device: &wgpu::Device,
201        len: u64,
202        usage: wgpu::BufferUsages,
203    ) -> Result<Self, BackendError> {
204        let allocation_len = aligned_len_u64(len, "persistent GPU allocation length")?;
205        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
206            label: Some("vyre persistent alloc"),
207            size: allocation_len,
208            usage,
209            mapped_at_creation: false,
210        });
211        let host_len = usize::try_from(len).map_err(|error| {
212            BackendError::new(format!(
213                "GpuBufferHandle::alloc received logical byte length {len} that does not fit usize on this host: {error}. Fix: shard the GPU buffer before allocating or run on a host with a wide enough address space."
214            ))
215        })?;
216        Ok(Self::from_parts(
217            Arc::new(buffer),
218            len,
219            allocation_len,
220            host_len,
221            usage,
222            None,
223        ))
224    }
225
226    /// Download this GPU buffer into `out`.
227    ///
228    /// This is intended for terminal output and test assertions, not hot-loop
229    /// dispatch. The buffer must have `COPY_SRC` usage.
230    ///
231    /// # Errors
232    ///
233    /// Returns a backend error when the handle is not copy-readable or the GPU
234    /// mapping fails.
235    pub fn readback(
236        &self,
237        device: &wgpu::Device,
238        queue: &wgpu::Queue,
239        out: &mut Vec<u8>,
240    ) -> Result<(), BackendError> {
241        self.readback_until(device, None, queue, out, None)
242    }
243
244    /// Download the first `len` logical bytes of this GPU buffer into `out`.
245    ///
246    /// Hot paths that publish a device-side count should read back only the
247    /// counted prefix instead of the whole capacity-sized buffer. The copy is
248    /// rounded up to wgpu's 4-byte copy granularity internally, then truncated
249    /// back to exactly `len` bytes before returning.
250    ///
251    /// # Errors
252    ///
253    /// Returns a backend error when the handle is not copy-readable, `len`
254    /// exceeds the logical buffer length, or the GPU mapping fails.
255    pub fn readback_prefix(
256        &self,
257        device: &wgpu::Device,
258        queue: &wgpu::Queue,
259        len: u64,
260        out: &mut Vec<u8>,
261    ) -> Result<(), BackendError> {
262        self.readback_prefix_until(device, None, queue, len, out, None)
263    }
264
265    /// Download `len` logical bytes starting at `byte_offset` into `out`.
266    ///
267    /// The internal GPU copy is alignment-expanded when necessary, then the
268    /// returned host slice is trimmed back to exactly the requested range.
269    ///
270    /// # Errors
271    ///
272    /// Returns a backend error when the handle is not copy-readable, the range
273    /// exceeds the logical buffer length, or the GPU mapping fails.
274    pub fn readback_range(
275        &self,
276        device: &wgpu::Device,
277        queue: &wgpu::Queue,
278        byte_offset: u64,
279        len: u64,
280        out: &mut Vec<u8>,
281    ) -> Result<(), BackendError> {
282        self.readback_range_until(device, None, queue, byte_offset, len, out, None)
283    }
284
285    pub(crate) fn readback_until(
286        &self,
287        device: &wgpu::Device,
288        pool: Option<&StagingBufferPool>,
289        queue: &wgpu::Queue,
290        out: &mut Vec<u8>,
291        deadline: Option<Instant>,
292    ) -> Result<(), BackendError> {
293        self.readback_prefix_until(device, pool, queue, self.byte_len(), out, deadline)
294    }
295
296    pub(crate) fn readback_prefix_until(
297        &self,
298        device: &wgpu::Device,
299        pool: Option<&StagingBufferPool>,
300        queue: &wgpu::Queue,
301        len: u64,
302        out: &mut Vec<u8>,
303        deadline: Option<Instant>,
304    ) -> Result<(), BackendError> {
305        self.readback_range_until(device, pool, queue, 0, len, out, deadline)
306    }
307
308    pub(crate) fn readback_range_until(
309        &self,
310        device: &wgpu::Device,
311        pool: Option<&StagingBufferPool>,
312        queue: &wgpu::Queue,
313        byte_offset: u64,
314        len: u64,
315        out: &mut Vec<u8>,
316        deadline: Option<Instant>,
317    ) -> Result<(), BackendError> {
318        if !self.usage().contains(wgpu::BufferUsages::COPY_SRC) {
319            return Err(BackendError::new(
320                "GpuBufferHandle readback requires COPY_SRC usage. Fix: allocate terminal-output buffers with COPY_SRC.",
321            ));
322        }
323        let logical_end = byte_offset.checked_add(len).ok_or_else(|| {
324            BackendError::new(format!(
325                "GpuBufferHandle range readback overflows u64 at offset {byte_offset} len {len}. Fix: split the readback range before dispatch."
326            ))
327        })?;
328        if logical_end > self.byte_len() {
329            return Err(BackendError::new(format!(
330                "GpuBufferHandle range readback requested bytes [{byte_offset}..{logical_end}) from a {} byte buffer. Fix: clamp the requested range to the device-published count.",
331                self.byte_len()
332            )));
333        }
334        if len == 0 {
335            out.clear();
336            return Ok(());
337        }
338        let copy_start = byte_offset & !3;
339        let trim_start = byte_offset - copy_start;
340        let visible_copy_len = trim_start.checked_add(len).ok_or_else(|| {
341            BackendError::new(format!(
342                "GpuBufferHandle range readback copy length overflows u64 at trim {trim_start} len {len}. Fix: split the readback range before dispatch."
343            ))
344        })?;
345        let read_len = aligned_len_u64(visible_copy_len, "GPU readback visible copy length")?;
346        let copy_end = copy_start.checked_add(read_len).ok_or_else(|| {
347            BackendError::new(format!(
348                "GpuBufferHandle range readback aligned copy overflows u64 at start {copy_start} len {read_len}. Fix: split the readback range before dispatch."
349            ))
350        })?;
351        if copy_end > self.inner.allocation_len {
352            return Err(BackendError::new(format!(
353                "GpuBufferHandle range readback rounded bytes [{byte_offset}..{logical_end}) to aligned bytes [{copy_start}..{copy_end}), beyond allocation length {}. Fix: allocate buffers with 4-byte padding.",
354                self.inner.allocation_len
355            )));
356        }
357        let readback_usage = wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ;
358        let readback = if let Some(pool) = pool {
359            pool.acquire(device, read_len, readback_usage)
360        } else {
361            device.create_buffer(&wgpu::BufferDescriptor {
362                label: Some("vyre persistent handle readback"),
363                size: read_len,
364                usage: readback_usage,
365                mapped_at_creation: false,
366            })
367        };
368        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
369            label: Some("vyre persistent handle readback encoder"),
370        });
371        encoder.copy_buffer_to_buffer(self.buffer(), copy_start, &readback, 0, read_len);
372        let submission = queue.submit(std::iter::once(encoder.finish()));
373        let slice = readback.slice(0..read_len);
374        let (sender, receiver) = std::sync::mpsc::channel();
375        slice.map_async(wgpu::MapMode::Read, move |result| {
376            if let Err(error) = sender.send(result) {
377                tracing::error!(
378                    ?error,
379                    "persistent buffer readback map_async result was lost because the receiver dropped"
380                );
381            }
382        });
383        let mapping = if let Some(deadline) = deadline {
384            let mut backoff = crate::wait_backoff::AdaptiveWaitBackoff::from_micros(64, 2, 50, 5);
385            loop {
386                crate::runtime::device::poll_device_once(device)?;
387                match receiver.try_recv() {
388                    Ok(result) => break result,
389                    Err(std::sync::mpsc::TryRecvError::Empty) => {}
390                    Err(std::sync::mpsc::TryRecvError::Disconnected) => {
391                        return Err(BackendError::new(
392                            "persistent buffer readback channel closed before completion. Fix: keep the GPU device alive until readback completes.",
393                        ));
394                    }
395                }
396                let now = Instant::now();
397                if now >= deadline {
398                    return Err(BackendError::new(
399                        "dispatch cancelled after DispatchConfig.timeout before readback completed. Fix: raise DispatchConfig.timeout or split the program into smaller chunks.",
400                    ));
401                }
402                backoff.idle_for(deadline.saturating_duration_since(now));
403            }
404        } else {
405            crate::runtime::device::poll_device_wait_for(device, submission)?;
406            receiver
407                .recv_timeout(std::time::Duration::from_secs(30))
408                .map_err(|source| {
409                    BackendError::new(format!(
410                        "persistent buffer readback callback did not complete after submission wait: {source}. Fix: keep the GPU device alive and inspect driver callback progress."
411                    ))
412                })?
413        };
414        let result = mapping.map_err(|source| {
415            BackendError::new(format!(
416                "persistent buffer readback mapping failed: {source:?}. Fix: use COPY_SRC handles and MAP_READ staging buffers."
417            ))
418        });
419        result?;
420        let mapped = slice.get_mapped_range();
421        let visible_len = usize::try_from(len).map_err(|source| {
422            BackendError::new(format!(
423                "persistent buffer prefix length {len} cannot fit usize: {source}. Fix: split the buffer before readback.",
424            ))
425        })?;
426        let trim_start = usize::try_from(trim_start).map_err(|source| {
427            BackendError::new(format!(
428                "persistent buffer range trim offset {trim_start} cannot fit usize: {source}. Fix: split the buffer before readback.",
429            ))
430        })?;
431        let trim_end = trim_start.checked_add(visible_len).ok_or_else(|| {
432            BackendError::new(format!(
433                "persistent buffer range trim overflows usize at offset {trim_start} len {visible_len}. Fix: split the buffer before readback."
434            ))
435        })?;
436        let visible = &mapped[trim_start..trim_end];
437        if out.len() == visible_len {
438            out.copy_from_slice(visible);
439        } else {
440            if visible_len > out.capacity() {
441                let additional = visible_len - out.capacity();
442                out.try_reserve_exact(additional).map_err(|source| {
443                    BackendError::new(format!(
444                        "persistent buffer readback could not reserve {visible_len} output bytes exactly: {source}. Fix: lower max_output_bytes or stream readback in smaller shards."
445                    ))
446                })?;
447            }
448            out.clear();
449            out.extend_from_slice(visible);
450        }
451        drop(mapped);
452        readback.unmap();
453        if let Some(pool) = pool {
454            pool.release(readback, read_len, readback_usage);
455        }
456        Ok(())
457    }
458
459    /// Stable process-local handle id used for cache signatures.
460    #[must_use]
461    pub fn id(&self) -> u64 {
462        self.inner.id
463    }
464
465    /// Stable process-local identity for the backing GPU allocation.
466    ///
467    /// Unlike [`Self::id`], this survives pool release/reacquire cycles for the
468    /// same underlying `wgpu::Buffer`. Bind-group caches must key on this value
469    /// plus the logical binding range; otherwise hot dispatches miss every time
470    /// a pooled allocation is wrapped in a fresh handle.
471    #[must_use]
472    pub(crate) fn allocation_identity(&self) -> u64 {
473        pointer_identity_key(Arc::as_ptr(&self.inner.buffer))
474    }
475
476    /// Resolve a process-local resident buffer id back into a live GPU handle.
477    #[must_use]
478    pub fn from_resident_id(id: u64) -> Option<Self> {
479        let registry = resident_buffers();
480        let entry = registry.get(&id)?;
481        let upgraded = entry.value().upgrade();
482        drop(entry);
483        match upgraded {
484            Some(inner) => Some(Self { inner }),
485            None => {
486                registry.remove(&id);
487                None
488            }
489        }
490    }
491
492    /// Underlying `wgpu::Buffer`.
493    #[must_use]
494    pub fn buffer(&self) -> &wgpu::Buffer {
495        &self.inner.buffer
496    }
497
498    /// Clone the internal `Arc<wgpu::Buffer>`  -  cheap, reference-
499    /// count only. Used by the indirect dispatch path (C-B4) which
500    /// needs to stash the buffer alongside other args.
501    #[must_use]
502    pub fn buffer_arc(&self) -> Arc<wgpu::Buffer> {
503        Arc::clone(&self.inner.buffer)
504    }
505
506    /// Logical byte length requested by the caller.
507    #[must_use]
508    pub fn byte_len(&self) -> u64 {
509        self.inner.byte_len
510    }
511
512    /// Backing allocation length.
513    #[must_use]
514    pub fn allocation_len(&self) -> u64 {
515        self.inner.allocation_len
516    }
517
518    /// Logical element count. Byte buffers report one element per byte.
519    #[must_use]
520    pub fn element_count(&self) -> usize {
521        self.inner.element_count
522    }
523
524    /// Actual usage flags on the underlying GPU allocation.
525    #[must_use]
526    pub fn usage(&self) -> wgpu::BufferUsages {
527        self.inner.usage
528    }
529
530    pub(crate) fn from_parts(
531        buffer: Arc<wgpu::Buffer>,
532        byte_len: u64,
533        allocation_len: u64,
534        element_count: usize,
535        usage: wgpu::BufferUsages,
536        pool_return: Option<PoolReturn>,
537    ) -> Self {
538        let inner = Arc::new(GpuBufferInner {
539            id: NEXT_BUFFER_ID.fetch_add(1, Ordering::Relaxed),
540            buffer,
541            byte_len,
542            allocation_len,
543            element_count,
544            usage,
545            pool_return,
546        });
547        resident_buffers().insert(inner.id, Arc::downgrade(&inner));
548        Self { inner }
549    }
550}
551
552impl std::fmt::Debug for GpuBufferHandle {
553    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
554        formatter
555            .debug_struct("GpuBufferHandle")
556            .field("id", &self.id())
557            .field("byte_len", &self.byte_len())
558            .field("allocation_len", &self.allocation_len())
559            .field("element_count", &self.element_count())
560            .field("usage", &self.usage())
561            .finish()
562    }
563}
564
565impl Drop for GpuBufferInner {
566    fn drop(&mut self) {
567        resident_buffers().remove(&self.id);
568        if let Some(pool_return) = self.pool_return.take() {
569            pool_return.release(
570                Arc::clone(&self.buffer),
571                self.byte_len,
572                self.allocation_len,
573                self.usage,
574            );
575        }
576    }
577}
578
579pub(crate) fn aligned_len(len: usize) -> Result<u64, BackendError> {
580    let padded = aligned_len_usize(len, "GPU buffer length")?;
581    u64::try_from(padded).map_err(|source| {
582        BackendError::new(format!(
583            "GPU buffer length {padded} cannot fit u64: {source}. Fix: split the dispatch input."
584        ))
585    })
586}
587
588fn aligned_len_u64(len: u64, label: &'static str) -> Result<u64, BackendError> {
589    crate::numeric::align_up_u64(len, 4, label)
590}
591
592fn aligned_len_usize(len: usize, label: &'static str) -> Result<usize, BackendError> {
593    crate::numeric::align_up_usize(len, 4, label)
594}
595
596pub(crate) fn write_padded(
597    queue: &wgpu::Queue,
598    buffer: &wgpu::Buffer,
599    bytes: &[u8],
600    allocation_len: u64,
601) -> Result<(), BackendError> {
602    crate::padded_upload::write_padded_and_zero_fill(queue, buffer, bytes, allocation_len)
603}
604
605/// Default cap for the [`BindGroupCache`] LRU.
606const BIND_GROUP_CACHE_CAP: usize = 256;
607
608/// Inline storage for bind-group cache keys: typical shaders use few bindings;
609/// `SmallVec` avoids a heap `Vec` on most `get_or_create` calls.
610type BindGroupHandleKey = SmallVec<[u64; 16]>;
611
612/// Bounded LRU cache for wgpu bind groups, keyed by layout identity and
613/// the ordered set of buffer handles bound to that layout.
614///
615/// wgpu bind-group creation is non-trivial; this cache eliminates the
616/// redundant cost on repeated dispatches that share the same buffer
617/// handles.  Capped at 256 entries with LRU eviction to prevent
618/// descriptor-heap exhaustion on long-running servers.
619#[derive(Clone)]
620pub struct BindGroupCache {
621    cache: Arc<Mutex<BindGroupCacheInner>>,
622    hits: Arc<AtomicUsize>,
623    misses: Arc<AtomicUsize>,
624    evictions: Arc<AtomicUsize>,
625}
626
627struct BindGroupCacheInner {
628    entries: FxHashMap<BindGroupCacheKey, BindGroupCacheEntry>,
629    lru: BinaryHeap<Reverse<BindGroupLruEntry>>,
630    cap: usize,
631    next_generation: u64,
632}
633
634struct BindGroupCacheEntry {
635    bind_group: Arc<wgpu::BindGroup>,
636    last_seen: u64,
637}
638
639#[derive(Clone, Debug, Eq, PartialEq)]
640struct BindGroupLruEntry {
641    last_seen: u64,
642    key: BindGroupCacheKey,
643}
644
645impl Ord for BindGroupLruEntry {
646    fn cmp(&self, other: &Self) -> CmpOrdering {
647        self.last_seen
648            .cmp(&other.last_seen)
649            .then_with(|| self.key.cmp(&other.key))
650    }
651}
652
653impl PartialOrd for BindGroupLruEntry {
654    fn partial_cmp(&self, other: &Self) -> Option<CmpOrdering> {
655        Some(self.cmp(other))
656    }
657}
658
659fn push_bind_group_handle_key(key: &mut BindGroupHandleKey, handle: &GpuBufferHandle) -> bool {
660    key.push(handle.allocation_identity());
661    let Ok(aligned_len) = aligned_len_u64(handle.byte_len(), "bind-group handle key byte length")
662    else {
663        key.pop();
664        return false;
665    };
666    key.push(aligned_len);
667    true
668}
669
670#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
671struct BindGroupCacheKey {
672    layout_id: usize,
673    handles: BindGroupHandleKey,
674}
675
676impl std::fmt::Debug for BindGroupCache {
677    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
678        f.debug_struct("BindGroupCache")
679            .field("hits", &self.hits.load(Ordering::Relaxed))
680            .field("misses", &self.misses.load(Ordering::Relaxed))
681            .field("evictions", &self.evictions.load(Ordering::Relaxed))
682            .field("entries", &self.lock_cache().entries.len())
683            .finish_non_exhaustive()
684    }
685}
686
687impl Default for BindGroupCache {
688    fn default() -> Self {
689        Self::new()
690    }
691}
692
693impl BindGroupCacheInner {
694    fn next_lru_generation(&mut self) -> u64 {
695        let generation = self.next_generation;
696        self.next_generation = self.next_generation.wrapping_add(1);
697        generation
698    }
699
700    fn touch_existing(&mut self, key: &BindGroupCacheKey) -> Option<Arc<wgpu::BindGroup>> {
701        let generation = self.next_lru_generation();
702        let bind_group = {
703            let entry = self.entries.get_mut(key)?;
704            entry.last_seen = generation;
705            Arc::clone(&entry.bind_group)
706        };
707        self.lru.push(Reverse(BindGroupLruEntry {
708            last_seen: generation,
709            key: key.clone(),
710        }));
711        self.compact_lru_if_needed();
712        Some(bind_group)
713    }
714
715    fn insert_entry(&mut self, key: BindGroupCacheKey, bind_group: Arc<wgpu::BindGroup>) {
716        let generation = self.next_lru_generation();
717        self.entries.insert(
718            key.clone(),
719            BindGroupCacheEntry {
720                bind_group,
721                last_seen: generation,
722            },
723        );
724        self.lru.push(Reverse(BindGroupLruEntry {
725            last_seen: generation,
726            key,
727        }));
728        self.compact_lru_if_needed();
729    }
730
731    fn evict_to_cap(&mut self, mut on_evict: impl FnMut()) {
732        while self.entries.len() > self.cap {
733            let Some(key) = self.pop_lru_key() else { break };
734            if self.entries.remove(&key).is_some() {
735                on_evict();
736            }
737        }
738    }
739
740    fn pop_lru_key(&mut self) -> Option<BindGroupCacheKey> {
741        while let Some(Reverse(entry)) = self.lru.pop() {
742            if self
743                .entries
744                .get(&entry.key)
745                .is_some_and(|current| current.last_seen == entry.last_seen)
746            {
747                return Some(entry.key);
748            }
749        }
750        None
751    }
752
753    fn compact_lru_if_needed(&mut self) {
754        let live = self.entries.len();
755        if let Some(limit) = stale_lru_limit(live) {
756            if self.lru.len() <= limit {
757                return;
758            }
759        }
760        self.lru.clear();
761        self.lru.extend(self.entries.iter().map(|(key, entry)| {
762            Reverse(BindGroupLruEntry {
763                last_seen: entry.last_seen,
764                key: key.clone(),
765            })
766        }));
767    }
768}
769
770fn stale_lru_limit(live: usize) -> Option<usize> {
771    live.checked_mul(4).map(|limit| limit.max(8))
772}
773
774impl BindGroupCache {
775    fn lock_cache(&self) -> MutexGuard<'_, BindGroupCacheInner> {
776        self.cache.lock().unwrap_or_else(|error| {
777            tracing::error!(
778                "Vyre WGPU bind-group cache lock was poisoned: {error}. Fix: discard the cache after a panic; continuing with recovered state."
779            );
780            error.into_inner()
781        })
782    }
783
784    /// Create a bind-group cache with the default 256-entry cap.
785    #[must_use]
786    pub fn new() -> Self {
787        Self::with_cap(BIND_GROUP_CACHE_CAP)
788    }
789
790    /// Create with an explicit cap (used by tests and consumers that
791    /// want to size the LRU against known working-set bounds).
792    #[must_use]
793    pub fn with_cap(cap: usize) -> Self {
794        Self {
795            cache: Arc::new(Mutex::new(BindGroupCacheInner {
796                entries: FxHashMap::default(),
797                lru: BinaryHeap::new(),
798                cap: cap.max(1),
799                next_generation: 0,
800            })),
801            hits: Arc::new(AtomicUsize::new(0)),
802            misses: Arc::new(AtomicUsize::new(0)),
803            evictions: Arc::new(AtomicUsize::new(0)),
804        }
805    }
806
807    /// Return a cached bind group or create one with `factory`.
808    ///
809    /// `layout_id` must uniquely identify the `wgpu::BindGroupLayout`
810    /// (e.g. `Arc::as_ptr(layout).addr()`).
811    /// `handles` must be in the same order as the `wgpu::BindGroupEntry`
812    /// slice that the caller will pass to `create_bind_group` so that
813    /// identical handle sets map to the same cache key.
814    pub fn get_or_create(
815        &self,
816        layout_id: usize,
817        handles: &[GpuBufferHandle],
818        factory: impl FnOnce() -> wgpu::BindGroup,
819    ) -> Arc<wgpu::BindGroup> {
820        let Some(key_part_count) = handles.len().checked_mul(2) else {
821            self.misses.fetch_add(1, Ordering::Relaxed);
822            return Arc::new(factory());
823        };
824        let mut key_parts = SmallVec::with_capacity(key_part_count);
825        for handle in handles {
826            if !push_bind_group_handle_key(&mut key_parts, handle) {
827                self.misses.fetch_add(1, Ordering::Relaxed);
828                return Arc::new(factory());
829            }
830        }
831        self.get_or_create_by_ids(layout_id, key_parts, factory)
832    }
833
834    pub(crate) fn get_or_create_by_ids(
835        &self,
836        layout_id: usize,
837        handles: SmallVec<[u64; 16]>,
838        factory: impl FnOnce() -> wgpu::BindGroup,
839    ) -> Arc<wgpu::BindGroup> {
840        let key = BindGroupCacheKey { layout_id, handles };
841        {
842            let mut cache = self.lock_cache();
843            if let Some(existing) = cache.touch_existing(&key) {
844                self.hits.fetch_add(1, Ordering::Relaxed);
845                return existing;
846            }
847        }
848        let bg = Arc::new(factory());
849        let mut cache = self.lock_cache();
850        cache.insert_entry(key, Arc::clone(&bg));
851        cache.evict_to_cap(|| {
852            self.evictions.fetch_add(1, Ordering::Relaxed);
853        });
854        self.misses.fetch_add(1, Ordering::Relaxed);
855        bg
856    }
857
858    pub(crate) fn get_by_ids(
859        &self,
860        layout_id: usize,
861        handles: &[u64],
862    ) -> Option<Arc<wgpu::BindGroup>> {
863        let key = BindGroupCacheKey {
864            layout_id,
865            handles: SmallVec::from_slice(handles),
866        };
867        let mut cache = self.lock_cache();
868        let existing = cache.touch_existing(&key)?;
869        self.hits.fetch_add(1, Ordering::Relaxed);
870        Some(existing)
871    }
872
873    pub(crate) fn insert_by_ids(
874        &self,
875        layout_id: usize,
876        handles: &[u64],
877        bind_group: wgpu::BindGroup,
878    ) -> Arc<wgpu::BindGroup> {
879        let key = BindGroupCacheKey {
880            layout_id,
881            handles: SmallVec::from_slice(handles),
882        };
883        let mut cache = self.lock_cache();
884        if let Some(existing) = cache.touch_existing(&key) {
885            self.hits.fetch_add(1, Ordering::Relaxed);
886            return existing;
887        }
888        let bg = Arc::new(bind_group);
889        cache.insert_entry(key, Arc::clone(&bg));
890        cache.evict_to_cap(|| {
891            self.evictions.fetch_add(1, Ordering::Relaxed);
892        });
893        self.misses.fetch_add(1, Ordering::Relaxed);
894        bg
895    }
896
897    /// Return cache statistics for diagnostics and tests.
898    #[must_use]
899    pub fn stats(&self) -> BindGroupCacheStats {
900        BindGroupCacheStats {
901            hits: self.hits.load(Ordering::Relaxed),
902            misses: self.misses.load(Ordering::Relaxed),
903            evictions: self.evictions.load(Ordering::Relaxed),
904            entries: self.lock_cache().entries.len(),
905        }
906    }
907}
908
909/// Bind-group cache statistics for a compiled wgpu pipeline.
910#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
911pub struct BindGroupCacheStats {
912    /// Number of cached bind-group hits.
913    pub hits: usize,
914    /// Number of bind-group creations caused by cache misses.
915    pub misses: usize,
916    /// Number of cached bind-group entries evicted to honor the cap.
917    pub evictions: usize,
918    /// Current number of entries held.
919    pub entries: usize,
920}
921
922#[cfg(test)]
923mod tests {
924    use super::*;
925
926    /// StagingBufferPool must reuse buffers across readback calls so that 100
927    /// readbacks of the same size allocate only ~1 buffer.
928    #[test]
929    fn staging_pool_reuses_buffers_on_hot_readback_loop() {
930        let arc = crate::runtime::cached_device()
931            .expect("Fix: GPU device is required for staging pool test");
932        let (device, queue) = &*arc;
933
934        // Create a small COPY_SRC buffer with known contents.
935        let contents: Vec<u8> = vec![0xAB; 64];
936        let handle =
937            GpuBufferHandle::upload(device, queue, &contents, wgpu::BufferUsages::COPY_SRC)
938                .expect("Fix: upload should succeed");
939
940        let pool = StagingBufferPool::new();
941
942        for _ in 0..100 {
943            let mut out = Vec::new();
944            handle
945                .readback_until(device, Some(&pool), queue, &mut out, None)
946                .expect("Fix: pooled readback should succeed");
947            assert_eq!(out, contents, "readback bytes must match uploaded bytes");
948        }
949
950        let stats = pool.stats();
951        assert!(
952            stats.allocations <= 2,
953            "hot loop of 100 identical readbacks should allocate at most 2 staging buffers, got {} allocations and {} hits",
954            stats.allocations,
955            stats.hits
956        );
957    }
958
959    /// Without a pool, readback must still work and always create fresh buffers.
960    #[test]
961    fn readback_without_pool_always_allocates() {
962        let arc = crate::runtime::cached_device()
963            .expect("Fix: GPU device is required for readback regression test");
964        let (device, queue) = &*arc;
965
966        let contents: Vec<u8> = vec![0xCD; 32];
967        let handle =
968            GpuBufferHandle::upload(device, queue, &contents, wgpu::BufferUsages::COPY_SRC)
969                .expect("Fix: upload should succeed");
970
971        for _ in 0..5 {
972            let mut out = Vec::new();
973            handle
974                .readback(device, queue, &mut out)
975                .expect("Fix: unpooled readback should succeed");
976            assert_eq!(out, contents);
977        }
978    }
979
980    #[test]
981    fn resident_registry_handles_concurrent_lookup_and_drop() {
982        let arc = crate::runtime::cached_device()
983            .expect("Fix: GPU device is required for resident registry concurrency test");
984        let (device, queue) = &*arc;
985        let handle =
986            GpuBufferHandle::upload(device, queue, &[1, 2, 3, 4], wgpu::BufferUsages::COPY_SRC)
987                .expect("Fix: upload should register a resident buffer");
988        let id = handle.id();
989
990        // Phase 1: while the handle is alive, 8 concurrent readers
991        // must always resolve the resident id. Join BEFORE the drop so
992        // there is no readers-vs-drop race producing flaky panics.
993        let readers = (0..8)
994            .map(|_| {
995                std::thread::spawn(move || {
996                    for _ in 0..1_000 {
997                        let resident = GpuBufferHandle::from_resident_id(id)
998                            .expect("Fix: resident id must resolve while handle is alive");
999                        assert_eq!(resident.id(), id);
1000                    }
1001                })
1002            })
1003            .collect::<Vec<_>>();
1004        for reader in readers {
1005            reader
1006                .join()
1007                .expect("Fix: concurrent resident lookups must not panic");
1008        }
1009
1010        // Phase 2: dropping the handle must remove the id from the
1011        // registry so subsequent lookups return None.
1012        drop(handle);
1013        assert!(
1014            GpuBufferHandle::from_resident_id(id).is_none(),
1015            "dropped handles must be removed from the resident registry"
1016        );
1017    }
1018
1019    #[test]
1020    fn poisoned_staging_pool_lock_recovers_without_aborting_dispatch_path() {
1021        let pool = StagingBufferPool::new();
1022        let poisoned = pool.clone();
1023        let _ = std::thread::spawn(move || {
1024            let _guard = poisoned.lock_inner();
1025            panic!("poison staging buffer pool");
1026        })
1027        .join();
1028
1029        std::panic::catch_unwind(|| {
1030            let _ = pool.stats();
1031        })
1032        .expect("Fix: poisoned staging pool must recover so GPU readback pooling does not abort");
1033    }
1034
1035    #[test]
1036    fn poisoned_bind_group_cache_lock_recovers_without_aborting_dispatch_path() {
1037        let cache = BindGroupCache::new();
1038        let poisoned = cache.clone();
1039        let _ = std::thread::spawn(move || {
1040            let _guard = poisoned.lock_cache();
1041            panic!("poison bind group cache");
1042        })
1043        .join();
1044
1045        std::panic::catch_unwind(|| {
1046            let _ = cache.stats();
1047        })
1048        .expect("Fix: poisoned bind-group cache must recover so GPU dispatch does not abort");
1049    }
1050
1051    #[test]
1052    fn bind_group_cache_lru_heap_stays_capacity_scale() {
1053        let arc = crate::runtime::cached_device()
1054            .expect("Fix: GPU device is required for bind-group cache test");
1055        let (device, _) = &*arc;
1056        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
1057            label: Some("vyre bind-group cache lru test layout"),
1058            entries: &[wgpu::BindGroupLayoutEntry {
1059                binding: 0,
1060                visibility: wgpu::ShaderStages::COMPUTE,
1061                ty: wgpu::BindingType::Buffer {
1062                    ty: wgpu::BufferBindingType::Storage { read_only: true },
1063                    has_dynamic_offset: false,
1064                    min_binding_size: wgpu::BufferSize::new(4),
1065                },
1066                count: None,
1067            }],
1068        });
1069        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
1070            label: Some("vyre bind-group cache lru test buffer"),
1071            size: 4,
1072            usage: wgpu::BufferUsages::STORAGE,
1073            mapped_at_creation: false,
1074        });
1075        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
1076            label: Some("vyre bind-group cache lru test bind group"),
1077            layout: &layout,
1078            entries: &[wgpu::BindGroupEntry {
1079                binding: 0,
1080                resource: buffer.as_entire_binding(),
1081            }],
1082        });
1083        let cache = BindGroupCache::with_cap(4);
1084
1085        for i in 0..64u64 {
1086            cache.insert_by_ids(1, &[i, 4], bind_group.clone());
1087        }
1088
1089        let inner = cache.lock_cache();
1090        assert_eq!(inner.entries.len(), 4);
1091        assert!(
1092            inner.lru.len() <= inner.entries.len().saturating_mul(4).max(8),
1093            "Fix: bind-group LRU heap must compact stale entries to cache-capacity scale"
1094        );
1095    }
1096}
vyre_driver_wgpu/buffer/handle.rs

vyre_driver_wgpu/buffer/
handle.rs
