Improve performance for many entities (#3078)

### What * Fixes #2285 (.. at least fixes it well enough for the moment I reckon) Improves both the per space view rendering and the space view adding heuristic. In particular on the later there's still significant gains to be made. Most importantly we should be scaling now much better with number of added systems. Went fairly early with the decision to have the primary datastructure to pass around be a `Map<System, Vec<Entity>>` (conceptual) instead of the other way round since this is what we need in the systems most of the time. This is another step towards a stronger contract of systems specifying what components they will query ahead of time! The way it is implemented on thr (rename from `BlueprintTree` 💥) `SpaceViewContents` also paves the way for ui selection of systems for a given entity path which complements the ongoing work on our new data ingestion interfaces. ----- Testcase: `examples/python/open_photogrammetry_format/main.py --no-frames`, **fully reset viewer** and then hiding **NOT REMOVING** the points (the world/pcl entity) and deselecting them. (The bold marked pieces are very important as they have an influence on what heuristics run - see #3077) Before: ![image](https://github.com/rerun-io/rerun/assets/1220815/cda9d555-8839-4d78-8507-cf4ed34bad70) After: ![image](https://github.com/rerun-io/rerun/assets/1220815/d34a00f8-05af-4160-85bf-b709a0d36ecb) (release runs, averaged over a bunch of frames on my M1 Max) Highlights: * `AppState::show`: 24.0ms ➡️ 14.6ms * `SpaceViewBlueprint::scene_ui`: 4.9ms ➡️ 2.6ms * this was the original primary optimization target! * `Viewport::on_frame_start`: 15.2ms ➡️ 7.4ms * `default_created_space_view`: 12.7ms ➡️ 5.1ms * quite some time spent on this, but we want to do a paradigm shift there: #3079 (slight number discrepancies from the screenshot are due to doing a different run) ---- Draft todo: Code sanity check. Suspecting some heuristics might be slightly broken, need to go through examples ### 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/3078) (if applicable) - [PR Build Summary](https://build.rerun.io/pr/3078) - [Docs preview](https://rerun.io/preview/fc5bb0b8908e2e8c666a6d5daa1dee105fe2aae8/docs)  - [Examples preview](https://rerun.io/preview/fc5bb0b8908e2e8c666a6d5daa1dee105fe2aae8/examples)  - [Recent benchmark results](https://ref.rerun.io/dev/bench/) - [Wasm size tracking](https://ref.rerun.io/dev/sizes/) --------- Co-authored-by: Jeremy Leibs <[email protected]>
rerun-io · Aug 29, 2023 · 41d2cbd · 41d2cbd
1 parent 9fa097b
commit 41d2cbd
Show file tree

Hide file tree

Showing 53 changed files with 978 additions and 541 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/re_space_view/src/lib.rs b/crates/re_space_view/src/lib.rs
@@ -3,10 +3,10 @@
 //! Types & utilities for defining Space View classes and communicating with the Viewport.
 
 pub mod controls;
-mod data_blueprint;
 mod screenshot;
+mod space_view_contents;
 mod unreachable_transform_reason;
 
-pub use data_blueprint::{DataBlueprintGroup, DataBlueprintTree};
 pub use screenshot::ScreenshotMode;
+pub use space_view_contents::{DataBlueprintGroup, SpaceViewContents};
 pub use unreachable_transform_reason::UnreachableTransformReason;
diff --git a/crates/re_space_view/src/data_blueprint.rs → .../re_space_view/src/space_view_contents.rs b/crates/re_space_view/src/data_blueprint.rs → .../re_space_view/src/space_view_contents.rs
@@ -1,8 +1,8 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 
-use nohash_hasher::{IntMap, IntSet};
+use nohash_hasher::IntMap;
 use re_data_store::{EntityPath, EntityProperties, EntityPropertyMap};
-use re_viewer_context::DataBlueprintGroupHandle;
+use re_viewer_context::{DataBlueprintGroupHandle, PerSystemEntities};
 use slotmap::SlotMap;
 use smallvec::{smallvec, SmallVec};
 
@@ -79,7 +79,7 @@ impl DataBlueprints {
 
 /// Tree of all data blueprint groups for a single space view.
 #[derive(Clone, serde::Deserialize, serde::Serialize)]
-pub struct DataBlueprintTree {
+pub struct SpaceViewContents {
     /// All data blueprint groups.
     groups: SlotMap<DataBlueprintGroupHandle, DataBlueprintGroup>,
 
@@ -92,11 +92,12 @@ pub struct DataBlueprintTree {
 
     /// List of all entities that we query via this data blueprint collection.
     ///
-    /// Two things to keep in sync:
-    /// * children on [`DataBlueprintGroup`] this is on
-    /// * elements in [`Self::path_to_group`]
-    /// TODO(andreas): Can we reduce the amount of these dependencies?
-    entity_paths: IntSet<EntityPath>,
+    /// Currently this is reset every frame in `SpaceViewBlueprint::reset_systems_per_entity_path`.
+    /// In the future, we may want to keep this around and only add/remove systems
+    /// for entities. But at this point we'd likely handle the heuristics a bit differently as well
+    /// and don't use serde here for serialization.
+    #[serde(skip)]
+    per_system_entity_list: PerSystemEntities,
 
     /// Root group, always exists as a placeholder
     root_group_handle: DataBlueprintGroupHandle,
@@ -105,12 +106,12 @@ pub struct DataBlueprintTree {
 }
 
 /// Determine whether this `DataBlueprintTree` has user-edits relative to another `DataBlueprintTree`
-impl DataBlueprintTree {
+impl SpaceViewContents {
     pub fn has_edits(&self, other: &Self) -> bool {
         let Self {
             groups,
             path_to_group,
-            entity_paths,
+            per_system_entity_list: _,
             root_group_handle,
             data_blueprints,
         } = self;
@@ -123,13 +124,12 @@ impl DataBlueprintTree {
                     .map_or(true, |other_val| val.has_edits(other_val))
             })
             || *path_to_group != other.path_to_group
-            || *entity_paths != other.entity_paths
             || *root_group_handle != other.root_group_handle
             || data_blueprints.has_edits(&other.data_blueprints)
     }
 }
 
-impl Default for DataBlueprintTree {
+impl Default for SpaceViewContents {
     fn default() -> Self {
         let mut groups = SlotMap::default();
         let root_group = groups.insert(DataBlueprintGroup::default());
@@ -140,14 +140,14 @@ impl Default for DataBlueprintTree {
         Self {
             groups,
             path_to_group: path_to_blueprint,
-            entity_paths: IntSet::default(),
+            per_system_entity_list: BTreeMap::default(),
             root_group_handle: root_group,
             data_blueprints: DataBlueprints::default(),
         }
     }
 }
 
-impl DataBlueprintTree {
+impl SpaceViewContents {
     /// Returns a handle to the root data blueprint.
     ///
     /// Even if there are no other groups, we always have a root group at the top.
@@ -207,12 +207,47 @@ impl DataBlueprintTree {
     }
 
     pub fn contains_entity(&self, path: &EntityPath) -> bool {
-        self.entity_paths.contains(path)
+        // If an entity is in path_to_group it is *likely* also an entity in the Space View.
+        // However, it could be that the path *only* refers to a group, not also an entity.
+        // So once we resolved the group, we need to check if it contains the entity of interest.
+        self.path_to_group
+            .get(path)
+            .and_then(|group| {
+                self.groups
+                    .get(*group)
+                    .and_then(|group| group.entities.get(path))
+            })
+            .is_some()
     }
 
     /// List of all entities that we query via this data blueprint collection.
-    pub fn entity_paths(&self) -> &IntSet<EntityPath> {
-        &self.entity_paths
+    pub fn entity_paths(&self) -> impl Iterator<Item = &EntityPath> {
+        // Each entity is only ever in one group, therefore collecting all entities from all groups, gives us all entities.
+        self.groups.values().flat_map(|group| group.entities.iter())
+    }
+
+    pub fn per_system_entities(&self) -> &PerSystemEntities {
+        &self.per_system_entity_list
+    }
+
+    pub fn per_system_entities_mut(&mut self) -> &mut PerSystemEntities {
+        &mut self.per_system_entity_list
+    }
+
+    pub fn contains_all_entities_from(&self, other: &SpaceViewContents) -> bool {
+        for (system, entities) in &other.per_system_entity_list {
+            let Some(self_entities) = self.per_system_entity_list.get(system) else {
+                if entities.is_empty() {
+                    continue;
+                } else {
+                    return false;
+                }
+            };
+            if !entities.is_subset(self_entities) {
+                return false;
+            }
+        }
+        true
     }
 
     /// Should be called on frame start.
@@ -225,7 +260,7 @@ impl DataBlueprintTree {
         // and/or when new entity paths are added, but such memoization would add complexity.
 
         fn project_tree(
-            tree: &mut DataBlueprintTree,
+            tree: &mut SpaceViewContents,
             parent_properties: &EntityProperties,
             group_handle: DataBlueprintGroupHandle,
         ) {
@@ -270,8 +305,6 @@ impl DataBlueprintTree {
         let mut new_leaf_groups = Vec::new();
 
         for path in paths {
-            self.entity_paths.insert(path.clone());
-
             // Is there already a group associated with this exact path? (maybe because a child was logged there earlier)
             // If so, we can simply move it under this existing group.
             let group_handle = if let Some(group_handle) = self.path_to_group.get(path) {
@@ -393,14 +426,19 @@ impl DataBlueprintTree {
     ///
     /// If the entity was not known by this data blueprint tree nothing happens.
     pub fn remove_entity(&mut self, path: &EntityPath) {
+        re_tracing::profile_function!();
+
         if let Some(group_handle) = self.path_to_group.get(path) {
             if let Some(group) = self.groups.get_mut(*group_handle) {
                 group.entities.remove(path);
                 self.remove_group_if_empty(*group_handle);
             }
         }
         self.path_to_group.remove(path);
-        self.entity_paths.remove(path);
+
+        for per_system_list in self.per_system_entity_list.values_mut() {
+            per_system_list.remove(path);
+        }
     }
 
     /// Removes a group and all its entities and subgroups from the blueprint tree
@@ -421,7 +459,9 @@ impl DataBlueprintTree {
 
         // Remove all child entities.
         for entity_path in &group.entities {
-            self.entity_paths.remove(entity_path);
+            for per_system_list in self.per_system_entity_list.values_mut() {
+                per_system_list.remove(entity_path);
+            }
         }
 
         // Remove from `path_to_group` map.

diff --git a/crates/re_space_view_bar_chart/src/view_part_system.rs b/crates/re_space_view_bar_chart/src/view_part_system.rs
@@ -6,8 +6,8 @@ use re_data_store::EntityPath;
 use re_log_types::{TimeInt, Timeline};
 use re_types::{ComponentName, Loggable as _};
 use re_viewer_context::{
-    ArchetypeDefinition, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem,
-    ViewQuery, ViewerContext,
+    ArchetypeDefinition, NamedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection,
+    ViewPartSystem, ViewQuery, ViewerContext,
 };
 
 /// A bar chart system, with everything needed to render it.
@@ -16,6 +16,12 @@ pub struct BarChartViewPartSystem {
     pub charts: BTreeMap<EntityPath, Tensor>,
 }
 
+impl NamedViewSystem for BarChartViewPartSystem {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "BarChartView".into()
+    }
+}
+
 impl ViewPartSystem for BarChartViewPartSystem {
     fn archetype(&self) -> ArchetypeDefinition {
         vec1::vec1![Tensor::name()]
@@ -51,7 +57,7 @@ impl ViewPartSystem for BarChartViewPartSystem {
 
         let store = &ctx.store_db.entity_db.data_store;
 
-        for (ent_path, props) in query.iter_entities() {
+        for (ent_path, props) in query.iter_entities_for_system(Self::name()) {
             if !props.visible {
                 continue;
             }

diff --git a/crates/re_space_view_spatial/src/contexts/annotation_context.rs b/crates/re_space_view_spatial/src/contexts/annotation_context.rs
@@ -1,9 +1,17 @@
 use re_types::{components::AnnotationContext, Loggable};
-use re_viewer_context::{AnnotationMap, ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{
+    AnnotationMap, ArchetypeDefinition, NamedViewSystem, ViewContextSystem, ViewSystemName,
+};
 
 #[derive(Default)]
 pub struct AnnotationSceneContext(pub AnnotationMap);
 
+impl NamedViewSystem for AnnotationSceneContext {
+    fn name() -> ViewSystemName {
+        "AnnotationSceneContext".into()
+    }
+}
+
 impl ViewContextSystem for AnnotationSceneContext {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![AnnotationContext::name()]]
@@ -15,11 +23,10 @@ impl ViewContextSystem for AnnotationSceneContext {
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
         re_tracing::profile_function!();
-        self.0.load(
-            ctx,
-            &query.latest_at_query(),
-            query.iter_entities().map(|(p, _)| p),
-        );
+        // We create a list of *all* entities here, do not only iterate over those with annotation context.
+        // TODO(andreas): But knowing ahead of time where we have annotation contexts could be used for optimization.
+        self.0
+            .load(ctx, &query.latest_at_query(), query.iter_entities());
     }
 
     fn as_any(&self) -> &dyn std::any::Any {

diff --git a/crates/re_space_view_spatial/src/contexts/depth_offsets.rs b/crates/re_space_view_spatial/src/contexts/depth_offsets.rs
@@ -4,7 +4,7 @@ use nohash_hasher::IntMap;
 
 use re_log_types::EntityPathHash;
 use re_types::{components::DrawOrder, Loggable as _};
-use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{ArchetypeDefinition, NamedViewSystem, ViewContextSystem};
 
 /// Context for creating a mapping from [`DrawOrder`] to [`re_renderer::DepthOffset`].
 #[derive(Default)]
@@ -19,6 +19,12 @@ pub struct EntityDepthOffsets {
     pub points: re_renderer::DepthOffset,
 }
 
+impl NamedViewSystem for EntityDepthOffsets {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "EntityDepthOffsets".into()
+    }
+}
+
 impl ViewContextSystem for EntityDepthOffsets {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![DrawOrder::name()]]
@@ -29,8 +35,6 @@ impl ViewContextSystem for EntityDepthOffsets {
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        re_tracing::profile_function!();
-
         #[derive(PartialEq, PartialOrd, Eq, Ord)]
         enum DrawOrderTarget {
             Entity(EntityPathHash),
@@ -44,7 +48,7 @@ impl ViewContextSystem for EntityDepthOffsets {
 
         // Use a BTreeSet for entity hashes to get a stable order.
         let mut entities_per_draw_order = BTreeMap::<DrawOrder, BTreeSet<DrawOrderTarget>>::new();
-        for (ent_path, _) in query.iter_entities() {
+        for (ent_path, _) in query.iter_entities_for_system(Self::name()) {
             if let Some(draw_order) = store.query_latest_component::<DrawOrder>(
                 ent_path,
                 &ctx.rec_cfg.time_ctrl.current_query(),

diff --git a/crates/re_space_view_spatial/src/contexts/mod.rs b/crates/re_space_view_spatial/src/contexts/mod.rs
@@ -15,7 +15,9 @@ pub use transform_context::{pinhole_camera_view_coordinates, TransformContext};
 // -----------------------------------------------------------------------------
 
 use re_renderer::DepthOffset;
-use re_viewer_context::{Annotations, SpaceViewClassRegistryError, ViewContextSystem};
+use re_viewer_context::{
+    Annotations, NamedViewSystem, SpaceViewClassRegistryError, ViewContextSystem,
+};
 
 /// Context objects for a single entity in a spatial scene.
 pub struct SpatialSceneEntityContext<'a> {
@@ -32,6 +34,12 @@ pub struct PrimitiveCounter {
     pub num_primitives: AtomicUsize,
 }
 
+impl NamedViewSystem for PrimitiveCounter {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "PrimitiveCounter".into()
+    }
+}
+
 impl ViewContextSystem for PrimitiveCounter {
     fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
         Vec::new()

diff --git a/crates/re_space_view_spatial/src/contexts/non_interactive_entities.rs b/crates/re_space_view_spatial/src/contexts/non_interactive_entities.rs
@@ -1,13 +1,19 @@
 use nohash_hasher::IntSet;
 use re_log_types::EntityPathHash;
-use re_viewer_context::ViewContextSystem;
+use re_viewer_context::{NamedViewSystem, ViewContextSystem};
 
 /// List of all non-interactive entities for lookup during picking evaluation.
 ///
 /// TODO(wumpf/jleibs): This is a temporary solution until the picking code can query propagated blueprint properties directly.
 #[derive(Default)]
 pub struct NonInteractiveEntities(pub IntSet<EntityPathHash>);
 
+impl NamedViewSystem for NonInteractiveEntities {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "NonInteractiveEntities".into()
+    }
+}
+
 impl ViewContextSystem for NonInteractiveEntities {
     fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
         Vec::new()
@@ -18,8 +24,6 @@ impl ViewContextSystem for NonInteractiveEntities {
         _ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        re_tracing::profile_function!();
-
         self.0 = query
             .entity_props_map
             .iter()

diff --git a/crates/re_space_view_spatial/src/contexts/shared_render_builders.rs b/crates/re_space_view_spatial/src/contexts/shared_render_builders.rs
@@ -1,6 +1,6 @@
 use parking_lot::{MappedMutexGuard, Mutex, MutexGuard};
 use re_renderer::{LineStripSeriesBuilder, PointCloudBuilder, RenderContext};
-use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{ArchetypeDefinition, NamedViewSystem, ViewContextSystem};
 
 use crate::parts::{
     SIZE_BOOST_IN_POINTS_FOR_LINE_OUTLINES, SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES,
@@ -14,6 +14,12 @@ pub struct SharedRenderBuilders {
     pub points: Mutex<Option<PointCloudBuilder>>,
 }
 
+impl NamedViewSystem for SharedRenderBuilders {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "SharedRenderBuilders".into()
+    }
+}
+
 impl SharedRenderBuilders {
     pub fn lines(&self) -> MappedMutexGuard<'_, LineStripSeriesBuilder> {
         MutexGuard::map(self.lines.lock(), |l| l.as_mut().unwrap())