gpu_readback/
gpu_readback.rs

1//! Simple example demonstrating the use of the [`Readback`] component to read back data from the GPU
2//! using both a storage buffer and texture.
3
4use bevy::{
5    asset::RenderAssetUsages,
6    prelude::*,
7    render::{
8        extract_resource::{ExtractResource, ExtractResourcePlugin},
9        gpu_readback::{Readback, ReadbackComplete, ReadbackOnce},
10        render_asset::RenderAssets,
11        render_resource::{
12            binding_types::{storage_buffer, texture_storage_2d},
13            *,
14        },
15        renderer::{RenderContext, RenderDevice, RenderGraph},
16        storage::{GpuShaderBuffer, ShaderBuffer},
17        texture::GpuImage,
18        Render, RenderApp, RenderStartup, RenderSystems,
19    },
20};
21
22/// This example uses a shader source file from the assets subdirectory
23const SHADER_ASSET_PATH: &str = "shaders/gpu_readback.wesl";
24
25// The length of the buffer sent to the gpu
26const BUFFER_LEN: usize = 16;
27
28fn main() {
29    App::new()
30        .add_plugins((
31            DefaultPlugins,
32            GpuReadbackPlugin,
33            ExtractResourcePlugin::<ReadbackBuffer>::default(),
34            ExtractResourcePlugin::<ReadbackImage>::default(),
35        ))
36        .insert_resource(ClearColor(Color::BLACK))
37        .add_systems(Startup, setup)
38        .add_systems(Update, trigger_once_readback)
39        .run();
40}
41
42// We need a plugin to organize all the systems and render node required for this example
43struct GpuReadbackPlugin;
44impl Plugin for GpuReadbackPlugin {
45    fn build(&self, app: &mut App) {
46        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
47            return;
48        };
49        render_app
50            .add_systems(RenderStartup, init_compute_pipeline)
51            .add_systems(
52                Render,
53                prepare_bind_group
54                    .in_set(RenderSystems::PrepareBindGroups)
55                    // We don't need to recreate the bind group every frame
56                    .run_if(not(resource_exists::<GpuBufferBindGroup>)),
57            )
58            .add_systems(RenderGraph, compute);
59    }
60}
61
62#[derive(Resource, ExtractResource, Clone)]
63#[extract_app(RenderApp)]
64struct ReadbackBuffer(Handle<ShaderBuffer>);
65
66fn trigger_once_readback(
67    keyboard: Res<ButtonInput<KeyCode>>,
68    readback_buffer: Res<ReadbackBuffer>,
69    mut commands: Commands,
70) {
71    if keyboard.just_pressed(KeyCode::Space) {
72        commands
73            .spawn(ReadbackOnce::buffer(readback_buffer.0.clone()))
74            .observe(|event: On<ReadbackComplete>| {
75                let data: Vec<u32> = event.to_shader_type();
76                info!("Buffer (once) {:?}", data);
77            });
78    }
79}
80
81#[derive(Resource, ExtractResource, Clone)]
82#[extract_app(RenderApp)]
83struct ReadbackImage(Handle<Image>);
84
85fn setup(
86    mut commands: Commands,
87    mut images: ResMut<Assets<Image>>,
88    mut buffers: ResMut<Assets<ShaderBuffer>>,
89) {
90    // Create a storage buffer with some data
91    let buffer: Vec<u32> = (0..BUFFER_LEN as u32).collect();
92    let mut buffer = ShaderBuffer::from(buffer);
93    // We need to enable the COPY_SRC usage so we can copy the buffer to the cpu
94    buffer.buffer_usage |= BufferUsages::COPY_SRC;
95    let buffer = buffers.add(buffer);
96
97    // Create a storage texture with some data
98    let size = Extent3d {
99        width: BUFFER_LEN as u32,
100        height: 1,
101        ..default()
102    };
103    // We create an uninitialized image since this texture will only be used for getting data out
104    // of the compute shader, not getting data in, so there's no reason for it to exist on the CPU
105    let mut image = Image::new_uninit(
106        size,
107        TextureDimension::D2,
108        TextureFormat::R32Uint,
109        RenderAssetUsages::RENDER_WORLD,
110    );
111    // We also need to enable the COPY_SRC, as well as STORAGE_BINDING so we can use it in the
112    // compute shader
113    image.texture_descriptor.usage |= TextureUsages::COPY_SRC | TextureUsages::STORAGE_BINDING;
114    let image = images.add(image);
115
116    // Spawn the readback components. For each frame, the data will be read back from the GPU
117    // asynchronously and trigger the `ReadbackComplete` event on this entity. Despawn the entity
118    // to stop reading back the data.
119    commands
120        .spawn(Readback::buffer(buffer.clone()))
121        .observe(|event: On<ReadbackComplete>| {
122            // This matches the type which was used to create the `ShaderBuffer` above,
123            // and is a convenient way to interpret the data.
124            let data: Vec<u32> = event.to_shader_type();
125            info!("Buffer {:?}", data);
126        });
127
128    // It is also possible to read only a range of the buffer.
129    commands
130        .spawn(Readback::buffer_range(
131            buffer.clone(),
132            4 * u32::SHADER_SIZE.get(), // skip the first four elements
133            8 * u32::SHADER_SIZE.get(), // read eight elements
134        ))
135        .observe(|event: On<ReadbackComplete>| {
136            let data: Vec<u32> = event.to_shader_type();
137            info!("Buffer range {:?}", data);
138        });
139
140    // This is just a simple way to pass the buffer handle to the render app for our compute node
141    commands.insert_resource(ReadbackBuffer(buffer));
142
143    // Textures can also be read back from the GPU. Pay careful attention to the format of the
144    // texture, as it will affect how the data is interpreted.
145    commands
146        .spawn(Readback::texture(image.clone()))
147        .observe(|event: On<ReadbackComplete>| {
148            // You probably want to interpret the data as a color rather than a `ShaderType`,
149            // but in this case we know the data is a single channel storage texture, so we can
150            // interpret it as a `Vec<u32>`
151            let data: Vec<u32> = event.to_shader_type();
152            info!("Image {:?}", data);
153        });
154    commands.insert_resource(ReadbackImage(image));
155}
156
157#[derive(Resource)]
158struct GpuBufferBindGroup(BindGroup);
159
160fn prepare_bind_group(
161    mut commands: Commands,
162    pipeline: Res<ComputePipeline>,
163    render_device: Res<RenderDevice>,
164    pipeline_cache: Res<PipelineCache>,
165    buffer: Res<ReadbackBuffer>,
166    image: Res<ReadbackImage>,
167    buffers: Res<RenderAssets<GpuShaderBuffer>>,
168    images: Res<RenderAssets<GpuImage>>,
169) {
170    let buffer = buffers.get(&buffer.0).unwrap();
171    let image = images.get(&image.0).unwrap();
172    let bind_group = render_device.create_bind_group(
173        None,
174        &pipeline_cache.get_bind_group_layout(&pipeline.layout),
175        &BindGroupEntries::sequential((
176            buffer.buffer.as_entire_buffer_binding(),
177            image.texture_view.into_binding(),
178        )),
179    );
180    commands.insert_resource(GpuBufferBindGroup(bind_group));
181}
182
183#[derive(Resource)]
184struct ComputePipeline {
185    layout: BindGroupLayoutDescriptor,
186    pipeline: CachedComputePipelineId,
187}
188
189fn init_compute_pipeline(
190    mut commands: Commands,
191    asset_server: Res<AssetServer>,
192    pipeline_cache: Res<PipelineCache>,
193) {
194    let layout = BindGroupLayoutDescriptor::new(
195        "",
196        &BindGroupLayoutEntries::sequential(
197            ShaderStages::COMPUTE,
198            (
199                storage_buffer::<Vec<u32>>(false),
200                texture_storage_2d(TextureFormat::R32Uint, StorageTextureAccess::WriteOnly),
201            ),
202        ),
203    );
204    let shader = asset_server.load(SHADER_ASSET_PATH);
205    let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
206        label: Some("GPU readback compute shader".into()),
207        layout: vec![layout.clone()],
208        shader: shader.clone(),
209        ..default()
210    });
211    commands.insert_resource(ComputePipeline { layout, pipeline });
212}
213
214fn compute(
215    mut render_context: RenderContext,
216    pipeline_cache: Res<PipelineCache>,
217    pipeline: Res<ComputePipeline>,
218    bind_group: Res<GpuBufferBindGroup>,
219) {
220    if let Some(init_pipeline) = pipeline_cache.get_compute_pipeline(pipeline.pipeline) {
221        let mut pass =
222            render_context
223                .command_encoder()
224                .begin_compute_pass(&ComputePassDescriptor {
225                    label: Some("GPU readback compute pass"),
226                    ..default()
227                });
228
229        pass.set_bind_group(0, &bind_group.0, &[]);
230        pass.set_pipeline(init_pipeline);
231        pass.dispatch_workgroups(BUFFER_LEN as u32, 1, 1);
232    }
233}
gpu_readback/gpu_readback.rs

gpu_readback/
gpu_readback.rs
