StagingBelt: add a test and document slice size requirements.

* Add documentation that every `allocate()` or `write_buffer()`
  operation must have a size that is a multiple of 4.

* Add assertions for those properties (this gives more helpful panics
  than leaving it to validation).

* Add a randomized test to exercise usage of StagingBelt.
  I was concerned that `StagingBelt` might have an alignment math bug;
  having written this test, I am much less concerned. There was previously
  no test for `StagingBelt` at all, if you don’t count examples/skybox.
This commit is contained in:
Kevin Reid 2025-10-30 21:39:57 -07:00 committed by Connor Fitzgerald
parent a682d9e15a
commit da927daf82
3 changed files with 65 additions and 0 deletions

View File

@ -2,3 +2,4 @@
mod api;
mod noop;
mod util;

View File

@ -0,0 +1,42 @@
//! Tests of [`wgpu::util`].
use nanorand::Rng;
/// Generate (deterministic) random staging belt operations to exercise its logic.
#[test]
fn staging_belt_random_test() {
let (device, queue) = wgpu::Device::noop(&wgpu::DeviceDescriptor::default());
let mut rng = nanorand::WyRand::new_seed(0xDEAD_BEEF);
let buffer_size = 1024;
let align = wgpu::COPY_BUFFER_ALIGNMENT;
let mut belt = wgpu::util::StagingBelt::new(buffer_size / 2);
let target_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: None,
size: buffer_size,
usage: wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
for _batch in 0..100 {
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
for _write in 0..5 {
let offset: u64 = rng.generate_range(0..=(buffer_size - align) / align) * align;
let size: u64 = rng.generate_range(1..=(buffer_size - offset) / align) * align;
println!("offset {offset} size {size}");
let mut slice = belt.write_buffer(
&mut encoder,
&target_buffer,
offset,
wgpu::BufferSize::new(size).unwrap(),
&device,
);
slice[0] = 1; // token amount of actual writing, just in case it makes a difference
}
belt.finish();
queue.submit([encoder.finish()]);
belt.recall();
}
}

View File

@ -6,6 +6,8 @@ use alloc::vec::Vec;
use core::fmt;
use std::sync::mpsc;
use crate::COPY_BUFFER_ALIGNMENT;
/// Efficiently performs many buffer writes by sharing and reusing temporary buffers.
///
/// Internally it uses a ring-buffer of staging buffers that are sub-allocated.
@ -63,6 +65,9 @@ impl StagingBelt {
/// Allocate a staging belt slice of `size` to be copied into the `target` buffer
/// at the specified offset.
///
/// `offset` and `size` must be multiples of [`COPY_BUFFER_ALIGNMENT`]
/// (as is required by the underlying buffer operations).
///
/// The upload will be placed into the provided command encoder. This encoder
/// must be submitted after [`StagingBelt::finish()`] is called and before
/// [`StagingBelt::recall()`] is called.
@ -70,6 +75,7 @@ impl StagingBelt {
/// If the `size` is greater than the size of any free internal buffer, a new buffer
/// will be allocated for it. Therefore, the `chunk_size` passed to [`StagingBelt::new()`]
/// should ideally be larger than every such size.
#[track_caller]
pub fn write_buffer(
&mut self,
encoder: &mut CommandEncoder,
@ -78,6 +84,14 @@ impl StagingBelt {
size: BufferSize,
device: &Device,
) -> BufferViewMut {
// Asserting this explicitly gives a usefully more specific, and more prompt, error than
// leaving it to regular API validation.
// We check only `offset`, not `size`, because `self.allocate()` will check the size.
assert!(
offset.is_multiple_of(COPY_BUFFER_ALIGNMENT),
"StagingBelt::write_buffer() offset {offset} must be a multiple of `COPY_BUFFER_ALIGNMENT`"
);
let slice_of_belt = self.allocate(
size,
const { BufferSize::new(crate::COPY_BUFFER_ALIGNMENT).unwrap() },
@ -95,6 +109,9 @@ impl StagingBelt {
/// Allocate a staging belt slice with the given `size` and `alignment` and return it.
///
/// `size` must be a multiple of [`COPY_BUFFER_ALIGNMENT`]
/// (as is required by the underlying buffer operations).
///
/// To use this slice, call [`BufferSlice::get_mapped_range_mut()`] and write your data into
/// that [`BufferViewMut`].
/// (The view must be dropped before [`StagingBelt::finish()`] is called.)
@ -112,12 +129,17 @@ impl StagingBelt {
/// The chosen slice will be positioned within the buffer at a multiple of `alignment`,
/// which may be used to meet alignment requirements for the operation you wish to perform
/// with the slice. This does not necessarily affect the alignment of the [`BufferViewMut`].
#[track_caller]
pub fn allocate(
&mut self,
size: BufferSize,
alignment: BufferSize,
device: &Device,
) -> BufferSlice<'_> {
assert!(
size.get().is_multiple_of(COPY_BUFFER_ALIGNMENT),
"StagingBelt allocation size {size} must be a multiple of `COPY_BUFFER_ALIGNMENT`"
);
assert!(
alignment.get().is_power_of_two(),
"alignment must be a power of two, not {alignment}"