docs/bevy/examples/shader/gpu_readback.rs 8.1 K · 233 lines · raw · history

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
4 use 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
23 const SHADER_ASSET_PATH: &str = "shaders/gpu_readback.wesl";
24
25 // The length of the buffer sent to the gpu
26 const BUFFER_LEN: usize = 16;
27
28 fn 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
43 struct GpuReadbackPlugin;
44 impl 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)]
64 struct ReadbackBuffer(Handle<ShaderBuffer>);
65
66 fn 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)]
83 struct ReadbackImage(Handle<Image>);
84
85 fn 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)]
158 struct GpuBufferBindGroup(BindGroup);
159
160 fn 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)]
184 struct ComputePipeline {
185 layout: BindGroupLayoutDescriptor,
186 pipeline: CachedComputePipelineId,
187 }
188
189 fn 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
214 fn 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 }