Gpu readback belt for fast & easy data readback from gpu

.gitignore

@@ -17,3 +17,6 @@ __pycache__
 perf.data*
 
 **/dataset/
+
+# Screenshots from samples etc.
+screenshot*.png

.vscode/settings.json

@@ -34,6 +34,7 @@
         "memoffset",
         "nyud",
         "objectron",
+        "Readback",
         "Skybox",
         "smallvec",
         "swapchain",

crates/re_renderer/examples/multiview.rs

@@ -12,31 +12,15 @@ use re_renderer::{
         GenericSkyboxDrawData, LineDrawData, LineStripFlags, MeshDrawData, MeshInstance,
         TestTriangleDrawData,
     },
-    view_builder::{OrthographicCameraMode, Projection, TargetConfiguration, ViewBuilder},
+    view_builder::{
+        OrthographicCameraMode, Projection, ScheduledScreenshot, TargetConfiguration, ViewBuilder,
+    },
     Color32, LineStripSeriesBuilder, PointCloudBuilder, RenderContext, Rgba, Size,
 };
 use winit::event::{ElementState, VirtualKeyCode};
 
 mod framework;
 
-fn draw_view<D: 'static + re_renderer::renderer::DrawData + Sync + Send + Clone>(
-    re_ctx: &mut RenderContext,
-    target_cfg: TargetConfiguration,
-    skybox: &GenericSkyboxDrawData,
-    draw_data: &D,
-) -> (ViewBuilder, wgpu::CommandBuffer) {
-    let mut view_builder = ViewBuilder::default();
-    let command_buffer = view_builder
-        .setup_view(re_ctx, target_cfg)
-        .unwrap()
-        .queue_draw(skybox)
-        .queue_draw(draw_data)
-        .draw(re_ctx, Rgba::TRANSPARENT)
-        .unwrap();
-
-    (view_builder, command_buffer)
-}
-
 fn build_mesh_instances(
     re_ctx: &mut RenderContext,
     model_mesh_instances: &[MeshInstance],
@@ -165,6 +149,9 @@ struct Multiview {
     random_points_positions: Vec<glam::Vec3>,
     random_points_radii: Vec<Size>,
     random_points_colors: Vec<Color32>,
+
+    take_screenshot_next_frame_for_view: Option<u32>,
+    scheduled_screenshots: Vec<ScheduledScreenshot>,
 }
 
 fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
@@ -177,6 +164,83 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
     .into()
 }
 
+impl Multiview {
+    fn draw_view<D: 'static + re_renderer::renderer::DrawData + Sync + Send + Clone>(
+        &mut self,
+        re_ctx: &mut RenderContext,
+        target_cfg: TargetConfiguration,
+        skybox: &GenericSkyboxDrawData,
+        draw_data: &D,
+        index: u32,
+    ) -> (ViewBuilder, wgpu::CommandBuffer) {
+        let mut view_builder = ViewBuilder::default();
+        view_builder.setup_view(re_ctx, target_cfg).unwrap();
+
+        if self
+            .take_screenshot_next_frame_for_view
+            .map_or(false, |i| i == index)
+        {
+            self.scheduled_screenshots
+                .push(view_builder.schedule_screenshot(re_ctx).unwrap());
+            re_log::info!("Scheduled screenshot for view {}", index);
+        }
+
+        let command_buffer = view_builder
+            .queue_draw(skybox)
+            .queue_draw(draw_data)
+            .draw(re_ctx, Rgba::TRANSPARENT)
+            .unwrap();
+
+        (view_builder, command_buffer)
+    }
+
+    fn handle_incoming_screenshots(&mut self, re_ctx: &mut RenderContext) {
+        re_ctx
+            .gpu_readback_belt
+            .lock()
+            .receive_data(|data, identifier| {
+                if let Some(index) = self
+                    .scheduled_screenshots
+                    .iter()
+                    .position(|s| s.identifier == identifier)
+                {
+                    let screenshot = self.scheduled_screenshots.swap_remove(index);
+
+                    // Need to do a memcpy to remove the padding.
+                    re_log::info!("Received screenshot. Total bytes {:?}", data.len());
+                    let row_info = screenshot.row_info;
+                    let mut buffer = Vec::with_capacity(
+                        (row_info.bytes_per_row_unpadded * screenshot.height) as usize,
+                    );
+                    for row in 0..screenshot.height {
+                        let offset = (row_info.bytes_per_row_padded * row) as usize;
+                        buffer.extend_from_slice(
+                            &data[offset..(offset + row_info.bytes_per_row_unpadded as usize)],
+                        );
+                    }
+
+                    // Get next available file name.
+                    let mut i = 0;
+                    let mut filename = "screenshot.png".to_owned();
+                    while std::path::Path::new(&filename).exists() {
+                        filename = format!("screenshot_{i}.png");
+                        i += 1;
+                    }
+
+                    #[cfg(not(target_arch = "wasm32"))]
+                    image::save_buffer(
+                        filename,
+                        &buffer,
+                        screenshot.width,
+                        screenshot.height,
+                        image::ColorType::Rgba8,
+                    )
+                    .expect("Failed to save screenshot");
+                }
+            });
+    }
+}
+
 impl Example for Multiview {
     fn title() -> &'static str {
         "Multiple Views"
@@ -229,6 +293,9 @@ impl Example for Multiview {
             random_points_positions,
             random_points_radii,
             random_points_colors,
+
+            take_screenshot_next_frame_for_view: None,
+            scheduled_screenshots: Vec::new(),
         }
     }
 
@@ -248,6 +315,8 @@ impl Example for Multiview {
             ) * 10.0;
         }
 
+        self.handle_incoming_screenshots(re_ctx);
+
         let seconds_since_startup = time.seconds_since_startup();
         let view_from_world =
             IsoTransform::look_at_rh(self.camera_position, Vec3::ZERO, Vec3::Y).unwrap();
@@ -295,7 +364,7 @@ impl Example for Multiview {
         #[rustfmt::skip]
         macro_rules! draw {
             ($name:ident @ split #$n:expr) => {{
-                let (view_builder, command_buffer) = draw_view(re_ctx,
+                let (view_builder, command_buffer) = self.draw_view(re_ctx,
                     TargetConfiguration {
                         name: stringify!($name).into(),
                         resolution_in_pixel: splits[$n].resolution_in_pixel,
@@ -305,7 +374,8 @@ impl Example for Multiview {
                         ..Default::default()
                     },
                     &skybox,
-                    &$name
+                    &$name,
+                    $n,
                 );
                 framework::ViewDrawResult {
                     view_builder,
@@ -315,12 +385,16 @@ impl Example for Multiview {
             }};
         }
 
-        vec![
+        let draw_results = vec![
             draw!(triangle @ split #0),
             draw!(lines @ split #1),
             draw!(meshes @ split #2),
             draw!(point_cloud @ split #3),
-        ]
+        ];
+
+        self.take_screenshot_next_frame_for_view = None;
+
+        draw_results
     }
 
     fn on_keyboard_input(&mut self, input: winit::event::KeyboardInput) {
@@ -336,6 +410,19 @@ impl Example for Multiview {
                 };
             }
 
+            (ElementState::Pressed, Some(VirtualKeyCode::Key1)) => {
+                self.take_screenshot_next_frame_for_view = Some(0);
+            }
+            (ElementState::Pressed, Some(VirtualKeyCode::Key2)) => {
+                self.take_screenshot_next_frame_for_view = Some(1);
+            }
+            (ElementState::Pressed, Some(VirtualKeyCode::Key3)) => {
+                self.take_screenshot_next_frame_for_view = Some(2);
+            }
+            (ElementState::Pressed, Some(VirtualKeyCode::Key4)) => {
+                self.take_screenshot_next_frame_for_view = Some(3);
+            }
+
             _ => {}
         }
     }

crates/re_renderer/src/allocator/cpu_write_gpu_read_belt.rs

@@ -303,6 +303,8 @@ impl CpuWriteGpuReadBelt {
     ) -> CpuWriteGpuReadBuffer<T> {
         crate::profile_function!();
 
+        debug_assert!(num_elements > 0, "Cannot allocate zero-sized buffer");
+
         // Potentially overestimate alignment with Self::MIN_ALIGNMENT, see Self::MIN_ALIGNMENT doc string.
         let alignment = (std::mem::align_of::<T>() as wgpu::BufferAddress).max(Self::MIN_ALIGNMENT);
         // Pad out the size of the used buffer to a multiple of Self::MIN_ALIGNMENT.
@@ -425,7 +427,7 @@ impl CpuWriteGpuReadBelt {
 
 impl std::fmt::Debug for CpuWriteGpuReadBelt {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("StagingBelt")
+        f.debug_struct("CpuWriteGpuReadBelt")
             .field("chunk_size", &self.chunk_size)
             .field("active_chunks", &self.active_chunks.len())
             .field("closed_chunks", &self.closed_chunks.len())