Skip to main content

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 fidget_core::render::ImageSize;
6use zerocopy::FromBytes;
7
8/// Handle around a growable GPU buffer
9///
10/// The buffer keeps track of both its current size and capacity (which may be
11/// larger).  It is used to prevent GPU buffer allocation churn.
12pub struct GenericFlexBuffer<T, B> {
13    /// Current size, which may be smaller than the buffer's capacity
14    size: B,
15    /// Actual GPU buffer
16    data: wgpu::Buffer,
17    /// Buffer label (to be used when reallocating)
18    name: String,
19    /// Marker for buffer tag type
20    _t: std::marker::PhantomData<T>,
21}
22
23/// Resizable array buffer
24pub type ArrayBuffer<T> = GenericFlexBuffer<T, usize>;
25
26/// Resizable image buffer
27pub type ImageBuffer<T> = GenericFlexBuffer<T, ImageSize>;
28
29/// Tag associated with a particular [`GenericFlexBuffer`]
30///
31/// The tag type serves two purposes:
32///
33/// - It declares the storage type and usage bits for the buffer
34/// - It makes buffers strongly typed, so that two buffers with equivalent
35///   storage and usage bits can be distinct types.
36pub trait BufferTag {
37    /// Data type stored in the buffer
38    type T;
39    /// Usage bits for the buffer
40    ///
41    /// This must be a union of [`wgpu::BufferUsages`] values
42    fn usage() -> u32;
43}
44
45/// Helper `struct` to make a mapped version of a storage buffer
46pub struct MappedBufferTag<T: BufferTag> {
47    _t: std::marker::PhantomData<T>,
48}
49impl<T: BufferTag> BufferTag for MappedBufferTag<T> {
50    type T = T::T;
51    fn usage() -> u32 {
52        wgpu::BufferUsages::COPY_DST.bits()
53            | wgpu::BufferUsages::MAP_READ.bits()
54    }
55}
56
57/// Helper macro to declare a buffer tag
58#[macro_export]
59macro_rules! tag {
60    ($vis:vis $name:ident,  $t:ty, $($flag:ident)|+ $(,$doc:expr)?) => {
61        $(#[doc = $doc])?
62        $vis struct $name;
63        impl $crate::buf::BufferTag for $name {
64            type T = $t;
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<T: BufferTag, B: BufferItemCount + Copy> GenericFlexBuffer<T, B> {
94    pub(crate) fn new(
95        device: &wgpu::Device,
96        name: String,
97        size: B,
98    ) -> Result<Self, BufferSizeError> {
99        Self::check_size(size)?;
100        let size_bytes = Self::calculate_buffer_size(size);
101        let usage = wgpu::BufferUsages::from_bits(T::usage()).unwrap();
102        let data = device.create_buffer(&wgpu::BufferDescriptor {
103            label: Some(name.as_str()),
104            size: size_bytes,
105            usage,
106            mapped_at_creation: false,
107        });
108        Ok(Self {
109            data,
110            size,
111            name,
112            _t: std::marker::PhantomData,
113        })
114    }
115
116    /// Calculate size from buffer item count
117    ///
118    /// Size is rounded up to the nearest multiple of 4 for alignment
119    fn calculate_buffer_size(item_count: B) -> u64 {
120        let out = u64::try_from(item_count.item_count())
121            .unwrap()
122            .checked_mul(u64::try_from(std::mem::size_of::<T::T>()).unwrap())
123            .unwrap();
124        out.next_multiple_of(4)
125    }
126
127    /// Returns the active buffer size (in bytes)
128    pub fn size_bytes(&self) -> u64 {
129        Self::calculate_buffer_size(self.size)
130    }
131
132    pub(crate) fn check_size(size: B) -> Result<(), BufferSizeError> {
133        let size = Self::calculate_buffer_size(size);
134        let usage = wgpu::BufferUsages::from_bits(T::usage()).unwrap();
135
136        let buf_ty = if usage.contains(wgpu::BufferUsages::STORAGE) {
137            BufferType::Storage
138        } else if usage.contains(wgpu::BufferUsages::UNIFORM) {
139            BufferType::Uniform
140        } else {
141            BufferType::Generic
142        };
143        buf_ty.check(size)
144    }
145
146    /// Grows the buffer to fit a particular size in bytes
147    ///
148    /// If the buffer already fits that size, then no allocation is performed,
149    /// but we always update the internal `item_count` (e.g. so that
150    /// [`bind_active`](Self::bind_active) returns the correct subset of the
151    /// buffer).
152    pub(crate) fn grow_to_fit(
153        &mut self,
154        device: &wgpu::Device,
155        size: B,
156    ) -> Result<(), BufferSizeError> {
157        Self::check_size(size)?;
158        let new_size = Self::calculate_buffer_size(size);
159        if new_size > self.capacity() {
160            let usage = self.data.usage();
161            self.data = device.create_buffer(&wgpu::BufferDescriptor {
162                label: Some(self.name.as_str()),
163                size: new_size,
164                usage,
165                mapped_at_creation: false,
166            });
167        }
168        self.size = size;
169        Ok(())
170    }
171
172    /// Returns a binding resource for the active slice of the buffer
173    pub fn bind_active(&self) -> wgpu::BindingResource<'_> {
174        self.data.slice(0..self.size_bytes()).into()
175    }
176
177    /// Returns the total buffer capacity (in bytes)
178    pub(crate) fn capacity(&self) -> u64 {
179        self.data.size()
180    }
181
182    /// Returns the buffer name
183    pub fn name(&self) -> &str {
184        &self.name
185    }
186
187    /// Maps the active portion of the buffer for reading
188    pub(crate) fn map_async(
189        &self,
190        callback: impl FnOnce(Result<(), wgpu::BufferAsyncError>)
191        + wgpu::WasmNotSend
192        + 'static,
193    ) -> wgpu::BufferSlice<'_> {
194        let slice = self.data.slice(0..self.size_bytes());
195        slice.map_async(wgpu::MapMode::Read, callback);
196        slice
197    }
198
199    /// Clears the active portion of the buffer
200    pub(crate) fn clear(&self, encoder: &mut wgpu::CommandEncoder) {
201        encoder.clear_buffer(&self.data, 0, Some(self.size_bytes()));
202    }
203
204    /// Returns a reference to the inner WGPU buffer
205    ///
206    /// Note that the whole buffer may not be active, since we allow for
207    /// oversized buffers!  Use [`size_bytes`](Self::size_bytes) to get the
208    /// active size, or [`bind_active`](Self::bind_active) to get a GPU binding.
209    pub fn data(&self) -> &wgpu::Buffer {
210        &self.data
211    }
212
213    /// Returns the size of the buffer (which is generic)
214    pub fn size(&self) -> B {
215        self.size
216    }
217}
218
219/// Buffer for reading data back from the GPU
220///
221/// Once mapped, this is wrapped by a [`MappedImage`]
222pub type ImageReadBuffer<T> = ImageBuffer<MappedBufferTag<T>>;
223
224/// Handle to a mapped [`ImageReadBuffer`], which unmaps the image when dropped
225pub struct MappedImage<'a, T: BufferTag> {
226    buf: &'a ImageReadBuffer<T>,
227    slice: wgpu::BufferSlice<'a>,
228}
229
230impl<T: BufferTag> Drop for MappedImage<'_, T> {
231    fn drop(&mut self) {
232        self.buf.data().unmap();
233    }
234}
235
236impl<'a, T: BufferTag> MappedImage<'a, T> {
237    /// Blocking function to build a new mapped image
238    pub fn map(
239        device: &wgpu::Device,
240        image: &'a mut ImageReadBuffer<T>,
241    ) -> Self {
242        let slice = image.map_async(|_| {});
243        device.poll(wgpu::PollType::wait_indefinitely()).unwrap();
244        MappedImage { buf: image, slice }
245    }
246
247    /// Returns the image's data
248    pub fn image(&self) -> fidget_raster::Image<u32, ImageSize> {
249        let result = <[u32]>::ref_from_bytes(&self.slice.get_mapped_range())
250            .unwrap()
251            .to_owned();
252        fidget_raster::Image::build(result, self.buf.size()).unwrap()
253    }
254}
255
256////////////////////////////////////////////////////////////////////////////////
257// Error handling zone!  This is perhaps a bit overengineered, but it meets the
258// desired behavior of function error types only containing errors that they can
259// actually return.
260
261/// Error type when resizing a buffer beyond its limit
262///
263/// We check against maximum buffer sizes (from the WebGPU spec) and return an
264/// error immediately, instead of deferring the error to the point where the
265/// buffer is used.
266#[derive(Debug, thiserror::Error)]
267pub enum BufferSizeError {
268    /// Buffer size is too large for the requested buffer usage
269    #[error(
270        "requested size {requested_size} exceeds maximum {} for \
271        {buffer_type} buffer",
272        buffer_type.max_size()
273    )]
274    TooLarge {
275        /// Size requested (in bytes)
276        requested_size: u64,
277        /// Buffer type (which determines the [max size](BufferType::max_size))
278        buffer_type: BufferType,
279    },
280}
281
282/// Buffer type for error reporting
283#[derive(Copy, Clone, Debug)]
284pub enum BufferType {
285    /// Uniform buffer ([`wgpu::BufferUsages::UNIFORM`])
286    Uniform,
287    /// Storage buffer ([`wgpu::BufferUsages::STORAGE`])
288    Storage,
289    /// Other buffer type (e.g. [`wgpu::BufferUsages::MAP_READ`])
290    Generic,
291}
292
293impl std::fmt::Display for BufferType {
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        let s = match self {
296            BufferType::Uniform => "uniform",
297            BufferType::Storage => "storage",
298            BufferType::Generic => "generic",
299        };
300        s.fmt(f)
301    }
302}
303
304impl BufferType {
305    /// Maximum size of this buffer type, per the WebGPU spec
306    pub const fn max_size(&self) -> u64 {
307        // These are copied from the spec, since we don't ask for anything extra
308        match self {
309            // maxUniformBufferBindingSize
310            BufferType::Uniform => 64 * 1024,
311            // maxStorageBufferBindingSize
312            BufferType::Storage => 128 * 1024 * 1024,
313            // maxBufferSize
314            BufferType::Generic => 256 * 1024 * 1024,
315        }
316    }
317
318    fn check(&self, requested_size: u64) -> Result<(), BufferSizeError> {
319        if requested_size > self.max_size() {
320            Err(BufferSizeError::TooLarge {
321                requested_size,
322                buffer_type: *self,
323            })
324        } else {
325            Ok(())
326        }
327    }
328}
329
330/// Helper function to make a uniform buffer binding
331pub(crate) fn buffer_uniform(binding: u32) -> wgpu::BindGroupLayoutEntry {
332    wgpu::BindGroupLayoutEntry {
333        binding,
334        visibility: wgpu::ShaderStages::COMPUTE,
335        ty: wgpu::BindingType::Buffer {
336            ty: wgpu::BufferBindingType::Uniform,
337            has_dynamic_offset: false,
338            min_binding_size: None,
339        },
340        count: None,
341    }
342}
343
344/// Helper function to make a read-only buffer binding
345pub(crate) fn buffer_ro(binding: u32) -> wgpu::BindGroupLayoutEntry {
346    wgpu::BindGroupLayoutEntry {
347        binding,
348        visibility: wgpu::ShaderStages::COMPUTE,
349        ty: wgpu::BindingType::Buffer {
350            ty: wgpu::BufferBindingType::Storage { read_only: true },
351            has_dynamic_offset: false,
352            min_binding_size: None,
353        },
354        count: None,
355    }
356}
357
358/// Helper function to make a read-only buffer binding with dynamic offset
359pub(crate) fn buffer_ro_dyn(binding: u32) -> wgpu::BindGroupLayoutEntry {
360    wgpu::BindGroupLayoutEntry {
361        binding,
362        visibility: wgpu::ShaderStages::COMPUTE,
363        ty: wgpu::BindingType::Buffer {
364            ty: wgpu::BufferBindingType::Storage { read_only: true },
365            has_dynamic_offset: true,
366            min_binding_size: None,
367        },
368        count: None,
369    }
370}
371
372/// Helper function to make a read-write buffer binding
373pub(crate) fn buffer_rw(binding: u32) -> wgpu::BindGroupLayoutEntry {
374    wgpu::BindGroupLayoutEntry {
375        binding,
376        visibility: wgpu::ShaderStages::COMPUTE,
377        ty: wgpu::BindingType::Buffer {
378            ty: wgpu::BufferBindingType::Storage { read_only: false },
379            has_dynamic_offset: false,
380            min_binding_size: None,
381        },
382        count: None,
383    }
384}