GpuReadbackPlugin: Allow reading only a part of a buffer (#20133)

# Objective

- So far only full buffer reads were supported. This adds the ability to
read a part of a buffer.

## Solution

- Allow passing in a start offset and a number of bytes to read when
creating the `Readback` component.
- I also removed the unused `src_start` and `dst_start` fields from
`ReadbackSource` as they were always 0.

## Testing

- Did you test these changes? If so, how?

I extended the example to also demonstrate partial reads.

- Are there any parts that need more testing?

Can't think of any.

- How can other people (reviewers) test your changes? Is there anything
specific they need to know?

Run the `gpu_readback` example. It now also reads and prints a partially
read buffer.

- If relevant, what platforms did you test these changes on, and are
there any important ones you can't test?

Only tested on Linux.

---

## Showcase

Example output:

<details>
  <summary>Click to view showcase</summary>

```
2025-07-14T14:05:15.614876Z  INFO gpu_readback: Buffer [257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272]
2025-07-14T14:05:15.614921Z  INFO gpu_readback: Buffer range [261, 262, 263, 264, 265, 266, 267, 268]
2025-07-14T14:05:15.614937Z  INFO gpu_readback: Image [257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
```

</details>
This commit is contained in:
Sven Niederberger 2025-07-14 23:01:16 +02:00 committed by GitHub
parent 0294fa89db
commit 3ce2f46ae5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 54 additions and 19 deletions

View File

@ -77,7 +77,10 @@ impl Plugin for GpuReadbackPlugin {
#[derive(Component, ExtractComponent, Clone, Debug)]
pub enum Readback {
Texture(Handle<Image>),
Buffer(Handle<ShaderStorageBuffer>),
Buffer {
buffer: Handle<ShaderStorageBuffer>,
start_offset_and_size: Option<(u64, u64)>,
},
}
impl Readback {
@ -86,9 +89,21 @@ impl Readback {
Self::Texture(image)
}
/// Create a readback component for a buffer using the given handle.
/// Create a readback component for a full buffer using the given handle.
pub fn buffer(buffer: Handle<ShaderStorageBuffer>) -> Self {
Self::Buffer(buffer)
Self::Buffer {
buffer,
start_offset_and_size: None,
}
}
/// Create a readback component for a buffer range using the given handle, a start offset in bytes
/// and a number of bytes to read.
pub fn buffer_range(buffer: Handle<ShaderStorageBuffer>, start_offset: u64, size: u64) -> Self {
Self::Buffer {
buffer,
start_offset_and_size: Some((start_offset, size)),
}
}
}
@ -193,9 +208,8 @@ enum ReadbackSource {
size: Extent3d,
},
Buffer {
src_start: u64,
dst_start: u64,
buffer: Buffer,
start_offset_and_size: Option<(u64, u64)>,
},
}
@ -266,16 +280,30 @@ fn prepare_buffers(
});
}
}
Readback::Buffer(buffer) => {
Readback::Buffer {
buffer,
start_offset_and_size,
} => {
if let Some(ssbo) = ssbos.get(buffer) {
let size = ssbo.buffer.size();
let full_size = ssbo.buffer.size();
let size = start_offset_and_size
.map(|(start, size)| {
let end = start + size;
if end > full_size {
panic!(
"Tried to read past the end of the buffer (start: {start}, \
size: {size}, buffer size: {full_size})."
);
}
size
})
.unwrap_or(full_size);
let buffer = buffer_pool.get(&render_device, size);
let (tx, rx) = async_channel::bounded(1);
readbacks.requested.push(GpuReadback {
entity: entity.id(),
src: ReadbackSource::Buffer {
src_start: 0,
dst_start: 0,
start_offset_and_size: *start_offset_and_size,
buffer: ssbo.buffer.clone(),
},
buffer,
@ -307,17 +335,11 @@ pub(crate) fn submit_readback_commands(world: &World, command_encoder: &mut Comm
);
}
ReadbackSource::Buffer {
src_start,
dst_start,
buffer,
start_offset_and_size,
} => {
command_encoder.copy_buffer_to_buffer(
buffer,
*src_start,
&readback.buffer,
*dst_start,
buffer.size(),
);
let (src_start, size) = start_offset_and_size.unwrap_or((0, buffer.size()));
command_encoder.copy_buffer_to_buffer(buffer, src_start, &readback.buffer, 0, size);
}
}
}

View File

@ -74,7 +74,7 @@ fn setup(
mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
) {
// Create a storage buffer with some data
let buffer = vec![0u32; BUFFER_LEN];
let buffer: Vec<u32> = (0..BUFFER_LEN as u32).collect();
let mut buffer = ShaderStorageBuffer::from(buffer);
// We need to enable the COPY_SRC usage so we can copy the buffer to the cpu
buffer.buffer_description.usage |= BufferUsages::COPY_SRC;
@ -110,6 +110,19 @@ fn setup(
let data: Vec<u32> = trigger.event().to_shader_type();
info!("Buffer {:?}", data);
});
// It is also possible to read only a range of the buffer.
commands
.spawn(Readback::buffer_range(
buffer.clone(),
4 * u32::SHADER_SIZE.get(), // skip the first four elements
8 * u32::SHADER_SIZE.get(), // read eight elements
))
.observe(|trigger: On<ReadbackComplete>| {
let data: Vec<u32> = trigger.event().to_shader_type();
info!("Buffer range {:?}", data);
});
// This is just a simple way to pass the buffer handle to the render app for our compute node
commands.insert_resource(ReadbackBuffer(buffer));