Skip to main content

bevy_render/render_resource/
uniform_buffer.rs

1use core::{marker::PhantomData, num::NonZero};
2
3use crate::{
4    render_resource::{make_buffer_label, Buffer},
5    renderer::{RenderDevice, RenderQueue},
6};
7use encase::{
8    internal::{AlignmentValue, BufferMut, WriteInto},
9    DynamicUniformBuffer as DynamicUniformBufferWrapper, ShaderType,
10    UniformBuffer as UniformBufferWrapper,
11};
12use wgpu::{
13    util::BufferInitDescriptor, BindingResource, BufferBinding, BufferDescriptor, BufferUsages,
14};
15
16use super::IntoBinding;
17
18/// Stores data to be transferred to the GPU and made accessible to shaders as a uniform buffer.
19///
20/// Uniform buffers are available to shaders on a read-only basis. Uniform buffers are commonly used to make available to shaders
21/// parameters that are constant during shader execution, and are best used for data that is relatively small in size as they are
22/// only guaranteed to support up to 16kB per binding.
23///
24/// The contained data is stored in system RAM. [`write_buffer`](UniformBuffer::write_buffer) queues
25/// copying of the data from system RAM to VRAM. Data in uniform buffers must follow [std140 alignment/padding requirements],
26/// which is automatically enforced by this structure. Per the WGPU spec, uniform buffers cannot store runtime-sized array
27/// (vectors), or structures with fields that are vectors.
28///
29/// Other options for storing GPU-accessible data are:
30/// * [`BufferVec`](crate::render_resource::BufferVec)
31/// * [`DynamicStorageBuffer`](crate::render_resource::DynamicStorageBuffer)
32/// * [`DynamicUniformBuffer`]
33/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
34/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
35/// * [`StorageBuffer`](crate::render_resource::StorageBuffer)
36/// * [`Texture`](crate::render_resource::Texture)
37///
38/// [std140 alignment/padding requirements]: https://www.w3.org/TR/WGSL/#address-spaces-uniform
39pub struct UniformBuffer<T: ShaderType> {
40    value: T,
41    scratch: UniformBufferWrapper<Vec<u8>>,
42    buffer: Option<Buffer>,
43    label: Option<String>,
44    changed: bool,
45    buffer_usage: BufferUsages,
46}
47
48impl<T: ShaderType> From<T> for UniformBuffer<T> {
49    fn from(value: T) -> Self {
50        Self {
51            value,
52            scratch: UniformBufferWrapper::new(Vec::new()),
53            buffer: None,
54            label: None,
55            changed: false,
56            buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
57        }
58    }
59}
60
61impl<T: ShaderType + Default> Default for UniformBuffer<T> {
62    fn default() -> Self {
63        Self {
64            value: T::default(),
65            scratch: UniformBufferWrapper::new(Vec::new()),
66            buffer: None,
67            label: None,
68            changed: false,
69            buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
70        }
71    }
72}
73
74impl<T: ShaderType + WriteInto> UniformBuffer<T> {
75    #[inline]
76    pub fn buffer(&self) -> Option<&Buffer> {
77        self.buffer.as_ref()
78    }
79
80    #[inline]
81    pub fn binding(&self) -> Option<BindingResource<'_>> {
82        Some(BindingResource::Buffer(
83            self.buffer()?.as_entire_buffer_binding(),
84        ))
85    }
86
87    /// Set the data the buffer stores.
88    pub fn set(&mut self, value: T) {
89        self.value = value;
90    }
91
92    pub fn get(&self) -> &T {
93        &self.value
94    }
95
96    pub fn get_mut(&mut self) -> &mut T {
97        &mut self.value
98    }
99
100    pub fn set_label(&mut self, label: Option<&str>) {
101        let label = label.map(str::to_string);
102
103        if label != self.label {
104            self.changed = true;
105        }
106
107        self.label = label;
108    }
109
110    pub fn get_label(&self) -> Option<&str> {
111        self.label.as_deref()
112    }
113
114    /// Add more [`BufferUsages`] to the buffer.
115    ///
116    /// This method only allows addition of flags to the default usage flags.
117    ///
118    /// The default values for buffer usage are `BufferUsages::COPY_DST` and `BufferUsages::UNIFORM`.
119    pub fn add_usages(&mut self, usage: BufferUsages) {
120        self.buffer_usage |= usage;
121        self.changed = true;
122    }
123
124    /// Queues writing of data from system RAM to VRAM using the [`RenderDevice`]
125    /// and the provided [`RenderQueue`], if a GPU-side backing buffer already exists.
126    ///
127    /// If a GPU-side buffer does not already exist for this data, such a buffer is initialized with currently
128    /// available data.
129    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
130        self.scratch.write(&self.value).unwrap();
131
132        if self.changed || self.buffer.is_none() {
133            self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
134                label: make_buffer_label::<Self>(&self.label),
135                usage: self.buffer_usage,
136                contents: self.scratch.as_ref(),
137            }));
138            self.changed = false;
139        } else if let Some(buffer) = &self.buffer {
140            queue.write_buffer(buffer, 0, self.scratch.as_ref());
141        }
142    }
143}
144
145impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a UniformBuffer<T> {
146    #[inline]
147    fn into_binding(self) -> BindingResource<'a> {
148        self.buffer()
149            .expect("Failed to get buffer")
150            .as_entire_buffer_binding()
151            .into_binding()
152    }
153}
154
155/// Stores data to be transferred to the GPU and made accessible to shaders as a dynamic uniform buffer.
156///
157/// Dynamic uniform buffers are available to shaders on a read-only basis. Dynamic uniform buffers are commonly used to make
158/// available to shaders runtime-sized arrays of parameters that are otherwise constant during shader execution, and are best
159/// suited to data that is relatively small in size as they are only guaranteed to support up to 16kB per binding.
160///
161/// The contained data is stored in system RAM. [`write_buffer`](DynamicUniformBuffer::write_buffer) queues
162/// copying of the data from system RAM to VRAM. Data in uniform buffers must follow [std140 alignment/padding requirements],
163/// which is automatically enforced by this structure. Per the WGPU spec, uniform buffers cannot store runtime-sized array
164/// (vectors), or structures with fields that are vectors.
165///
166/// Other options for storing GPU-accessible data are:
167/// * [`BufferVec`](crate::render_resource::BufferVec)
168/// * [`DynamicStorageBuffer`](crate::render_resource::DynamicStorageBuffer)
169/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
170/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
171/// * [`StorageBuffer`](crate::render_resource::StorageBuffer)
172/// * [`Texture`](crate::render_resource::Texture)
173/// * [`UniformBuffer`]
174///
175/// [std140 alignment/padding requirements]: https://www.w3.org/TR/WGSL/#address-spaces-uniform
176pub struct DynamicUniformBuffer<T: ShaderType> {
177    scratch: DynamicUniformBufferWrapper<Vec<u8>>,
178    buffer: Option<Buffer>,
179    label: Option<String>,
180    changed: bool,
181    buffer_usage: BufferUsages,
182    _marker: PhantomData<fn() -> T>,
183}
184
185impl<T: ShaderType> Default for DynamicUniformBuffer<T> {
186    fn default() -> Self {
187        Self {
188            scratch: DynamicUniformBufferWrapper::new(Vec::new()),
189            buffer: None,
190            label: None,
191            changed: false,
192            buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
193            _marker: PhantomData,
194        }
195    }
196}
197
198impl<T: ShaderType> DynamicUniformBuffer<T> {
199    pub fn new_with_alignment(alignment: u64) -> Self {
200        Self {
201            scratch: DynamicUniformBufferWrapper::new_with_alignment(Vec::new(), alignment),
202            buffer: None,
203            label: None,
204            changed: false,
205            buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
206            _marker: PhantomData,
207        }
208    }
209
210    #[inline]
211    pub fn buffer(&self) -> Option<&Buffer> {
212        self.buffer.as_ref()
213    }
214
215    #[inline]
216    pub fn binding(&self) -> Option<BindingResource<'_>> {
217        Some(BindingResource::Buffer(BufferBinding {
218            buffer: self.buffer()?,
219            offset: 0,
220            size: Some(T::min_size()),
221        }))
222    }
223
224    #[inline]
225    pub fn is_empty(&self) -> bool {
226        self.scratch.as_ref().is_empty()
227    }
228
229    pub fn set_label(&mut self, label: Option<&str>) {
230        let label = label.map(str::to_string);
231
232        if label != self.label {
233            self.changed = true;
234        }
235
236        self.label = label;
237    }
238
239    pub fn get_label(&self) -> Option<&str> {
240        self.label.as_deref()
241    }
242
243    /// Add more [`BufferUsages`] to the buffer.
244    ///
245    /// This method only allows addition of flags to the default usage flags.
246    ///
247    /// The default values for buffer usage are `BufferUsages::COPY_DST` and `BufferUsages::UNIFORM`.
248    pub fn add_usages(&mut self, usage: BufferUsages) {
249        self.buffer_usage |= usage;
250        self.changed = true;
251    }
252
253    /// Creates a writer that can be used to directly write elements into the target buffer.
254    ///
255    /// This method uses less memory and performs fewer memory copies using over [`push`] and [`write_buffer`].
256    ///
257    /// `max_count` *must* be greater than or equal to the number of elements that are to be written to the buffer, or
258    /// the writer will panic while writing.  Dropping the writer will schedule the buffer write into the provided
259    /// [`RenderQueue`].
260    ///
261    /// If there is no GPU-side buffer allocated to hold the data currently stored, or if a GPU-side buffer previously
262    /// allocated does not have enough capacity to hold `max_count` elements, a new GPU-side buffer is created.
263    ///
264    /// Returns `None` if there is no allocated GPU-side buffer, and `max_count` is 0.
265    ///
266    /// [`push`]: Self::push
267    /// [`write_buffer`]: Self::write_buffer
268    #[inline]
269    pub fn get_writer<'a>(
270        &'a mut self,
271        max_count: usize,
272        device: &RenderDevice,
273        queue: &'a RenderQueue,
274    ) -> Option<DynamicUniformBufferWriter<T>> {
275        let alignment = if cfg!(target_abi = "sim") {
276            // On iOS simulator on silicon macs, metal validation check that the host OS alignment
277            // is respected, but the device reports the correct value for iOS, which is smaller.
278            // Use the larger value.
279            // See https://github.com/gfx-rs/wgpu/issues/7057 - remove if it's not needed anymore.
280            AlignmentValue::new(256)
281        } else {
282            AlignmentValue::new(device.limits().min_uniform_buffer_offset_alignment as u64)
283        };
284
285        let mut capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
286        let size = alignment
287            .round_up(T::min_size().get())
288            .checked_mul(max_count as u64)
289            .unwrap();
290
291        if capacity < size || (self.changed && size > 0) {
292            let buffer = device.create_buffer(&BufferDescriptor {
293                label: make_buffer_label::<Self>(&self.label),
294                usage: self.buffer_usage,
295                size,
296                mapped_at_creation: false,
297            });
298            capacity = buffer.size();
299            self.buffer = Some(buffer);
300            self.changed = false;
301        }
302
303        if let Some(buffer) = self.buffer.as_deref() {
304            let buffer_view =
305                queue.write_buffer_with(buffer, 0, NonZero::<u64>::new(buffer.size())?)?;
306            Some(DynamicUniformBufferWriter {
307                buffer: encase::DynamicUniformBuffer::new_with_alignment(
308                    QueueWriteBufferViewWrapper {
309                        capacity: capacity as usize,
310                        buffer_view,
311                    },
312                    alignment.get(),
313                ),
314                _marker: PhantomData,
315            })
316        } else {
317            None
318        }
319    }
320
321    /// Queues writing of data from system RAM to VRAM using the [`RenderDevice`]
322    /// and the provided [`RenderQueue`].
323    ///
324    /// If there is no GPU-side buffer allocated to hold the data currently stored, or if a GPU-side buffer previously
325    /// allocated does not have enough capacity, a new GPU-side buffer is created.
326    #[inline]
327    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
328        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
329        let size = self.scratch.as_ref().len() as u64;
330
331        if capacity < size || (self.changed && size > 0) {
332            self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
333                label: self.label.as_deref(),
334                usage: self.buffer_usage,
335                contents: self.scratch.as_ref(),
336            }));
337            self.changed = false;
338        } else if let Some(buffer) = &self.buffer {
339            queue.write_buffer(buffer, 0, self.scratch.as_ref());
340        }
341    }
342
343    #[inline]
344    pub fn clear(&mut self) {
345        self.scratch.as_mut().clear();
346        self.scratch.set_offset(0);
347    }
348}
349
350impl<T: ShaderType + WriteInto> DynamicUniformBuffer<T> {
351    /// Push data into the `DynamicUniformBuffer`'s internal vector (residing on system RAM).
352    #[inline]
353    pub fn push(&mut self, value: &T) -> u32 {
354        self.scratch.write(value).unwrap() as u32
355    }
356}
357
358/// A writer that can be used to directly write elements into the target buffer.
359///
360/// For more information, see [`DynamicUniformBuffer::get_writer`].
361pub struct DynamicUniformBufferWriter<T> {
362    buffer: encase::DynamicUniformBuffer<QueueWriteBufferViewWrapper>,
363    _marker: PhantomData<fn() -> T>,
364}
365
366impl<T: ShaderType + WriteInto> DynamicUniformBufferWriter<T> {
367    pub fn write(&mut self, value: &T) -> u32 {
368        self.buffer.write(value).unwrap() as u32
369    }
370}
371
372/// A wrapper to work around the orphan rule so that [`wgpu::QueueWriteBufferView`] can implement
373/// [`BufferMut`].
374struct QueueWriteBufferViewWrapper {
375    buffer_view: wgpu::QueueWriteBufferView,
376    // Must be kept separately and cannot be retrieved from buffer_view, as the read-only access will
377    // invoke a panic.
378    capacity: usize,
379}
380
381impl BufferMut for QueueWriteBufferViewWrapper {
382    #[inline]
383    fn capacity(&self) -> usize {
384        self.capacity
385    }
386
387    #[inline]
388    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
389        self.buffer_view
390            .slice(offset..offset + val.len())
391            .copy_from_slice(val);
392    }
393
394    #[inline]
395    fn write_slice(&mut self, offset: usize, val: &[u8]) {
396        self.buffer_view
397            .slice(offset..offset + val.len())
398            .copy_from_slice(val);
399    }
400}
401
402impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a DynamicUniformBuffer<T> {
403    #[inline]
404    fn into_binding(self) -> BindingResource<'a> {
405        self.binding().unwrap()
406    }
407}