bevy_render/
gpu_readback.rs

1use crate::{
2    extract_component::ExtractComponentPlugin,
3    render_asset::RenderAssets,
4    render_resource::{
5        Buffer, BufferUsages, CommandEncoder, Extent3d, TexelCopyBufferLayout, Texture,
6        TextureFormat,
7    },
8    renderer::RenderDevice,
9    storage::{GpuShaderBuffer, ShaderBuffer},
10    sync_world::MainEntity,
11    texture::GpuImage,
12    ExtractSchedule, MainWorld, Render, RenderApp, RenderSystems,
13};
14use async_channel::{Receiver, Sender};
15use bevy_app::{App, First, Plugin};
16use bevy_asset::Handle;
17use bevy_derive::{Deref, DerefMut};
18use bevy_ecs::{
19    bundle::Bundle,
20    change_detection::ResMut,
21    entity::Entity,
22    event::EntityEvent,
23    prelude::{Component, Resource, World},
24    query::With,
25    system::{Commands, Query, Res},
26};
27use bevy_ecs::{schedule::IntoScheduleConfigs, template::FromTemplate};
28use bevy_extract_macros::ExtractComponent;
29use bevy_image::{Image, TextureFormatPixelInfo};
30use bevy_log::{debug, warn};
31use bevy_platform::collections::HashMap;
32use bevy_reflect::Reflect;
33use encase::internal::ReadFrom;
34use encase::private::Reader;
35use encase::ShaderType;
36
37/// A plugin that enables reading back gpu buffers and textures to the cpu.
38pub struct GpuReadbackPlugin {
39    /// Describes the number of frames a buffer can be unused before it is removed from the pool in
40    /// order to avoid unnecessary reallocations.
41    max_unused_frames: usize,
42}
43
44impl Default for GpuReadbackPlugin {
45    fn default() -> Self {
46        Self {
47            max_unused_frames: 10,
48        }
49    }
50}
51
52impl Plugin for GpuReadbackPlugin {
53    fn build(&self, app: &mut App) {
54        app.add_plugins(ExtractComponentPlugin::<Readback>::default())
55            .register_type::<ReadbackOnce>()
56            .add_systems(
57                First,
58                cleanup_readback_once.after(bevy_ecs::message::MessageUpdateSystems),
59            );
60
61        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
62            render_app
63                .init_resource::<GpuReadbackBufferPool>()
64                .init_resource::<GpuReadbacks>()
65                .insert_resource(GpuReadbackMaxUnusedFrames(self.max_unused_frames))
66                .add_systems(ExtractSchedule, sync_readbacks.ambiguous_with_all())
67                .add_systems(
68                    Render,
69                    (
70                        prepare_buffers.in_set(RenderSystems::PrepareResources),
71                        // TODO: this should be in the graph somehow
72                        map_buffers.in_set(RenderSystems::Cleanup),
73                    ),
74                );
75        }
76    }
77}
78
79/// A component that registers the wrapped handle for gpu readback, either a texture or a buffer.
80///
81/// Data is read asynchronously and will be triggered on the entity via the [`ReadbackComplete`] event
82/// when complete. If this component is not removed, the readback will be attempted every frame
83#[derive(Component, ExtractComponent, Clone, Debug, FromTemplate)]
84#[extract_app(RenderApp)]
85pub enum Readback {
86    #[default]
87    Texture(Handle<Image>),
88    Buffer {
89        buffer: Handle<ShaderBuffer>,
90        start_offset_and_size: Option<(u64, u64)>,
91    },
92}
93
94impl Readback {
95    /// Create a readback component for a texture using the given handle.
96    pub fn texture(image: Handle<Image>) -> Self {
97        Self::Texture(image)
98    }
99
100    /// Create a readback component for a full buffer using the given handle.
101    pub fn buffer(buffer: Handle<ShaderBuffer>) -> Self {
102        Self::Buffer {
103            buffer,
104            start_offset_and_size: None,
105        }
106    }
107
108    /// Create a readback component for a buffer range using the given handle, a start offset in bytes
109    /// and a number of bytes to read.
110    pub fn buffer_range(buffer: Handle<ShaderBuffer>, start_offset: u64, size: u64) -> Self {
111        Self::Buffer {
112            buffer,
113            start_offset_and_size: Some((start_offset, size)),
114        }
115    }
116}
117
118/// Marker component that pairs with a [`Readback`] to indicate the entity should be
119/// despawned after the first [`ReadbackComplete`] event fires for it.
120///
121/// ```ignore
122/// commands.spawn(ReadbackOnce::buffer(handle));
123/// ```
124#[derive(Component, Clone, Debug, Default, Reflect)]
125pub struct ReadbackOnce;
126
127impl ReadbackOnce {
128    /// Create a one-shot readback bundle for a texture using the given handle.
129    pub fn texture(image: Handle<Image>) -> impl Bundle {
130        (Readback::texture(image), Self)
131    }
132
133    /// Create a one-shot readback bundle for a full buffer using the given handle.
134    pub fn buffer(buffer: Handle<ShaderBuffer>) -> impl Bundle {
135        (Readback::buffer(buffer), Self)
136    }
137
138    /// Create a one-shot readback bundle for a buffer range using the given handle, a start
139    /// offset in bytes and a number of bytes to read.
140    pub fn buffer_range(buffer: Handle<ShaderBuffer>, start_offset: u64, size: u64) -> impl Bundle {
141        (Readback::buffer_range(buffer, start_offset, size), Self)
142    }
143}
144
145fn cleanup_readback_once(mut commands: Commands, query: Query<Entity, With<ReadbackOnce>>) {
146    for entity in &query {
147        commands.entity(entity).remove::<(Readback, ReadbackOnce)>();
148    }
149}
150
151/// An event that is triggered when a gpu readback is complete.
152///
153/// The event contains the data as a `Vec<u8>`, which can be interpreted as the raw bytes of the
154/// requested buffer or texture.
155#[derive(EntityEvent, Deref, DerefMut, Reflect, Debug)]
156#[reflect(Debug)]
157pub struct ReadbackComplete {
158    pub entity: Entity,
159    #[deref]
160    pub data: Vec<u8>,
161}
162
163impl ReadbackComplete {
164    /// Convert the raw bytes of the event to a shader type.
165    pub fn to_shader_type<T: ShaderType + ReadFrom + Default>(&self) -> T {
166        let mut val = T::default();
167        let mut reader = Reader::new::<T>(&self.data, 0).expect("Failed to create Reader");
168        T::read_from(&mut val, &mut reader);
169        val
170    }
171}
172
173#[derive(Resource)]
174struct GpuReadbackMaxUnusedFrames(usize);
175
176struct GpuReadbackBuffer {
177    buffer: Buffer,
178    taken: bool,
179    frames_unused: usize,
180}
181
182#[derive(Resource, Default)]
183struct GpuReadbackBufferPool {
184    // Map of buffer size to list of buffers, with a flag for whether the buffer is taken and how
185    // many frames it has been unused for.
186    // TODO: We could ideally write all readback data to one big buffer per frame, the assumption
187    // here is that very few entities well actually be read back at once, and their size is
188    // unlikely to change.
189    buffers: HashMap<u64, Vec<GpuReadbackBuffer>>,
190}
191
192impl GpuReadbackBufferPool {
193    fn get(&mut self, render_device: &RenderDevice, size: u64) -> Buffer {
194        let buffers = self.buffers.entry(size).or_default();
195
196        // find an untaken buffer for this size
197        if let Some(buf) = buffers.iter_mut().find(|x| !x.taken) {
198            buf.taken = true;
199            buf.frames_unused = 0;
200            return buf.buffer.clone();
201        }
202
203        let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
204            label: Some("Readback Buffer"),
205            size,
206            usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
207            mapped_at_creation: false,
208        });
209        buffers.push(GpuReadbackBuffer {
210            buffer: buffer.clone(),
211            taken: true,
212            frames_unused: 0,
213        });
214        buffer
215    }
216
217    // Returns the buffer to the pool so it can be used in a future frame
218    fn return_buffer(&mut self, buffer: &Buffer) {
219        let size = buffer.size();
220        let buffers = self
221            .buffers
222            .get_mut(&size)
223            .expect("Returned buffer of untracked size");
224        if let Some(buf) = buffers.iter_mut().find(|x| x.buffer.id() == buffer.id()) {
225            buf.taken = false;
226        } else {
227            warn!("Returned buffer that was not allocated");
228        }
229    }
230
231    fn update(&mut self, max_unused_frames: usize) {
232        for (_, buffers) in &mut self.buffers {
233            // Tick all the buffers
234            for buf in &mut *buffers {
235                if !buf.taken {
236                    buf.frames_unused += 1;
237                }
238            }
239
240            // Remove buffers that haven't been used for MAX_UNUSED_FRAMES
241            buffers.retain(|x| x.frames_unused < max_unused_frames);
242        }
243
244        // Remove empty buffer sizes
245        self.buffers.retain(|_, buffers| !buffers.is_empty());
246    }
247}
248
249enum ReadbackSource {
250    Texture {
251        texture: Texture,
252        layout: TexelCopyBufferLayout,
253        size: Extent3d,
254    },
255    Buffer {
256        buffer: Buffer,
257        start_offset_and_size: Option<(u64, u64)>,
258    },
259}
260
261#[derive(Resource, Default)]
262struct GpuReadbacks {
263    requested: Vec<GpuReadback>,
264    mapped: Vec<GpuReadback>,
265}
266
267struct GpuReadback {
268    pub entity: Entity,
269    pub src: ReadbackSource,
270    pub buffer: Buffer,
271    pub rx: Receiver<(Entity, Buffer, Vec<u8>)>,
272    pub tx: Sender<(Entity, Buffer, Vec<u8>)>,
273}
274
275fn sync_readbacks(
276    mut main_world: ResMut<MainWorld>,
277    mut buffer_pool: ResMut<GpuReadbackBufferPool>,
278    mut readbacks: ResMut<GpuReadbacks>,
279    max_unused_frames: Res<GpuReadbackMaxUnusedFrames>,
280) {
281    readbacks.mapped.retain(|readback| {
282        if let Ok((entity, buffer, data)) = readback.rx.try_recv() {
283            main_world.trigger(ReadbackComplete { data, entity });
284            buffer_pool.return_buffer(&buffer);
285            false
286        } else {
287            true
288        }
289    });
290
291    buffer_pool.update(max_unused_frames.0);
292}
293
294fn prepare_buffers(
295    render_device: Res<RenderDevice>,
296    mut readbacks: ResMut<GpuReadbacks>,
297    mut buffer_pool: ResMut<GpuReadbackBufferPool>,
298    gpu_images: Res<RenderAssets<GpuImage>>,
299    ssbos: Res<RenderAssets<GpuShaderBuffer>>,
300    handles: Query<(&MainEntity, &Readback)>,
301) {
302    for (entity, readback) in handles.iter() {
303        match readback {
304            Readback::Texture(image) => {
305                if let Some(gpu_image) = gpu_images.get(image)
306                    && let Ok(pixel_size) = gpu_image.texture_descriptor.format.pixel_size()
307                {
308                    let layout = layout_data(
309                        gpu_image.texture_descriptor.size,
310                        gpu_image.texture_descriptor.format,
311                    );
312                    let buffer = buffer_pool.get(
313                        &render_device,
314                        get_aligned_size(gpu_image.texture_descriptor.size, pixel_size as u32)
315                            as u64,
316                    );
317                    let (tx, rx) = async_channel::bounded(1);
318                    readbacks.requested.push(GpuReadback {
319                        entity: entity.id(),
320                        src: ReadbackSource::Texture {
321                            texture: gpu_image.texture.clone(),
322                            layout,
323                            size: gpu_image.texture_descriptor.size,
324                        },
325                        buffer,
326                        rx,
327                        tx,
328                    });
329                }
330            }
331            Readback::Buffer {
332                buffer,
333                start_offset_and_size,
334            } => {
335                if let Some(ssbo) = ssbos.get(buffer) {
336                    let full_size = ssbo.buffer.size();
337                    let size = start_offset_and_size
338                        .map(|(start, size)| {
339                            let end = start + size;
340                            if end > full_size {
341                                panic!(
342                                    "Tried to read past the end of the buffer (start: {start}, \
343                                    size: {size}, buffer size: {full_size})."
344                                );
345                            }
346                            size
347                        })
348                        .unwrap_or(full_size);
349                    let buffer = buffer_pool.get(&render_device, size);
350                    let (tx, rx) = async_channel::bounded(1);
351                    readbacks.requested.push(GpuReadback {
352                        entity: entity.id(),
353                        src: ReadbackSource::Buffer {
354                            start_offset_and_size: *start_offset_and_size,
355                            buffer: ssbo.buffer.clone(),
356                        },
357                        buffer,
358                        rx,
359                        tx,
360                    });
361                }
362            }
363        }
364    }
365}
366
367pub(crate) fn submit_readback_commands(world: &World, command_encoder: &mut CommandEncoder) {
368    let readbacks = world.resource::<GpuReadbacks>();
369    for readback in &readbacks.requested {
370        match &readback.src {
371            ReadbackSource::Texture {
372                texture,
373                layout,
374                size,
375            } => {
376                command_encoder.copy_texture_to_buffer(
377                    texture.as_image_copy(),
378                    wgpu::TexelCopyBufferInfo {
379                        buffer: &readback.buffer,
380                        layout: *layout,
381                    },
382                    *size,
383                );
384            }
385            ReadbackSource::Buffer {
386                buffer,
387                start_offset_and_size,
388            } => {
389                let (src_start, size) = start_offset_and_size.unwrap_or((0, buffer.size()));
390                command_encoder.copy_buffer_to_buffer(buffer, src_start, &readback.buffer, 0, size);
391            }
392        }
393    }
394}
395
396/// Move requested readbacks to mapped readbacks after commands have been submitted in render system
397#[expect(
398    clippy::drain_collect,
399    reason = "draining preserves the capacity of `requested`, which is refilled every frame"
400)]
401fn map_buffers(mut readbacks: ResMut<GpuReadbacks>) {
402    let requested = readbacks.requested.drain(..).collect::<Vec<GpuReadback>>();
403    for readback in requested {
404        let slice = readback.buffer.slice(..);
405        let entity = readback.entity;
406        let buffer = readback.buffer.clone();
407        let tx = readback.tx.clone();
408        slice.map_async(wgpu::MapMode::Read, move |res| {
409            res.expect("Failed to map buffer");
410            let buffer_slice = buffer.slice(..);
411            let data = buffer_slice.get_mapped_range().unwrap();
412            let result = Vec::from(&*data);
413            drop(data);
414            buffer.unmap();
415            if let Err(e) = tx.try_send((entity, buffer, result)) {
416                debug!("Failed to send readback result: {}", e);
417            }
418        });
419        readbacks.mapped.push(readback);
420    }
421}
422
423// Utils
424
425/// Round up a given value to be a multiple of [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`].
426pub(crate) const fn align_byte_size(value: u32) -> u32 {
427    RenderDevice::align_copy_bytes_per_row(value as usize) as u32
428}
429
430/// Get the size of an image when the size of each row has been rounded up to [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`].
431pub(crate) const fn get_aligned_size(extent: Extent3d, pixel_size: u32) -> u32 {
432    extent.height * align_byte_size(extent.width * pixel_size) * extent.depth_or_array_layers
433}
434
435/// Get a [`TexelCopyBufferLayout`] aligned such that the image can be copied into a buffer.
436pub(crate) fn layout_data(extent: Extent3d, format: TextureFormat) -> TexelCopyBufferLayout {
437    TexelCopyBufferLayout {
438        bytes_per_row: if extent.height > 1 || extent.depth_or_array_layers > 1 {
439            if let Ok(pixel_size) = format.pixel_size() {
440                // 1 = 1 row
441                Some(get_aligned_size(
442                    Extent3d {
443                        width: extent.width,
444                        height: 1,
445                        depth_or_array_layers: 1,
446                    },
447                    pixel_size as u32,
448                ))
449            } else {
450                None
451            }
452        } else {
453            None
454        },
455        rows_per_image: if extent.depth_or_array_layers > 1 {
456            let (_, block_dimension_y) = format.block_dimensions();
457            Some(extent.height / block_dimension_y)
458        } else {
459            None
460        },
461        offset: 0,
462    }
463}
bevy_render/gpu_readback.rs

bevy_render/
gpu_readback.rs
