fidget_wgpu/
buf.rs

1//! Strongly-typed buffers
2//!
3//! This module is mostly internal to the crate, but is public because its types
4//! appear as return values and arguments.
5use crate::Gpu;
6use fidget_core::render::{ImageSize, VoxelSize};
7use fidget_raster::RenderSize;
8use zerocopy::{FromBytes, IntoBytes};
9
10/// Handle around a growable GPU buffer
11///
12/// The buffer keeps track of both its current size and capacity (which may be
13/// larger).  It is used to prevent GPU buffer allocation churn.
14pub struct FlexBuffer<T: BufferTag> {
15    /// Current size, which may be smaller than the buffer's capacity
16    size: T::S,
17    /// Actual GPU buffer
18    data: wgpu::Buffer,
19    /// Buffer label (to be used when reallocating)
20    name: String,
21    /// Marker for buffer tag type
22    _t: std::marker::PhantomData<T>,
23}
24
25/// Tag associated with a particular [`FlexBuffer`]
26///
27/// The tag type serves two purposes:
28///
29/// - It declares the storage type, usage bits, and size type for the buffer
30/// - It makes buffers strongly typed, so that two buffers with equivalent
31///   storage and usage bits can be distinct types.
32pub trait BufferTag {
33    /// Data type stored in the buffer
34    type T;
35    /// Size type
36    type S: BufferItemCount + Copy;
37    /// Usage bits for the buffer
38    ///
39    /// This must be a union of [`wgpu::BufferUsages`] values
40    fn usage() -> u32;
41}
42
43/// Helper `struct` to make a mapped version of a storage buffer
44pub struct MappedBufferTag<T: BufferTag> {
45    _t: std::marker::PhantomData<T>,
46}
47impl<T: BufferTag> BufferTag for MappedBufferTag<T> {
48    type T = T::T;
49    type S = T::S;
50    fn usage() -> u32 {
51        wgpu::BufferUsages::COPY_DST.bits()
52            | wgpu::BufferUsages::MAP_READ.bits()
53    }
54}
55
56/// Helper macro to declare a buffer tag
57#[macro_export]
58macro_rules! tag {
59    ($vis:vis $name:ident, $t:ty, $s:ty, $($flag:ident)|+ $(,$doc:expr)?) => {
60        $(#[doc = $doc])?
61        $vis struct $name;
62        impl $crate::buf::BufferTag for $name {
63            type T = $t;
64            type S = $s;
65            fn usage() -> u32 {
66                $( wgpu::BufferUsages::$flag.bits() )|+
67            }
68        }
69    }
70}
71
72/// Trait for types which have a certain number of items
73pub trait BufferItemCount {
74    /// The number of items
75    fn item_count(&self) -> usize;
76}
77
78impl BufferItemCount for usize {
79    fn item_count(&self) -> usize {
80        *self
81    }
82}
83
84impl BufferItemCount for ImageSize {
85    fn item_count(&self) -> usize {
86        usize::try_from(self.width())
87            .unwrap()
88            .checked_mul(usize::try_from(self.height()).unwrap())
89            .unwrap()
90    }
91}
92
93impl BufferItemCount for VoxelSize {
94    fn item_count(&self) -> usize {
95        ImageSize::from(*self).item_count()
96    }
97}
98
99/// Tag for a flex buffer which contains a config `C` followed by many `T`
100///
101/// Item count is reported in bytes
102pub(crate) struct FlexConfigSize<C, T> {
103    count: usize,
104    _c: std::marker::PhantomData<C>,
105    _t: std::marker::PhantomData<T>,
106}
107
108impl<C, T> From<usize> for FlexConfigSize<C, T> {
109    fn from(value: usize) -> Self {
110        Self {
111            count: value,
112            _c: std::marker::PhantomData,
113            _t: std::marker::PhantomData,
114        }
115    }
116}
117
118impl<C, T> Clone for FlexConfigSize<C, T> {
119    fn clone(&self) -> Self {
120        *self
121    }
122}
123impl<C, T> Copy for FlexConfigSize<C, T> {}
124
125impl<C, T> BufferItemCount for FlexConfigSize<C, T> {
126    fn item_count(&self) -> usize {
127        self.count * std::mem::size_of::<T>() + std::mem::size_of::<C>()
128    }
129}
130
131pub(crate) struct FlexConfigTag<C, T> {
132    _c: std::marker::PhantomData<C>,
133    _t: std::marker::PhantomData<T>,
134}
135
136impl<C, T> BufferTag for FlexConfigTag<C, T>
137where
138    T: Copy,
139{
140    type T = u8;
141    type S = FlexConfigSize<C, T>;
142    fn usage() -> u32 {
143        wgpu::BufferUsages::COPY_DST.bits() | wgpu::BufferUsages::STORAGE.bits()
144    }
145}
146
147/// Buffer which contains one config `C` followed by many `T`
148pub(crate) type FlexConfigBuffer<C, T> = FlexBuffer<FlexConfigTag<C, T>>;
149
150impl<C, T> FlexConfigBuffer<C, T>
151where
152    C: zerocopy::IntoBytes + zerocopy::Immutable,
153    T: zerocopy::IntoBytes + zerocopy::Immutable + Copy,
154{
155    pub(crate) fn write(
156        &mut self,
157        gpu: &Gpu,
158        c: &C,
159        data: &[T],
160    ) -> Result<ConfigBufferWrite, BufferSizeError> {
161        let r = self.grow_to_fit(&gpu.device, data.len().into())?;
162        let config_len = std::mem::size_of::<C>();
163        let data_len = std::mem::size_of_val(data);
164        let mut writer = gpu
165            .queue
166            .write_buffer_with(
167                &self.data,
168                0,
169                ((config_len + data_len) as u64).try_into().unwrap(),
170            )
171            .unwrap();
172        writer.slice(..config_len).copy_from_slice(c.as_bytes());
173        writer.slice(config_len..).copy_from_slice(data.as_bytes());
174        Ok(match r {
175            std::cmp::Ordering::Equal => ConfigBufferWrite::BufferUnchanged,
176            std::cmp::Ordering::Less | std::cmp::Ordering::Greater => {
177                ConfigBufferWrite::BufferChanged
178            }
179        })
180    }
181}
182
183#[derive(Copy, Clone, Debug)]
184#[must_use]
185pub(crate) enum ConfigBufferWrite {
186    BufferChanged,
187    BufferUnchanged,
188}
189
190impl<T: BufferTag> FlexBuffer<T> {
191    pub(crate) fn new(
192        device: &wgpu::Device,
193        name: impl AsRef<str>,
194        size: T::S,
195    ) -> Result<Self, BufferSizeError> {
196        Self::check_size(size)?;
197        let size_bytes = Self::calculate_buffer_size(size);
198        let usage = wgpu::BufferUsages::from_bits(T::usage()).unwrap();
199        let data = device.create_buffer(&wgpu::BufferDescriptor {
200            label: Some(name.as_ref()),
201            size: size_bytes,
202            usage,
203            mapped_at_creation: false,
204        });
205        Ok(Self {
206            data,
207            size,
208            name: name.as_ref().to_owned(),
209            _t: std::marker::PhantomData,
210        })
211    }
212
213    /// Calculate size from buffer item count
214    ///
215    /// Size is rounded up to the nearest multiple of 4 for alignment
216    pub fn calculate_buffer_size(item_count: T::S) -> u64 {
217        let out = u64::try_from(item_count.item_count())
218            .unwrap()
219            .checked_mul(u64::try_from(std::mem::size_of::<T::T>()).unwrap())
220            .unwrap();
221        out.next_multiple_of(4)
222    }
223
224    /// Returns the active buffer size (in bytes)
225    ///
226    /// Note that this is rounded up to the nearest multiple of 4, per
227    /// [`calculate_buffer_size`](Self::calculate_buffer_size)
228    pub fn size_bytes(&self) -> u64 {
229        Self::calculate_buffer_size(self.size)
230    }
231
232    pub(crate) fn check_size(size: T::S) -> Result<(), BufferSizeError> {
233        let size = Self::calculate_buffer_size(size);
234        let usage = wgpu::BufferUsages::from_bits(T::usage()).unwrap();
235
236        let buf_ty = if usage.contains(wgpu::BufferUsages::STORAGE) {
237            BufferType::Storage
238        } else if usage.contains(wgpu::BufferUsages::UNIFORM) {
239            BufferType::Uniform
240        } else {
241            BufferType::Generic
242        };
243        buf_ty.check(size)
244    }
245
246    /// Grows the buffer to fit a particular size in bytes
247    ///
248    /// If the buffer already fits that size, then no allocation is performed,
249    /// but we always update the internal `item_count` (e.g. so that
250    /// [`bind_active`](Self::bind_active) returns the correct subset of the
251    /// buffer).
252    ///
253    /// Returns a comparison between the previous item count and the new item
254    /// count (set by `size`).
255    pub(crate) fn grow_to_fit(
256        &mut self,
257        device: &wgpu::Device,
258        size: T::S,
259    ) -> Result<std::cmp::Ordering, BufferSizeError> {
260        Self::check_size(size)?;
261        let r = self.size.item_count().cmp(&size.item_count());
262        let new_size = Self::calculate_buffer_size(size);
263        if new_size > self.capacity() {
264            let usage = self.data.usage();
265            self.data = device.create_buffer(&wgpu::BufferDescriptor {
266                label: Some(self.name.as_str()),
267                size: new_size,
268                usage,
269                mapped_at_creation: false,
270            });
271        }
272        self.size = size;
273        Ok(r)
274    }
275
276    /// Returns a binding resource for the active slice of the buffer
277    pub fn bind_active(&self) -> wgpu::BindingResource<'_> {
278        self.data.slice(0..self.size_bytes()).try_into().unwrap()
279    }
280
281    /// Returns the total buffer capacity (in bytes)
282    pub(crate) fn capacity(&self) -> u64 {
283        self.data.size()
284    }
285
286    /// Returns the buffer name
287    pub fn name(&self) -> &str {
288        &self.name
289    }
290
291    /// Maps the active portion of the buffer for reading
292    pub(crate) fn map_async(
293        &self,
294        callback: impl FnOnce(Result<(), wgpu::BufferAsyncError>)
295        + wgpu::WasmNotSend
296        + 'static,
297    ) -> wgpu::BufferSlice<'_> {
298        let slice = self.data.slice(0..self.size_bytes());
299        slice.map_async(wgpu::MapMode::Read, callback);
300        slice
301    }
302
303    /// Clears the active portion of the buffer
304    pub(crate) fn clear(&self, encoder: &mut wgpu::CommandEncoder) {
305        encoder.clear_buffer(&self.data, 0, Some(self.size_bytes()));
306    }
307
308    /// Returns a reference to the inner WGPU buffer
309    ///
310    /// Note that the whole buffer may not be active, since we allow for
311    /// oversized buffers!  Use [`size_bytes`](Self::size_bytes) to get the
312    /// active size, or [`bind_active`](Self::bind_active) to get a GPU binding.
313    pub fn data(&self) -> &wgpu::Buffer {
314        &self.data
315    }
316
317    /// Returns the size of the buffer (which is generic)
318    pub fn size(&self) -> T::S {
319        self.size
320    }
321}
322
323/// Buffer for reading data back from the GPU
324///
325/// Once mapped, this can be wrapped by a [`MappedImage`] (if the tag's size
326/// parameter is an image size).
327pub type ReadBuffer<T> = FlexBuffer<MappedBufferTag<T>>;
328
329/// Handle to a mapped [`ReadBuffer`]
330///
331/// The buffer is automatically unmapped when this handle is dropped
332pub struct MappedBuffer<'a, T: BufferTag> {
333    buf: &'a ReadBuffer<T>,
334    slice: wgpu::BufferSlice<'a>,
335}
336
337impl<T: BufferTag> Drop for MappedBuffer<'_, T> {
338    fn drop(&mut self) {
339        self.buf.data().unmap();
340    }
341}
342
343impl<'a, T> MappedBuffer<'a, T>
344where
345    T: BufferTag,
346    T::T: zerocopy::FromBytes + zerocopy::Immutable + Copy,
347{
348    /// Basic constructor
349    pub(crate) fn new(
350        buf: &'a ReadBuffer<T>,
351        slice: wgpu::BufferSlice<'a>,
352    ) -> Self {
353        Self { buf, slice }
354    }
355
356    /// Copies the buffer data into a `Vec`
357    pub fn to_vec(&self) -> Vec<T::T> {
358        // We don't want to use the buffer's size, because it's rounded up to a
359        // multiple of 4; use the raw item size (in bytes) instead.
360        let slice = self.slice.get_mapped_range().unwrap();
361        let n = self.buf.size().item_count() * std::mem::size_of::<T::T>();
362        <[T::T]>::ref_from_bytes(&slice[..n]).unwrap().to_owned()
363    }
364}
365
366/// Handle to a mapped [`ReadBuffer`] containing an image
367///
368/// The buffer is automatically unmapped when this handle is dropped
369pub struct MappedImage<'a, T: BufferTag>(MappedBuffer<'a, T>)
370where
371    T::S: RenderSize;
372
373impl<'a, T> From<MappedBuffer<'a, T>> for MappedImage<'a, T>
374where
375    T: BufferTag,
376    T::S: RenderSize,
377{
378    fn from(value: MappedBuffer<'a, T>) -> Self {
379        Self(value)
380    }
381}
382
383impl<'a, T> MappedImage<'a, T>
384where
385    T: BufferTag,
386    T::T: zerocopy::FromBytes + zerocopy::Immutable + Copy,
387    T::S: RenderSize,
388{
389    /// Returns the image's data
390    pub fn image(&self) -> fidget_raster::Image<T::T, T::S> {
391        fidget_raster::Image::build(self.0.to_vec(), self.0.buf.size()).unwrap()
392    }
393}
394
395////////////////////////////////////////////////////////////////////////////////
396// Error handling zone!  This is perhaps a bit overengineered, but it meets the
397// desired behavior of function error types only containing errors that they can
398// actually return.
399
400/// Error type when resizing a buffer beyond its limit
401///
402/// We check against maximum buffer sizes (from the WebGPU spec) and return an
403/// error immediately, instead of deferring the error to the point where the
404/// buffer is used.
405#[derive(Debug, thiserror::Error)]
406pub enum BufferSizeError {
407    /// Buffer size is too large for the requested buffer usage
408    #[error(
409        "requested size {requested_size} exceeds maximum {} for \
410        {buffer_type} buffer",
411        buffer_type.max_size()
412    )]
413    TooLarge {
414        /// Size requested (in bytes)
415        requested_size: u64,
416        /// Buffer type (which determines the [max size](BufferType::max_size))
417        buffer_type: BufferType,
418    },
419}
420
421/// Buffer type for error reporting
422#[derive(Copy, Clone, Debug)]
423pub enum BufferType {
424    /// Uniform buffer ([`wgpu::BufferUsages::UNIFORM`])
425    Uniform,
426    /// Storage buffer ([`wgpu::BufferUsages::STORAGE`])
427    Storage,
428    /// Other buffer type (e.g. [`wgpu::BufferUsages::MAP_READ`])
429    Generic,
430}
431
432impl std::fmt::Display for BufferType {
433    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
434        let s = match self {
435            BufferType::Uniform => "uniform",
436            BufferType::Storage => "storage",
437            BufferType::Generic => "generic",
438        };
439        s.fmt(f)
440    }
441}
442
443impl BufferType {
444    /// Maximum size of this buffer type, per the WebGPU spec
445    pub const fn max_size(&self) -> u64 {
446        // These are copied from the spec, since we don't ask for anything extra
447        match self {
448            // maxUniformBufferBindingSize
449            BufferType::Uniform => 64 * 1024,
450            // maxStorageBufferBindingSize
451            BufferType::Storage => 128 * 1024 * 1024,
452            // maxBufferSize
453            BufferType::Generic => 256 * 1024 * 1024,
454        }
455    }
456
457    fn check(&self, requested_size: u64) -> Result<(), BufferSizeError> {
458        if requested_size > self.max_size() {
459            Err(BufferSizeError::TooLarge {
460                requested_size,
461                buffer_type: *self,
462            })
463        } else {
464            Ok(())
465        }
466    }
467}
468
469/// Helper function to make a uniform buffer binding
470pub(crate) fn buffer_uniform(binding: u32) -> wgpu::BindGroupLayoutEntry {
471    wgpu::BindGroupLayoutEntry {
472        binding,
473        visibility: wgpu::ShaderStages::COMPUTE,
474        ty: wgpu::BindingType::Buffer {
475            ty: wgpu::BufferBindingType::Uniform,
476            has_dynamic_offset: false,
477            min_binding_size: None,
478        },
479        count: None,
480    }
481}
482
483/// Helper function to make a read-only buffer binding
484pub(crate) fn buffer_ro(binding: u32) -> wgpu::BindGroupLayoutEntry {
485    wgpu::BindGroupLayoutEntry {
486        binding,
487        visibility: wgpu::ShaderStages::COMPUTE,
488        ty: wgpu::BindingType::Buffer {
489            ty: wgpu::BufferBindingType::Storage { read_only: true },
490            has_dynamic_offset: false,
491            min_binding_size: None,
492        },
493        count: None,
494    }
495}
496
497/// Helper function to make a read-only buffer binding with dynamic offset
498pub(crate) fn buffer_ro_dyn(binding: u32) -> wgpu::BindGroupLayoutEntry {
499    wgpu::BindGroupLayoutEntry {
500        binding,
501        visibility: wgpu::ShaderStages::COMPUTE,
502        ty: wgpu::BindingType::Buffer {
503            ty: wgpu::BufferBindingType::Storage { read_only: true },
504            has_dynamic_offset: true,
505            min_binding_size: None,
506        },
507        count: None,
508    }
509}
510
511/// Helper function to make a read-write buffer binding
512pub(crate) fn buffer_rw(binding: u32) -> wgpu::BindGroupLayoutEntry {
513    wgpu::BindGroupLayoutEntry {
514        binding,
515        visibility: wgpu::ShaderStages::COMPUTE,
516        ty: wgpu::BindingType::Buffer {
517            ty: wgpu::BufferBindingType::Storage { read_only: false },
518            has_dynamic_offset: false,
519            min_binding_size: None,
520        },
521        count: None,
522    }
523}
fidget_wgpu/buf.rs

fidget_wgpu/
buf.rs
