Skip to content

Commit

Permalink
Optimize gathering of point cloud colors (#3730)
Browse files Browse the repository at this point in the history
Saves around 2ms of native CPU-time on OPF. On native this is mostly
hidden by the parallelism, but on web it is a real saving.

Found by looking at the in-browser profiler.

### What

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/3730) (if
applicable)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG

- [PR Build Summary](https://build.rerun.io/pr/3730)
- [Docs
preview](https://rerun.io/preview/a55c07912e4d24103187b05e34ffa4dc4a606da2/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/a55c07912e4d24103187b05e34ffa4dc4a606da2/examples)
<!--EXAMPLES-PREVIEW-->
- [Recent benchmark results](https://ref.rerun.io/dev/bench/)
- [Wasm size tracking](https://ref.rerun.io/dev/sizes/)
  • Loading branch information
emilk authored Oct 10, 2023
1 parent 3dfc221 commit 1fe15f0
Show file tree
Hide file tree
Showing 13 changed files with 428 additions and 156 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions crates/re_arrow_store/benches/arrow2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use std::sync::Arc;

use arrow2::array::{Array, PrimitiveArray, StructArray, UnionArray};
use criterion::{criterion_group, Criterion};
use criterion::Criterion;
use itertools::Itertools;

use re_log_types::{DataCell, SizeBytes as _};
Expand All @@ -19,7 +19,7 @@ use re_types::{

// ---

criterion_group!(benches, erased_clone, estimated_size_bytes);
criterion::criterion_group!(benches, erased_clone, estimated_size_bytes);

#[cfg(not(feature = "core_benchmarks_only"))]
criterion::criterion_main!(benches);
Expand Down
10 changes: 9 additions & 1 deletion crates/re_log_types/src/data_cell.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::sync::Arc;

use arrow2::datatypes::DataType;
use re_types::{Component, ComponentName, DeserializationError};
use re_types::{Component, ComponentBatch, ComponentName, DeserializationError};

use crate::SizeBytes;

Expand Down Expand Up @@ -164,6 +164,14 @@ pub struct DataCellInner {
// TODO(#1696): Check that the array is indeed a leaf / component type when building a cell from an
// arrow payload.
impl DataCell {
/// Builds a new `DataCell` from a component batch.
#[inline]
pub fn from_component_batch(batch: &dyn ComponentBatch) -> re_types::SerializationResult<Self> {
batch
.to_arrow()
.map(|arrow| DataCell::from_arrow(batch.name(), arrow))
}

/// Builds a new `DataCell` from a uniform iterable of native component values.
///
/// Fails if the given iterable cannot be serialized to arrow, which should never happen when
Expand Down
44 changes: 43 additions & 1 deletion crates/re_log_types/src/data_row.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use ahash::HashSetExt;
use nohash_hasher::IntSet;
use re_types::ComponentName;
use re_types::{AsComponents, ComponentName};
use smallvec::SmallVec;

use crate::{DataCell, DataCellError, DataTable, EntityPath, SizeBytes, TableId, TimePoint};
Expand Down Expand Up @@ -266,6 +266,48 @@ pub struct DataRow {
}

impl DataRow {
/// Builds a new `DataRow` from anything implementing [`AsComponents`].
pub fn from_archetype(
row_id: RowId,
timepoint: TimePoint,
entity_path: EntityPath,
as_components: &dyn AsComponents,
) -> anyhow::Result<Self> {
re_tracing::profile_function!();

let batches = as_components.as_component_batches();
Self::from_component_batches(
row_id,
timepoint,
entity_path,
batches.iter().map(|batch| batch.as_ref()),
)
}

/// Builds a new `DataRow` from anything implementing [`AsComponents`].
pub fn from_component_batches<'a>(
row_id: RowId,
timepoint: TimePoint,
entity_path: EntityPath,
comp_batches: impl IntoIterator<Item = &'a dyn re_types::ComponentBatch>,
) -> anyhow::Result<Self> {
re_tracing::profile_function!();

let data_cells = comp_batches
.into_iter()
.map(|batch| DataCell::from_component_batch(batch))
.collect::<Result<Vec<DataCell>, _>>()?;

// TODO(emilk): should `DataRow::from_cells` calculate `num_instances` instead?
let num_instances = data_cells.iter().map(|cell| cell.num_instances()).max();
let num_instances = num_instances.unwrap_or(0);

let mut row =
DataRow::from_cells(row_id, timepoint, entity_path, num_instances, data_cells)?;
row.compute_all_size_bytes();
Ok(row)
}

/// Builds a new `DataRow` from an iterable of [`DataCell`]s.
///
/// Fails if:
Expand Down
12 changes: 12 additions & 0 deletions crates/re_space_view_spatial/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,15 @@ parking_lot.workspace = true
rayon.workspace = true
serde = "1"
smallvec = { workspace = true, features = ["serde"] }


[dev-dependencies]
criterion.workspace = true
mimalloc.workspace = true

[lib]
bench = false

[[bench]]
name = "bench_points"
harness = false
134 changes: 134 additions & 0 deletions crates/re_space_view_spatial/benches/bench_points.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
//! High-level benchmark of the CPU-side of our `Points3D` rendering.
use re_arrow_store::{DataStore, LatestAtQuery};
use re_log_types::{DataRow, EntityPath, RowId, TimeInt, TimePoint, Timeline};
use re_space_view_spatial::LoadedPoints;
use re_types::{
archetypes::Points3D,
components::{Color, InstanceKey, Position3D},
Loggable as _,
};
use re_viewer_context::Annotations;

#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

criterion::criterion_main!(benches);
criterion::criterion_group!(benches, bench_points);

// ---

#[cfg(not(debug_assertions))]
const NUM_POINTS: usize = 1_000_000;

// `cargo test` also runs the benchmark setup code, so make sure they run quickly:
#[cfg(debug_assertions)]
const NUM_POINTS: usize = 10;

// ---

/// Mimics `examples/python/open_photogrammetry_format/main.py`
fn bench_points(c: &mut criterion::Criterion) {
let timeline = Timeline::log_time();
let ent_path = EntityPath::from("points");

let store = {
let mut store = DataStore::new(InstanceKey::name(), Default::default());

let positions = vec![Position3D::new(0.1, 0.2, 0.3); NUM_POINTS];
let colors = vec![Color::from(0xffffffff); NUM_POINTS];
let points = Points3D::new(positions).with_colors(colors);
let mut timepoint = TimePoint::default();
timepoint.insert(timeline, TimeInt::from_seconds(0));
let data_row =
DataRow::from_archetype(RowId::random(), timepoint, ent_path.clone(), &points).unwrap();
store.insert_row(&data_row).unwrap();
store
};

let latest_at = LatestAtQuery::latest(timeline);
let annotations = Annotations::missing();

{
let mut group = c.benchmark_group("Points3D");
group.bench_function("query_archetype", |b| {
b.iter(|| {
let arch_view =
re_query::query_archetype::<Points3D>(&store, &latest_at, &ent_path).unwrap();
assert_eq!(arch_view.num_instances(), NUM_POINTS);
arch_view
});
});
}

let arch_view = re_query::query_archetype::<Points3D>(&store, &latest_at, &ent_path).unwrap();
assert_eq!(arch_view.num_instances(), NUM_POINTS);

{
let mut group = c.benchmark_group("Points3D");
group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
group.bench_function("load_all", |b| {
b.iter(|| {
let points =
LoadedPoints::load(&arch_view, &ent_path, latest_at.at, &annotations).unwrap();
assert_eq!(points.positions.len(), NUM_POINTS);
assert_eq!(points.colors.len(), NUM_POINTS);
assert_eq!(points.radii.len(), NUM_POINTS); // NOTE: we don't log radii, but we should get a list of defaults!
points
});
});
}

{
let mut group = c.benchmark_group("Points3D");
group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
group.bench_function("load_positions", |b| {
b.iter(|| {
let positions = LoadedPoints::load_positions(&arch_view).unwrap();
assert_eq!(positions.len(), NUM_POINTS);
positions
});
});
}

{
let points = LoadedPoints::load(&arch_view, &ent_path, latest_at.at, &annotations).unwrap();

let mut group = c.benchmark_group("Points3D");
group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
group.bench_function("load_colors", |b| {
b.iter(|| {
let colors =
LoadedPoints::load_colors(&arch_view, &ent_path, &points.annotation_infos)
.unwrap();
assert_eq!(colors.len(), NUM_POINTS);
colors
});
});
}

// NOTE: we don't log radii!
{
let mut group = c.benchmark_group("Points3D");
group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
group.bench_function("load_radii", |b| {
b.iter(|| {
let radii = LoadedPoints::load_radii(&arch_view, &ent_path).unwrap();
assert_eq!(radii.len(), NUM_POINTS);
radii
});
});
}

{
let mut group = c.benchmark_group("Points3D");
group.throughput(criterion::Throughput::Elements(NUM_POINTS as _));
group.bench_function("load_picking_ids", |b| {
b.iter(|| {
let picking_ids = LoadedPoints::load_picking_ids(&arch_view);
assert_eq!(picking_ids.len(), NUM_POINTS);
picking_ids
});
});
}
}
3 changes: 3 additions & 0 deletions crates/re_space_view_spatial/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ mod ui_3d;
pub use space_view_2d::SpatialSpaceView2D;
pub use space_view_3d::SpatialSpaceView3D;

#[doc(hidden)] // Public for benchmarks
pub use parts::LoadedPoints;

// ---

mod view_kind {
Expand Down
15 changes: 3 additions & 12 deletions crates/re_space_view_spatial/src/parts/images.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,10 +260,7 @@ impl ImagesPart {
.annotations
.resolved_class_description(None)
.annotation_info()
.color(
color.map(|c| c.to_array()).as_ref(),
DefaultColor::OpaqueWhite,
);
.color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);

if let Some(textured_rect) = to_textured_rect(
ctx,
Expand Down Expand Up @@ -388,10 +385,7 @@ impl ImagesPart {
.annotations
.resolved_class_description(None)
.annotation_info()
.color(
color.map(|c| c.to_array()).as_ref(),
DefaultColor::OpaqueWhite,
);
.color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);

if let Some(textured_rect) = to_textured_rect(
ctx,
Expand Down Expand Up @@ -484,10 +478,7 @@ impl ImagesPart {
.annotations
.resolved_class_description(None)
.annotation_info()
.color(
color.map(|c| c.to_array()).as_ref(),
DefaultColor::OpaqueWhite,
);
.color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite);

if let Some(textured_rect) = to_textured_rect(
ctx,
Expand Down
20 changes: 10 additions & 10 deletions crates/re_space_view_spatial/src/parts/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@ use re_types::components::Text;
pub use spatial_view_part::SpatialViewPartData;
pub use transform3d_arrows::{add_axis_arrows, Transform3DArrowsPart};

#[doc(hidden)] // Public for benchmarks
pub use points3d::LoadedPoints;

use ahash::HashMap;
use std::sync::Arc;

use re_data_store::{EntityPath, InstancePathHash};
use re_types::components::{Color, InstanceKey};
Expand Down Expand Up @@ -118,7 +120,7 @@ pub fn process_colors<'a, A: Archetype>(
arch_view.iter_optional_component::<Color>()?,
)
.map(move |(annotation_info, color)| {
annotation_info.color(color.map(move |c| c.to_array()).as_ref(), default_color)
annotation_info.color(color.map(|c| c.to_array()), default_color)
}))
}

Expand All @@ -134,9 +136,7 @@ pub fn process_labels<'a, A: Archetype>(
annotation_infos.iter(),
arch_view.iter_optional_component::<Text>()?,
)
.map(move |(annotation_info, text)| {
annotation_info.label(text.as_ref().map(move |t| t.as_str()))
}))
.map(move |(annotation_info, text)| annotation_info.label(text.as_ref().map(|t| t.as_str()))))
}

/// Process [`re_types::components::Radius`] components to [`re_renderer::Size`] using auto size
Expand Down Expand Up @@ -171,22 +171,22 @@ pub fn process_radii<'a, A: Archetype>(
fn process_annotations<Primary, A: Archetype>(
query: &ViewQuery<'_>,
arch_view: &re_query::ArchetypeView<A>,
annotations: &Arc<Annotations>,
annotations: &Annotations,
) -> Result<ResolvedAnnotationInfos, re_query::QueryError>
where
Primary: re_types::Component + Clone,
{
process_annotations_and_keypoints(query, arch_view, annotations, |_: &Primary| {
process_annotations_and_keypoints(query.latest_at, arch_view, annotations, |_: &Primary| {
glam::Vec3::ZERO
})
.map(|(a, _)| a)
}

/// Resolves all annotations and keypoints for the given entity view.
fn process_annotations_and_keypoints<Primary, A: Archetype>(
query: &ViewQuery<'_>,
latest_at: re_log_types::TimeInt,
arch_view: &re_query::ArchetypeView<A>,
annotations: &Arc<Annotations>,
annotations: &Annotations,
mut primary_into_position: impl FnMut(&Primary) -> glam::Vec3,
) -> Result<(ResolvedAnnotationInfos, Keypoints), re_query::QueryError>
where
Expand Down Expand Up @@ -220,7 +220,7 @@ where

if let (Some(keypoint_id), Some(class_id), primary) = (keypoint_id, class_id, primary) {
keypoints
.entry((class_id, query.latest_at.as_i64()))
.entry((class_id, latest_at.as_i64()))
.or_default()
.insert(keypoint_id.0, primary_into_position(&primary));
class_description.annotation_info_with_keypoint(keypoint_id.0)
Expand Down
Loading

0 comments on commit 1fe15f0

Please sign in to comment.