vyre_driver_wgpu/
device_buffer.rs

1//! WGPU concrete `DeviceBuffer` impl wrapping the existing
2//! `GpuBufferHandle`. Lets consumers allocate one persistent
3//! `Box<dyn DeviceBuffer>`, upload host bytes once, dispatch many
4//! times reusing the same device-resident allocation, and download
5//! when the loop ends  -  paying host↔device copy cost only at the
6//! boundary instead of per-dispatch.
7
8use crate::buffer::{write_padded, GpuBufferHandle};
9use crate::WgpuBackend;
10use vyre_driver::{BackendError, DeviceBuffer};
11
12/// Backend id string registered for `WgpuDeviceBuffer`. Matches
13/// `WgpuBackend::id()`.
14pub const WGPU_BACKEND_ID: &str = "wgpu";
15
16/// Concrete `DeviceBuffer` impl over a `GpuBufferHandle`.
17///
18/// Constructed via `WgpuBackend::allocate_device_buffer`. Public so
19/// downstream code can `downcast_ref::<WgpuDeviceBuffer>()` and
20/// reach the underlying `wgpu::Buffer` for advanced use; opaque to
21/// callers that only hold `Box<dyn DeviceBuffer>`.
22#[derive(Debug)]
23pub struct WgpuDeviceBuffer {
24    backend_id: &'static str,
25    handle: GpuBufferHandle,
26    logical_byte_len: usize,
27    label: Option<String>,
28}
29
30impl WgpuDeviceBuffer {
31    /// Borrow the underlying handle for advanced wgpu work (custom
32    /// bind groups, manual readback, etc.). Most callers should
33    /// stick to the `DeviceBuffer` API.
34    #[must_use]
35    pub fn handle(&self) -> &GpuBufferHandle {
36        &self.handle
37    }
38}
39
40impl DeviceBuffer for WgpuDeviceBuffer {
41    fn backend_id(&self) -> &'static str {
42        self.backend_id
43    }
44
45    fn byte_len(&self) -> usize {
46        self.logical_byte_len
47    }
48
49    fn debug_label(&self) -> Option<&str> {
50        self.label.as_deref()
51    }
52
53    fn as_any(&self) -> &dyn std::any::Any {
54        self
55    }
56
57    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
58        self
59    }
60}
61
62impl WgpuBackend {
63    /// Allocate a new GPU-resident buffer of `byte_len` bytes. The
64    /// buffer is created with STORAGE | COPY_SRC | COPY_DST so it can
65    /// participate in dispatch as either input or output and round-
66    /// trip through `upload_device_buffer` / `download_device_buffer`.
67    ///
68    /// # Errors
69    /// Returns a backend error if the underlying wgpu allocation
70    /// fails (e.g. byte_len exceeds device limits).
71    pub fn allocate_wgpu_device_buffer(
72        &self,
73        byte_len: usize,
74    ) -> Result<Box<dyn DeviceBuffer>, BackendError> {
75        let device_queue = self.current_device_queue();
76        let usage = wgpu::BufferUsages::STORAGE
77            | wgpu::BufferUsages::COPY_SRC
78            | wgpu::BufferUsages::COPY_DST;
79        let len = u64::try_from(byte_len).map_err(|_| {
80            BackendError::new(format!(
81                "Fix: WgpuBackend::allocate_device_buffer received byte_len {byte_len} that does not fit in u64."
82            ))
83        })?;
84        let handle = GpuBufferHandle::alloc(&device_queue.0, len, usage)?;
85        Ok(Box::new(WgpuDeviceBuffer {
86            backend_id: WGPU_BACKEND_ID,
87            handle,
88            logical_byte_len: byte_len,
89            label: None,
90        }))
91    }
92
93    /// Upload `bytes` into a previously-allocated wgpu DeviceBuffer.
94    /// Bytes shorter than the allocation are written at offset 0; the
95    /// remainder of the buffer is left as-is. Bytes longer than the
96    /// allocation are an error.
97    ///
98    /// # Errors
99    /// Returns a backend error when the buffer was not allocated by
100    /// this backend, when `bytes` exceeds the buffer's allocation, or
101    /// when the wgpu queue write fails.
102    pub fn upload_wgpu_device_buffer(
103        &self,
104        buffer: &mut dyn DeviceBuffer,
105        bytes: &[u8],
106    ) -> Result<(), BackendError> {
107        let backend_id = buffer.backend_id().to_string();
108        let wgpu_buf = buffer
109            .as_any_mut()
110            .downcast_mut::<WgpuDeviceBuffer>()
111            .ok_or_else(|| {
112                BackendError::new(format!(
113                    "Fix: upload_device_buffer expected a WgpuDeviceBuffer (allocated by `wgpu` backend) but got buffer owned by `{backend_id}`."
114                ))
115            })?;
116        let byte_len = u64::try_from(bytes.len()).map_err(|source| {
117            BackendError::new(format!(
118                "Fix: upload_device_buffer byte length cannot fit u64: {source}. Shard the upload before writing to a WGPU buffer."
119            ))
120        })?;
121        if byte_len > wgpu_buf.handle.byte_len() {
122            return Err(BackendError::new(format!(
123                "Fix: upload_device_buffer received {} bytes, exceeds logical length {} bytes for WgpuDeviceBuffer.",
124                bytes.len(),
125                wgpu_buf.handle.byte_len()
126            )));
127        }
128        let device_queue = self.current_device_queue();
129        write_padded(
130            &device_queue.1,
131            wgpu_buf.handle.buffer(),
132            bytes,
133            wgpu_buf.handle.allocation_len(),
134        )
135    }
136
137    /// Download the full byte_len of a previously-allocated wgpu
138    /// DeviceBuffer into a fresh `Vec<u8>`.
139    ///
140    /// # Errors
141    /// Returns a backend error when the buffer was not allocated by
142    /// this backend or when the readback fails (typically: buffer
143    /// missing COPY_SRC, which the standard allocator path includes).
144    pub fn download_wgpu_device_buffer(
145        &self,
146        buffer: &dyn DeviceBuffer,
147    ) -> Result<Vec<u8>, BackendError> {
148        let wgpu_buf = buffer
149            .as_any()
150            .downcast_ref::<WgpuDeviceBuffer>()
151            .ok_or_else(|| {
152                BackendError::new(format!(
153                    "Fix: download_device_buffer expected a WgpuDeviceBuffer (allocated by `wgpu` backend) but got buffer owned by `{}`.",
154                    buffer.backend_id()
155                ))
156        })?;
157        let device_queue = self.current_device_queue();
158        let mut out = Vec::new();
159        wgpu_buf
160            .handle
161            .readback(&device_queue.0, &device_queue.1, &mut out)?;
162        Ok(out)
163    }
164
165    /// Free a previously-allocated wgpu DeviceBuffer. The wgpu
166    /// allocation is released when the underlying Arc<wgpu::Buffer>
167    /// reaches zero references  -  dropping the box here is sufficient.
168    ///
169    /// # Errors
170    /// Returns a backend error when the buffer was not allocated by
171    /// this backend.
172    pub fn free_wgpu_device_buffer(
173        &self,
174        buffer: Box<dyn DeviceBuffer>,
175    ) -> Result<(), BackendError> {
176        let backend_id = buffer.backend_id().to_string();
177        if buffer.as_any().downcast_ref::<WgpuDeviceBuffer>().is_none() {
178            return Err(BackendError::new(format!(
179                "Fix: free_device_buffer expected a WgpuDeviceBuffer but got buffer owned by `{backend_id}`."
180            )));
181        }
182        drop(buffer);
183        Ok(())
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    fn required_backend() -> WgpuBackend {
192        WgpuBackend::new().unwrap_or_else(|error| {
193            panic!(
194                "Fix: WGPU device-buffer tests require a working GPU adapter on this fleet; repair adapter probing/driver configuration instead of silently skipping: {error}"
195            )
196        })
197    }
198
199    #[test]
200    fn allocate_round_trip_when_gpu_present() {
201        let backend = required_backend();
202        let mut buffer = backend.allocate_wgpu_device_buffer(64).expect(
203            "Fix: GPU-resident allocation should succeed for 64 bytes on a healthy adapter",
204        );
205        assert_eq!(buffer.backend_id(), WGPU_BACKEND_ID);
206        assert!(buffer.byte_len() >= 64);
207        let payload: Vec<u8> = (0..64).collect();
208        backend
209            .upload_wgpu_device_buffer(buffer.as_mut(), &payload)
210            .expect("Fix: upload must succeed for 64 bytes within allocation.");
211        let read = backend
212            .download_wgpu_device_buffer(buffer.as_ref())
213            .expect("Fix: download must succeed after upload on a freshly-allocated buffer.");
214        assert!(read.starts_with(&payload), "round-trip must preserve bytes");
215        backend
216            .free_wgpu_device_buffer(buffer)
217            .expect("Fix: free must accept a buffer the same backend allocated.");
218    }
219
220    #[test]
221    fn upload_rejects_oversize() {
222        let backend = required_backend();
223        let mut buffer = backend
224            .allocate_wgpu_device_buffer(16)
225            .expect("Fix: 16-byte allocation must succeed.");
226        let too_big = vec![0u8; 4096];
227        let err = backend
228            .upload_wgpu_device_buffer(buffer.as_mut(), &too_big)
229            .expect_err("Fix: oversize upload must error, not silently truncate.");
230        let msg = format!("{err}");
231        assert!(
232            msg.contains("exceeds logical length"),
233            "error must explain the size mismatch, got: {msg}"
234        );
235    }
236
237    #[test]
238    fn cross_backend_buffer_rejected() {
239        let backend = required_backend();
240        let mut alien = vyre_driver::HostShimBuffer::allocate("not-wgpu", 32);
241        let err = backend
242            .upload_wgpu_device_buffer(alien.as_mut(), &[1u8; 4])
243            .expect_err("Fix: WgpuBackend must reject buffers owned by other backends.");
244        let msg = format!("{err}");
245        assert!(
246            msg.contains("not-wgpu"),
247            "error must name the offending backend id, got: {msg}"
248        );
249    }
250
251    #[test]
252    fn upload_rejects_padding_overwrite_beyond_logical_length() {
253        let backend = required_backend();
254        let mut buffer = backend
255            .allocate_wgpu_device_buffer(17)
256            .expect("Fix: 17-byte allocation must succeed.");
257        let padding_overwrite = vec![0u8; 20];
258        let err = backend
259            .upload_wgpu_device_buffer(buffer.as_mut(), &padding_overwrite)
260            .expect_err("Fix: WgpuBackend must reject writes past the logical DeviceBuffer length even when the allocation is padded.");
261        let msg = format!("{err}");
262        assert!(
263            msg.contains("logical length 17"),
264            "error must explain the logical length boundary, got: {msg}"
265        );
266    }
267
268    #[test]
269    fn device_buffer_source_has_no_release_path_panic_or_padded_upload_boundary() {
270        let source = include_str!("device_buffer.rs");
271        let production = source
272            .split("#[cfg(test)]")
273            .next()
274            .expect("Fix: device buffer production source must precede tests");
275        assert!(
276            !production.contains(concat!("panic", "!("))
277                && !production.contains(".unwrap_or_else(")
278                && !production.contains("Vec::with_capacity"),
279            "Fix: WGPU DeviceBuffer release path must not panic or preallocate padded readback capacity."
280        );
281        assert!(
282            production.contains("logical_byte_len")
283                && production.contains("byte_len > wgpu_buf.handle.byte_len()")
284                && production.contains("let mut out = Vec::new();"),
285            "Fix: WGPU DeviceBuffer must preserve logical byte length and reject padded overwrite attempts."
286        );
287    }
288}
vyre_driver_wgpu/device_buffer.rs

vyre_driver_wgpu/
device_buffer.rs
