From 2b7a508ad7ac4bc3c886aaa6fb4a5233dbdb0310 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 28 Aug 2024 10:34:23 -0400 Subject: [PATCH] Speed up handling of large numbers of transform entities (#7300) ### What Resolves: https://github.com/rerun-io/rerun/issues/7292 Halves the time spent in `execute_systems` when using many Transforms. - TransformContext uses a store subscriper to track which components are non-empty and only queries those from the store. - Arrows visualizer no longer queries for overrides on the unused transform components since those are fetched from the transform context manager anyways. Before: ![image](https://github.com/user-attachments/assets/90137375-087a-47ee-8e76-6cb13b46f4cf) After: ![image](https://github.com/user-attachments/assets/beab559a-2d72-49da-8403-48c057e2a811) ### 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 the web demo (if applicable): * Using examples from latest `main` build: [rerun.io/viewer](https://rerun.io/viewer/pr/7300?manifest_url=https://app.rerun.io/version/main/examples_manifest.json) * Using full set of examples from `nightly` build: [rerun.io/viewer](https://rerun.io/viewer/pr/7300?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG * [x] If applicable, add a new check to the [release checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)! * [x] If have noted any breaking changes to the log API in `CHANGELOG.md` and the migration guide - [PR Build Summary](https://build.rerun.io/pr/7300) - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) To run all checks from `main`, comment on the PR with `@rerun-bot full-check`. --- .../src/contexts/transform_context.rs | 51 ++++++------------- .../src/transform_component_tracker.rs | 51 ++++++++++++++----- .../src/visualizers/transform3d_arrows.rs | 16 ++++-- 3 files changed, 67 insertions(+), 51 deletions(-) diff --git a/crates/viewer/re_space_view_spatial/src/contexts/transform_context.rs b/crates/viewer/re_space_view_spatial/src/contexts/transform_context.rs index f5ba0c709e28..ed9e6e171f39 100644 --- a/crates/viewer/re_space_view_spatial/src/contexts/transform_context.rs +++ b/crates/viewer/re_space_view_spatial/src/contexts/transform_context.rs @@ -530,27 +530,17 @@ fn query_and_resolve_tree_transform_at_entity( entity_db: &EntityDb, query: &LatestAtQuery, ) -> Option { - if !TransformComponentTracker::access(entity_db.store_id(), |tracker| { - tracker.is_potentially_transformed_transform3d(entity_path) - }) - .unwrap_or(false) - { + let Some(transform3d_components) = + TransformComponentTracker::access(entity_db.store_id(), |tracker| { + tracker.transform3d_components(entity_path).cloned() + }) + .flatten() + else { return None; - } + }; // TODO(#6743): Doesn't take into account overrides. - let result = entity_db.latest_at( - query, - entity_path, - [ - Translation3D::name(), - RotationAxisAngle::name(), - RotationQuat::name(), - Scale3D::name(), - TransformMat3x3::name(), - TransformRelation::name(), - ], - ); + let result = entity_db.latest_at(query, entity_path, transform3d_components); if result.components.is_empty() { return None; } @@ -586,26 +576,17 @@ fn query_and_resolve_instance_poses_at_entity( entity_db: &EntityDb, query: &LatestAtQuery, ) -> Vec { - if !TransformComponentTracker::access(entity_db.store_id(), |tracker| { - tracker.is_potentially_transformed_pose3d(entity_path) - }) - .unwrap_or(false) - { + let Some(pose3d_components) = + TransformComponentTracker::access(entity_db.store_id(), |tracker| { + tracker.pose3d_components(entity_path).cloned() + }) + .flatten() + else { return Vec::new(); - } + }; // TODO(#6743): Doesn't take into account overrides. - let result = entity_db.latest_at( - query, - entity_path, - [ - PoseTranslation3D::name(), - PoseRotationAxisAngle::name(), - PoseRotationQuat::name(), - PoseScale3D::name(), - PoseTransformMat3x3::name(), - ], - ); + let result = entity_db.latest_at(query, entity_path, pose3d_components); let max_count = result .components diff --git a/crates/viewer/re_space_view_spatial/src/transform_component_tracker.rs b/crates/viewer/re_space_view_spatial/src/transform_component_tracker.rs index 1174aa99b3cc..4e13e9c95217 100644 --- a/crates/viewer/re_space_view_spatial/src/transform_component_tracker.rs +++ b/crates/viewer/re_space_view_spatial/src/transform_component_tracker.rs @@ -1,7 +1,7 @@ use ahash::HashMap; use once_cell::sync::OnceCell; -use nohash_hasher::IntSet; +use nohash_hasher::{IntMap, IntSet}; use re_chunk_store::{ ChunkStore, ChunkStoreDiffKind, ChunkStoreEvent, ChunkStoreSubscriber, ChunkStoreSubscriberHandle, @@ -20,11 +20,13 @@ use re_types::ComponentName; /// This is a huge performance improvement in practice, especially in recordings with many entities. #[derive(Default)] pub struct TransformComponentTracker { - /// Which entities have had any `Transform3D` component at any point in time. - transform3d_entities: IntSet, + /// Which entities have had any `Transform3D` component at any point in time, and which + /// components they actually make use of. + transform3d_entities: IntMap>, - /// Which entities have had any `InstancePoses3D` components at any point in time. - pose3d_entities: IntSet, + /// Which entities have had any `InstancePoses3D` components at any point in time, and + /// which components they actually make use of. + pose3d_entities: IntMap>, } impl TransformComponentTracker { @@ -41,13 +43,16 @@ impl TransformComponentTracker { } #[inline] - pub fn is_potentially_transformed_transform3d(&self, entity_path: &EntityPath) -> bool { - self.transform3d_entities.contains(&entity_path.hash()) + pub fn transform3d_components( + &self, + entity_path: &EntityPath, + ) -> Option<&IntSet> { + self.transform3d_entities.get(&entity_path.hash()) } #[inline] - pub fn is_potentially_transformed_pose3d(&self, entity_path: &EntityPath) -> bool { - self.pose3d_entities.contains(&entity_path.hash()) + pub fn pose3d_components(&self, entity_path: &EntityPath) -> Option<&IntSet> { + self.pose3d_entities.get(&entity_path.hash()) } } @@ -119,15 +124,35 @@ impl ChunkStoreSubscriber for TransformComponentTrackerStoreSubscriber { let entity_path_hash = event.chunk.entity_path().hash(); for component_name in event.chunk.component_names() { - if self.transform_components.contains(&component_name) { + if self.transform_components.contains(&component_name) + && event + .chunk + .components() + .get(&component_name) + .map_or(false, |list_array| { + list_array.offsets().lengths().any(|len| len > 0) + }) + { transform_component_tracker .transform3d_entities - .insert(entity_path_hash); + .entry(entity_path_hash) + .or_default() + .insert(component_name); } - if self.pose_components.contains(&component_name) { + if self.pose_components.contains(&component_name) + && event + .chunk + .components() + .get(&component_name) + .map_or(false, |list_array| { + list_array.offsets().lengths().any(|len| len > 0) + }) + { transform_component_tracker .pose3d_entities - .insert(entity_path_hash); + .entry(entity_path_hash) + .or_default() + .insert(component_name); } } } diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/transform3d_arrows.rs b/crates/viewer/re_space_view_spatial/src/visualizers/transform3d_arrows.rs index 758e0843ac77..61a0ca8a3223 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/transform3d_arrows.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/transform3d_arrows.rs @@ -1,7 +1,7 @@ use egui::Color32; use nohash_hasher::IntSet; use re_log_types::{EntityPath, Instance}; -use re_space_view::DataResultQuery; +use re_space_view::{latest_at_with_blueprint_resolved_data, DataResultQuery}; use re_types::{ archetypes::{Pinhole, Transform3D}, components::{AxisLength, ImagePlaneDistance}, @@ -122,8 +122,18 @@ impl VisualizerSystem for Transform3DArrowsVisualizer { .single_entity_transform_required(&data_result.entity_path, "Transform3DArrows") }; - let results = data_result - .latest_at_with_blueprint_resolved_data::(ctx, &latest_at_query); + // Note, we use this interface instead of `data_result.latest_at_with_blueprint_resolved_data` to avoid querying + // for a bunch of unused components. The actual transform data comes out of the context manager and can't be + // overridden via blueprint anyways. + let results = latest_at_with_blueprint_resolved_data( + ctx, + None, + &latest_at_query, + data_result, + std::iter::once(AxisLength::name()), + false, + ); + let axis_length: f32 = results.get_mono_with_fallback::().into(); if axis_length == 0.0 {