Skip to content

Commit

Permalink
Improved readback data handling (rerun-io#1734)
Browse files Browse the repository at this point in the history 
* user or re_renderer is no longer forced to consume all data in one go, instead we'll search if the data you request is available and return it then
    * this leaves the possibility for stale data, which might cause heavy memory leaks, so we clean stale stuff up automatically! (this can happen easily in a healthy application! Think closing a view for which you request picking data)
* identifiers are user-chosen and not required to be unique (space view id is a good & valid identifier!)
* readback data can cary arbitrary user data
* the readback belt itself is a strictly internal datastructure now. Higher level systems like `ScreenshotProcessor` (new! shifting further towards a class-per-draw-pass; this is still an ongoing evolution!) or `PickingLayerProcessor` wrap your userdata and provide everything you need to know about for their readback data
   * this keeps the readback belt agnostic while presenting high level constructs where needed! 
* Do smaller readback chunks
  • Loading branch information
@Wumpf
Wumpf authored Mar 31, 2023
1 parent 118630c commit 6dbc5b9
Show file tree
Hide file tree
Showing 20 changed files with 719 additions and 575 deletions.
103 changes: 46 additions & 57 deletions crates/re_renderer/examples/multiview.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,9 @@ use re_renderer::{
GenericSkyboxDrawData, LineDrawData, LineStripFlags, MeshDrawData, MeshInstance,
TestTriangleDrawData,
},
view_builder::{
OrthographicCameraMode, Projection, ScheduledScreenshot, TargetConfiguration, ViewBuilder,
},
Color32, LineStripSeriesBuilder, PointCloudBuilder, RenderContext, Rgba, Size,
view_builder::{OrthographicCameraMode, Projection, TargetConfiguration, ViewBuilder},
Color32, GpuReadbackIdentifier, LineStripSeriesBuilder, PointCloudBuilder, RenderContext, Rgba,
ScreenshotProcessor, Size,
};
use winit::event::{ElementState, VirtualKeyCode};

Expand Down Expand Up @@ -151,7 +150,6 @@ struct Multiview {
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 {
Expand All @@ -164,6 +162,45 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
.into()
}

/// Readback identifier for screenshots.
/// Identifiers don't need to be unique and we don't have anything interesting to distinguish here!
const READBACK_IDENTIFIER: GpuReadbackIdentifier = 0;

fn handle_incoming_screenshots(re_ctx: &RenderContext) {
ScreenshotProcessor::next_readback_result(
re_ctx,
READBACK_IDENTIFIER,
|data, _extent, view_idx: u32| {
re_log::info!(
"Received screenshot for view {view_idx}. Total bytes {:?}",
data.len()
);

#[cfg(not(target_arch = "wasm32"))]
{
// Get next available file name.
let mut i = 1;
let filename = loop {
let filename = format!("screenshot_{i}.png");
if !std::path::Path::new(&filename).exists() {
break filename;
}
i += 1;
};

image::save_buffer(
filename,
data,
_extent.x,
_extent.y,
image::ColorType::Rgba8,
)
.expect("Failed to save screenshot");
}
},
);
}

impl Multiview {
fn draw_view<D: 'static + re_renderer::renderer::DrawData + Sync + Send + Clone>(
&mut self,
Expand All @@ -180,8 +217,9 @@ impl Multiview {
.take_screenshot_next_frame_for_view
.map_or(false, |i| i == index)
{
self.scheduled_screenshots
.push(view_builder.schedule_screenshot(re_ctx).unwrap());
view_builder
.schedule_screenshot(re_ctx, READBACK_IDENTIFIER, index)
.unwrap();
re_log::info!("Scheduled screenshot for view {}", index);
}

Expand All @@ -193,54 +231,6 @@ impl Multiview {

(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)
{
re_log::info!(
"Received screenshot. Total bytes {:?}. Identifier {identifier}",
data.len()
);

#[cfg(target_arch = "wasm32")]
self.scheduled_screenshots.remove(index);

#[cfg(not(target_arch = "wasm32"))]
{
let screenshot = self.scheduled_screenshots.swap_remove(index);

re_log::info!("Received screenshot. Total bytes {}", data.len());

// Get next available file name.
let mut i = 1;
let filename = loop {
let filename = format!("screenshot_{i}.png");
if !std::path::Path::new(&filename).exists() {
break filename;
}
i += 1;
};

#[cfg(not(target_arch = "wasm32"))]
image::save_buffer(
filename,
&screenshot.row_info.remove_padding(data),
screenshot.extent.x,
screenshot.extent.y,
image::ColorType::Rgba8,
)
.expect("Failed to save screenshot");
}
}
});
}
}

impl Example for Multiview {
Expand Down Expand Up @@ -297,7 +287,6 @@ impl Example for Multiview {
random_points_colors,

take_screenshot_next_frame_for_view: None,
scheduled_screenshots: Vec::new(),
}
}

Expand All @@ -317,7 +306,7 @@ impl Example for Multiview {
) * 10.0;
}

self.handle_incoming_screenshots(re_ctx);
handle_incoming_screenshots(re_ctx);

let seconds_since_startup = time.seconds_since_startup();
let view_from_world =
Expand Down
75 changes: 26 additions & 49 deletions crates/re_renderer/examples/picking.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
use ahash::HashMap;
use itertools::Itertools as _;
use rand::Rng;
use re_renderer::{
view_builder::{Projection, TargetConfiguration, ViewBuilder},
Color32, GpuReadbackBufferIdentifier, IntRect, PickingLayerId, PickingLayerInstanceId,
PointCloudBuilder, RenderContext, ScheduledPickingRect, Size,
Color32, GpuReadbackIdentifier, IntRect, PickingLayerInstanceId, PickingLayerProcessor,
PointCloudBuilder, Size,
};

mod framework;
Expand All @@ -18,7 +17,6 @@ struct PointSet {

struct Picking {
point_sets: Vec<PointSet>,
scheduled_picking_rects: HashMap<GpuReadbackBufferIdentifier, ScheduledPickingRect>,
picking_position: glam::UVec2,
}

Expand All @@ -32,36 +30,9 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
.into()
}

impl Picking {
#[allow(clippy::unused_self)]
fn handle_incoming_picking_data(&mut self, re_ctx: &mut RenderContext, _time: f32) {
re_ctx
.gpu_readback_belt
.lock()
.receive_data(|data, identifier| {
if let Some(picking_rect_info) = self.scheduled_picking_rects.remove(&identifier) {
// TODO(andreas): Move this into a utility function?
let picking_data_without_padding =
picking_rect_info.row_info.remove_padding(data);
let picking_data: &[PickingLayerId] =
bytemuck::cast_slice(&picking_data_without_padding);

// Grab the middle pixel. usually we'd want to do something clever that snaps the the closest object of interest.
let picked_pixel = picking_data[(picking_rect_info.rect.extent.x / 2
+ (picking_rect_info.rect.extent.y / 2) * picking_rect_info.rect.extent.x)
as usize];

if picked_pixel.object.0 != 0 {
let point_set = &mut self.point_sets[picked_pixel.object.0 as usize - 1];
point_set.radii[picked_pixel.instance.0 as usize] = Size::new_scene(0.1);
point_set.colors[picked_pixel.instance.0 as usize] = Color32::DEBUG_COLOR;
}
} else {
re_log::error!("Received picking data for unknown identifier");
}
});
}
}
/// Readback identifier for picking rects.
/// Identifiers don't need to be unique and we don't have anything interesting to distinguish here!
const READBACK_IDENTIFIER: GpuReadbackIdentifier = 0;

impl framework::Example for Picking {
fn title() -> &'static str {
Expand Down Expand Up @@ -103,7 +74,6 @@ impl framework::Example for Picking {

Picking {
point_sets,
scheduled_picking_rects: HashMap::default(),
picking_position: glam::UVec2::ZERO,
}
}
Expand All @@ -112,10 +82,23 @@ impl framework::Example for Picking {
&mut self,
re_ctx: &mut re_renderer::RenderContext,
resolution: [u32; 2],
time: &framework::Time,
_time: &framework::Time,
pixels_from_point: f32,
) -> Vec<framework::ViewDrawResult> {
self.handle_incoming_picking_data(re_ctx, time.seconds_since_startup());
while let Some(picking_result) =
PickingLayerProcessor::next_readback_result::<()>(re_ctx, READBACK_IDENTIFIER)
{
// Grab the middle pixel. usually we'd want to do something clever that snaps the the closest object of interest.
let picked_pixel = picking_result.picking_data[(picking_result.rect.extent.x / 2
+ (picking_result.rect.extent.y / 2) * picking_result.rect.extent.x)
as usize];

if picked_pixel.object.0 != 0 {
let point_set = &mut self.point_sets[picked_pixel.object.0 as usize - 1];
point_set.radii[picked_pixel.instance.0 as usize] = Size::new_scene(0.1);
point_set.colors[picked_pixel.instance.0 as usize] = Color32::DEBUG_COLOR;
}
}

let mut view_builder = ViewBuilder::default();

Expand Down Expand Up @@ -148,19 +131,13 @@ impl framework::Example for Picking {
// Use an uneven number of pixels for the picking rect so that there is a clearly defined middle-pixel.
// (for this sample a size of 1 would be sufficient, but for a real application you'd want to use a larger size to allow snapping)
let picking_rect_size = 31;

let picking_rect = view_builder
.schedule_picking_readback(
re_ctx,
IntRect::from_middle_and_extent(
self.picking_position.as_ivec2(),
glam::uvec2(picking_rect_size, picking_rect_size),
),
false,
)
let picking_rect = IntRect::from_middle_and_extent(
self.picking_position.as_ivec2(),
glam::uvec2(picking_rect_size, picking_rect_size),
);
view_builder
.schedule_picking_rect(re_ctx, picking_rect, READBACK_IDENTIFIER, (), false)
.unwrap();
self.scheduled_picking_rects
.insert(picking_rect.identifier, picking_rect);

let mut builder = PointCloudBuilder::<()>::new(re_ctx);

Expand Down
Loading

0 comments on commit 6dbc5b9

Please sign in to comment.