From 075a3080616f2c4836fee13e4772b3f6642e45dc Mon Sep 17 00:00:00 2001 From: Robert Swain Date: Wed, 2 Mar 2022 07:44:07 +0100 Subject: [PATCH 1/4] bevy_render: Add StorageBuffer helper for storage buffer bindings StorageBuffer takes two types, the first is the type of the body of the shader struct, excluding any final variable-sized array. The second is the type of the variable-sized array items. --- crates/bevy_render/src/render_resource/mod.rs | 2 + .../src/render_resource/storage_buffer.rs | 145 ++++++++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 crates/bevy_render/src/render_resource/storage_buffer.rs diff --git a/crates/bevy_render/src/render_resource/mod.rs b/crates/bevy_render/src/render_resource/mod.rs index 8089199f790daf..fd606116859576 100644 --- a/crates/bevy_render/src/render_resource/mod.rs +++ b/crates/bevy_render/src/render_resource/mod.rs @@ -6,6 +6,7 @@ mod pipeline; mod pipeline_cache; mod pipeline_specializer; mod shader; +mod storage_buffer; mod texture; mod uniform_vec; @@ -17,6 +18,7 @@ pub use pipeline::*; pub use pipeline_cache::*; pub use pipeline_specializer::*; pub use shader::*; +pub use storage_buffer::*; pub use texture::*; pub use uniform_vec::*; diff --git a/crates/bevy_render/src/render_resource/storage_buffer.rs b/crates/bevy_render/src/render_resource/storage_buffer.rs new file mode 100644 index 00000000000000..b3d00c923a2889 --- /dev/null +++ b/crates/bevy_render/src/render_resource/storage_buffer.rs @@ -0,0 +1,145 @@ +use std::num::NonZeroU64; + +use bevy_crevice::std430::{self, AsStd430, Std430}; +use bevy_utils::tracing::warn; +use wgpu::{BindingResource, BufferBinding, BufferDescriptor, BufferUsages}; + +use crate::renderer::{RenderDevice, RenderQueue}; + +use super::Buffer; + +/// A helper for a storage buffer binding with a body, or a variable-sized array, or both. +pub struct StorageBuffer { + body: T, + values: Vec, + scratch: Vec, + storage_buffer: Option, + capacity: usize, + item_size: usize, +} + +impl Default for StorageBuffer { + fn default() -> Self { + Self { + body: T::default(), + values: Vec::new(), + scratch: Vec::new(), + storage_buffer: None, + capacity: 0, + item_size: U::std430_size_static(), + } + } +} + +impl StorageBuffer { + #[inline] + pub fn storage_buffer(&self) -> Option<&Buffer> { + self.storage_buffer.as_ref() + } + + #[inline] + pub fn binding(&self) -> Option { + Some(BindingResource::Buffer(BufferBinding { + buffer: self.storage_buffer()?, + offset: 0, + size: Some(NonZeroU64::new((self.size()) as u64).unwrap()), + })) + } + + #[inline] + pub fn set_body(&mut self, body: T) { + self.body = body; + } + + #[inline] + pub fn len(&self) -> usize { + self.values.len() + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + #[inline] + pub fn capacity(&self) -> usize { + self.capacity + } + + pub fn push(&mut self, value: U) -> usize { + let index = self.values.len(); + self.values.push(value); + index + } + + pub fn get_mut(&mut self, index: usize) -> &mut U { + &mut self.values[index] + } + + pub fn reserve(&mut self, capacity: usize, device: &RenderDevice) -> bool { + if self.storage_buffer.is_none() || capacity > self.capacity { + self.capacity = capacity; + let size = self.size(); + self.scratch.resize(size, 0); + self.storage_buffer = Some(device.create_buffer(&BufferDescriptor { + label: None, + size: size as wgpu::BufferAddress, + usage: BufferUsages::COPY_DST | BufferUsages::STORAGE, + mapped_at_creation: false, + })); + true + } else { + false + } + } + + fn size(&self) -> usize { + let mut size = 0; + size += T::std430_size_static(); + if size > 0 { + // Pad according to the array item type's alignment + size = (size + ::Output::ALIGNMENT - 1) + & !(::Output::ALIGNMENT - 1); + } + // Variable size arrays must have at least 1 element + size += self.item_size * self.capacity.max(1); + size + } + + pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) { + self.reserve(self.values.len(), device); + if let Some(storage_buffer) = &self.storage_buffer { + let range = 0..self.size(); + let mut writer = std430::Writer::new(&mut self.scratch[range.clone()]); + let mut offset = 0; + // First write the struct body if there is one + if T::std430_size_static() > 0 { + if let Ok(new_offset) = writer.write(&self.body).map_err(|e| warn!("{:?}", e)) { + offset = new_offset; + } + } + if self.values.is_empty() { + for i in offset..self.size() { + self.scratch[i] = 0; + } + } else { + // Then write the array. Note that padding bytes may be added between the body + // and the array in order to align the array to the alignment requirements of its + // items + writer + .write(self.values.as_slice()) + .map_err(|e| warn!("{:?}", e)) + .ok(); + } + queue.write_buffer(storage_buffer, 0, &self.scratch[range]); + } + } + + pub fn clear(&mut self) { + self.values.clear(); + } + + pub fn values(&self) -> &[U] { + &self.values + } +} From 3237758cedeaf671e74c617ab374e2386431a72b Mon Sep 17 00:00:00 2001 From: Robert Swain Date: Wed, 2 Mar 2022 07:49:54 +0100 Subject: [PATCH 2/4] bevy_crevice: Add impl Std430 for () This allows specification of the empty () type for StorageBuffer bodies or variable-sized arrays. --- crates/bevy_crevice/src/std430/traits.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/crates/bevy_crevice/src/std430/traits.rs b/crates/bevy_crevice/src/std430/traits.rs index 7f2967f3b421a5..04f9f526d7f3a8 100644 --- a/crates/bevy_crevice/src/std430/traits.rs +++ b/crates/bevy_crevice/src/std430/traits.rs @@ -39,6 +39,14 @@ pub unsafe trait Std430: Copy + Zeroable + Pod { } } +unsafe impl Std430 for () { + const ALIGNMENT: usize = 0; + + const PAD_AT_END: bool = false; + + type Padded = (); +} + /// Trait specifically for Std430::Padded, implements conversions between padded type and base type. pub trait Std430Convertible: Copy { /// Convert from self to Std430 From 952819acbfd0ed1bd1cd84fa928ffd1eb5ec1842 Mon Sep 17 00:00:00 2001 From: Robert Swain Date: Wed, 2 Mar 2022 08:02:31 +0100 Subject: [PATCH 3/4] bevy_render: Handle whether the storage buffer has a variable-sized array --- .../src/render_resource/storage_buffer.rs | 39 +++++++++++-------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/crates/bevy_render/src/render_resource/storage_buffer.rs b/crates/bevy_render/src/render_resource/storage_buffer.rs index b3d00c923a2889..f750ba879faeda 100644 --- a/crates/bevy_render/src/render_resource/storage_buffer.rs +++ b/crates/bevy_render/src/render_resource/storage_buffer.rs @@ -96,13 +96,15 @@ impl StorageBuffer { fn size(&self) -> usize { let mut size = 0; size += T::std430_size_static(); - if size > 0 { - // Pad according to the array item type's alignment - size = (size + ::Output::ALIGNMENT - 1) - & !(::Output::ALIGNMENT - 1); + if self.item_size > 0 { + if size > 0 { + // Pad according to the array item type's alignment + size = (size + ::Output::ALIGNMENT - 1) + & !(::Output::ALIGNMENT - 1); + } + // Variable size arrays must have at least 1 element + size += self.item_size * self.capacity.max(1); } - // Variable size arrays must have at least 1 element - size += self.item_size * self.capacity.max(1); size } @@ -118,18 +120,21 @@ impl StorageBuffer { offset = new_offset; } } - if self.values.is_empty() { - for i in offset..self.size() { - self.scratch[i] = 0; + if self.item_size > 0 { + if self.values.is_empty() { + // Zero-out the padding and dummy array item in the case of the array being empty + for i in offset..self.size() { + self.scratch[i] = 0; + } + } else { + // Then write the array. Note that padding bytes may be added between the body + // and the array in order to align the array to the alignment requirements of its + // items + writer + .write(self.values.as_slice()) + .map_err(|e| warn!("{:?}", e)) + .ok(); } - } else { - // Then write the array. Note that padding bytes may be added between the body - // and the array in order to align the array to the alignment requirements of its - // items - writer - .write(self.values.as_slice()) - .map_err(|e| warn!("{:?}", e)) - .ok(); } queue.write_buffer(storage_buffer, 0, &self.scratch[range]); } From 136e146b759b11dbdd01cea362fb781b0c85b0ca Mon Sep 17 00:00:00 2001 From: Robert Swain Date: Wed, 2 Mar 2022 13:40:14 +0100 Subject: [PATCH 4/4] bevy_render: Make StorageBuffer consistent with changes in #3532 --- .../src/render_resource/storage_buffer.rs | 80 ++++++++----------- 1 file changed, 34 insertions(+), 46 deletions(-) diff --git a/crates/bevy_render/src/render_resource/storage_buffer.rs b/crates/bevy_render/src/render_resource/storage_buffer.rs index f750ba879faeda..2132390818cc96 100644 --- a/crates/bevy_render/src/render_resource/storage_buffer.rs +++ b/crates/bevy_render/src/render_resource/storage_buffer.rs @@ -1,4 +1,7 @@ -use std::num::NonZeroU64; +use std::{ + num::NonZeroU64, + ops::{Deref, DerefMut}, +}; use bevy_crevice::std430::{self, AsStd430, Std430}; use bevy_utils::tracing::warn; @@ -14,33 +17,38 @@ pub struct StorageBuffer { values: Vec, scratch: Vec, storage_buffer: Option, - capacity: usize, - item_size: usize, } impl Default for StorageBuffer { + /// Creates a new [`StorageBuffer`] + /// + /// This does not immediately allocate system/video RAM buffers. fn default() -> Self { Self { body: T::default(), values: Vec::new(), scratch: Vec::new(), storage_buffer: None, - capacity: 0, - item_size: U::std430_size_static(), } } } impl StorageBuffer { + // NOTE: AsStd430::std430_size_static() uses size_of internally but trait functions cannot be + // marked as const functions + const BODY_SIZE: usize = std::mem::size_of::(); + const ITEM_SIZE: usize = std::mem::size_of::(); + + /// Gets the reference to the underlying buffer, if one has been allocated. #[inline] - pub fn storage_buffer(&self) -> Option<&Buffer> { + pub fn buffer(&self) -> Option<&Buffer> { self.storage_buffer.as_ref() } #[inline] pub fn binding(&self) -> Option { Some(BindingResource::Buffer(BufferBinding { - buffer: self.storage_buffer()?, + buffer: self.buffer()?, offset: 0, size: Some(NonZeroU64::new((self.size()) as u64).unwrap()), })) @@ -51,35 +59,9 @@ impl StorageBuffer { self.body = body; } - #[inline] - pub fn len(&self) -> usize { - self.values.len() - } - - #[inline] - pub fn is_empty(&self) -> bool { - self.values.is_empty() - } - - #[inline] - pub fn capacity(&self) -> usize { - self.capacity - } - - pub fn push(&mut self, value: U) -> usize { - let index = self.values.len(); - self.values.push(value); - index - } - - pub fn get_mut(&mut self, index: usize) -> &mut U { - &mut self.values[index] - } - - pub fn reserve(&mut self, capacity: usize, device: &RenderDevice) -> bool { - if self.storage_buffer.is_none() || capacity > self.capacity { - self.capacity = capacity; - let size = self.size(); + fn reserve_buffer(&mut self, device: &RenderDevice) -> bool { + let size = self.size(); + if self.storage_buffer.is_none() || size > self.scratch.len() { self.scratch.resize(size, 0); self.storage_buffer = Some(device.create_buffer(&BufferDescriptor { label: None, @@ -95,32 +77,32 @@ impl StorageBuffer { fn size(&self) -> usize { let mut size = 0; - size += T::std430_size_static(); - if self.item_size > 0 { + size += Self::BODY_SIZE; + if Self::ITEM_SIZE > 0 { if size > 0 { // Pad according to the array item type's alignment size = (size + ::Output::ALIGNMENT - 1) & !(::Output::ALIGNMENT - 1); } // Variable size arrays must have at least 1 element - size += self.item_size * self.capacity.max(1); + size += Self::ITEM_SIZE * self.values.len().max(1); } size } pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) { - self.reserve(self.values.len(), device); + self.reserve_buffer(device); if let Some(storage_buffer) = &self.storage_buffer { let range = 0..self.size(); let mut writer = std430::Writer::new(&mut self.scratch[range.clone()]); let mut offset = 0; // First write the struct body if there is one - if T::std430_size_static() > 0 { + if Self::BODY_SIZE > 0 { if let Ok(new_offset) = writer.write(&self.body).map_err(|e| warn!("{:?}", e)) { offset = new_offset; } } - if self.item_size > 0 { + if Self::ITEM_SIZE > 0 { if self.values.is_empty() { // Zero-out the padding and dummy array item in the case of the array being empty for i in offset..self.size() { @@ -139,12 +121,18 @@ impl StorageBuffer { queue.write_buffer(storage_buffer, 0, &self.scratch[range]); } } +} - pub fn clear(&mut self) { - self.values.clear(); - } +impl Deref for StorageBuffer { + type Target = Vec; - pub fn values(&self) -> &[U] { + fn deref(&self) -> &Self::Target { &self.values } } + +impl DerefMut for StorageBuffer { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.values + } +}