Parallize Space View system evaluation (#4460)

### What * Fixes #1325 Two levels of parallism actually: * All context & part systems for a single space view each run in parallel * All space views now in parallel first setup their query and then kick of their (parallized!) system execution This gives us major speedups whenever: * There is several heavy systems in a scene (e.g. a lot of points and a lot of arrows) * There is several space views As such this is a very fundamental & widely applicable parallelization! ![image](https://github.com/rerun-io/rerun/assets/1220815/14bc6ca8-abb0-480b-91dc-2da51253dd92) (a rather silly "points with arrows" scene duplicated a bunch of times. Note that we still don't have any caching, so all these space views might as well be animated individually!) <img width="1421" alt="image" src="https://github.com/rerun-io/rerun/assets/1220815/57e13e15-9d8f-4a30-bf74-d8174676986d"> (busy rayon worker threads!) Related things that are not parallel yet: * ui methods per space view (run one by one) * draw call recording (done as part of the ui method, but we could likely do this in parallel, only joining on these tasks when egui wants to queue up the command buffers) * processing _within_ a system depends on the system at hand. Today, only point cloud has some parallization ### 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): * Full build: [app.rerun.io](https://app.rerun.io/pr/4460/index.html) * Partial build: [app.rerun.io](https://app.rerun.io/pr/4460/index.html?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json) - Useful for quick testing when changes do not affect examples in any way * [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/4460) - [Docs preview](https://rerun.io/preview/cb7e9a9902bb9923ade41271ccc43ed01a2a4d33/docs)  - [Examples preview](https://rerun.io/preview/cb7e9a9902bb9923ade41271ccc43ed01a2a4d33/examples)  - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) --- End of a series on refactors leading to this parallization: * #4387 * #4404 * #4389 * #4421 * #4422 * #4430 * #4438 * #4460 --------- Co-authored-by: Emil Ernerfeldt <[email protected]>
rerun-io · Dec 11, 2023 · ea4dff1 · ea4dff1
1 parent a1c11f3
commit ea4dff1
Show file tree

Hide file tree

Showing 36 changed files with 352 additions and 215 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/re_renderer/src/allocator/cpu_write_gpu_read_belt.rs b/crates/re_renderer/src/allocator/cpu_write_gpu_read_belt.rs
@@ -38,7 +38,7 @@ pub enum CpuWriteGpuReadError {
 /// Note that the "vec like behavior" further encourages
 /// * not leaving holes
 /// * keeping writes sequential
-pub struct CpuWriteGpuReadBuffer<T: bytemuck::Pod + 'static> {
+pub struct CpuWriteGpuReadBuffer<T: bytemuck::Pod + Send + Sync> {
     /// Write view into the relevant buffer portion.
     ///
     /// UNSAFE: The lifetime is transmuted to be `'static`.
@@ -55,9 +55,17 @@ pub struct CpuWriteGpuReadBuffer<T: bytemuck::Pod + 'static> {
     _type: std::marker::PhantomData<T>,
 }
 
+#[allow(unsafe_code)]
+// SAFETY: TODO(gfx-rs/wgpu#4818): Upstream wgpu allows `wgpu::BufferViewMut` to be Send.
+unsafe impl<T> Send for CpuWriteGpuReadBuffer<T> where T: bytemuck::Pod + Send + Sync {}
+
+#[allow(unsafe_code)]
+// SAFETY: TODO(gfx-rs/wgpu#4818): Upstream wgpu allows `wgpu::BufferViewMut` to be Sync.
+unsafe impl<T> Sync for CpuWriteGpuReadBuffer<T> where T: bytemuck::Pod + Send + Sync {}
+
 impl<T> CpuWriteGpuReadBuffer<T>
 where
-    T: bytemuck::Pod + 'static,
+    T: bytemuck::Pod + Send + Sync,
 {
     /// Memory as slice.
     ///
@@ -283,7 +291,7 @@ impl Chunk {
     }
 
     /// Caller needs to make sure that there is enough space.
-    fn allocate<T: bytemuck::Pod>(
+    fn allocate<T: bytemuck::Pod + Send + Sync>(
         &mut self,
         num_elements: usize,
         size_in_bytes: u64,
@@ -419,7 +427,7 @@ impl CpuWriteGpuReadBelt {
     }
 
     /// Allocates a cpu writable buffer for `num_elements` instances of type `T`.
-    pub fn allocate<T: bytemuck::Pod>(
+    pub fn allocate<T: bytemuck::Pod + Send + Sync>(
         &mut self,
         device: &wgpu::Device,
         buffer_pool: &GpuBufferPool,

diff --git a/crates/re_renderer/src/allocator/uniform_buffer_fill.rs b/crates/re_renderer/src/allocator/uniform_buffer_fill.rs
@@ -42,7 +42,7 @@ impl<T> UniformBufferAlignmentCheck<T> {
 /// For subsequent frames, this will automatically not allocate any resources (thanks to our buffer pooling mechanism).
 ///
 /// TODO(#1383): We could do this on a more complex stack allocator.
-pub fn create_and_fill_uniform_buffer_batch<T: bytemuck::Pod>(
+pub fn create_and_fill_uniform_buffer_batch<T: bytemuck::Pod + Send + Sync>(
     ctx: &RenderContext,
     label: DebugLabel,
     content: impl ExactSizeIterator<Item = T>,
@@ -97,7 +97,7 @@ pub fn create_and_fill_uniform_buffer_batch<T: bytemuck::Pod>(
 }
 
 /// See [`create_and_fill_uniform_buffer`].
-pub fn create_and_fill_uniform_buffer<T: bytemuck::Pod>(
+pub fn create_and_fill_uniform_buffer<T: bytemuck::Pod + Send + Sync>(
     ctx: &RenderContext,
     label: DebugLabel,
     content: T,

diff --git a/crates/re_space_view_bar_chart/src/space_view_class.rs b/crates/re_space_view_bar_chart/src/space_view_class.rs
@@ -5,8 +5,7 @@ use re_space_view::controls;
 use re_types::datatypes::TensorBuffer;
 use re_viewer_context::{
     auto_color, SpaceViewClass, SpaceViewClassRegistryError, SpaceViewId,
-    SpaceViewSystemExecutionError, ViewContextCollection, ViewPartCollection, ViewQuery,
-    ViewerContext,
+    SpaceViewSystemExecutionError, ViewQuery, ViewerContext,
 };
 
 use super::view_part_system::BarChartViewPartSystem;
@@ -129,14 +128,15 @@ impl SpaceViewClass for BarChartSpaceView {
         ui: &mut egui::Ui,
         _state: &mut Self::State,
         root_entity_properties: &EntityProperties,
-        _view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         _query: &ViewQuery<'_>,
-        _draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
         use egui_plot::{Bar, BarChart, Legend, Plot};
 
-        let charts = &parts.get::<BarChartViewPartSystem>()?.charts;
+        let charts = &system_output
+            .view_systems
+            .get::<BarChartViewPartSystem>()?
+            .charts;
 
         let zoom_both_axis = !ui.input(|i| i.modifiers.contains(controls::ASPECT_SCROLL_MODIFIER));
 

diff --git a/crates/re_space_view_spatial/src/space_view_2d.rs b/crates/re_space_view_spatial/src/space_view_2d.rs
@@ -2,8 +2,7 @@ use re_data_store::EntityProperties;
 use re_log_types::EntityPath;
 use re_viewer_context::{
     AutoSpawnHeuristic, PerSystemEntities, SpaceViewClass, SpaceViewClassRegistryError,
-    SpaceViewId, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartCollection,
-    ViewQuery, ViewerContext,
+    SpaceViewId, SpaceViewSystemExecutionError, ViewQuery, ViewerContext,
 };
 
 use crate::{
@@ -141,17 +140,17 @@ impl SpaceViewClass for SpatialSpaceView2D {
         ui: &mut egui::Ui,
         state: &mut Self::State,
         _root_entity_properties: &EntityProperties,
-        view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         query: &ViewQuery<'_>,
-        draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
-        state.scene_bbox = calculate_bounding_box(parts, &mut state.scene_bbox_accum);
-        state.scene_num_primitives = view_ctx
+        state.scene_bbox =
+            calculate_bounding_box(&system_output.view_systems, &mut state.scene_bbox_accum);
+        state.scene_num_primitives = system_output
+            .context_systems
             .get::<PrimitiveCounter>()?
             .num_primitives
             .load(std::sync::atomic::Ordering::Relaxed);
 
-        crate::ui_2d::view_2d(ctx, ui, state, view_ctx, parts, query, draw_data)
+        crate::ui_2d::view_2d(ctx, ui, state, query, system_output)
     }
 }
diff --git a/crates/re_space_view_spatial/src/space_view_3d.rs b/crates/re_space_view_spatial/src/space_view_3d.rs
@@ -2,8 +2,8 @@ use re_data_store::EntityProperties;
 use re_log_types::EntityPath;
 use re_viewer_context::{
     AutoSpawnHeuristic, IdentifiedViewSystem as _, PerSystemEntities, SpaceViewClass,
-    SpaceViewClassRegistryError, SpaceViewId, SpaceViewSystemExecutionError, ViewContextCollection,
-    ViewPartCollection, ViewQuery, ViewerContext,
+    SpaceViewClassRegistryError, SpaceViewId, SpaceViewSystemExecutionError, ViewQuery,
+    ViewerContext,
 };
 
 use crate::{
@@ -116,19 +116,19 @@ impl SpaceViewClass for SpatialSpaceView3D {
         ui: &mut egui::Ui,
         state: &mut Self::State,
         _root_entity_properties: &EntityProperties,
-        view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         query: &ViewQuery<'_>,
-        draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
         re_tracing::profile_function!();
 
-        state.scene_bbox = calculate_bounding_box(parts, &mut state.scene_bbox_accum);
-        state.scene_num_primitives = view_ctx
+        state.scene_bbox =
+            calculate_bounding_box(&system_output.view_systems, &mut state.scene_bbox_accum);
+        state.scene_num_primitives = system_output
+            .context_systems
             .get::<PrimitiveCounter>()?
             .num_primitives
             .load(std::sync::atomic::Ordering::Relaxed);
 
-        crate::ui_3d::view_3d(ctx, ui, state, view_ctx, parts, query, draw_data)
+        crate::ui_3d::view_3d(ctx, ui, state, query, system_output)
     }
 }
diff --git a/crates/re_space_view_spatial/src/ui_2d.rs b/crates/re_space_view_spatial/src/ui_2d.rs
@@ -9,8 +9,8 @@ use re_renderer::view_builder::{TargetConfiguration, ViewBuilder};
 use re_space_view::controls::{DRAG_PAN2D_BUTTON, RESET_VIEW_BUTTON_TEXT, ZOOM_SCROLL_MODIFIER};
 use re_types::{archetypes::Pinhole, components::ViewCoordinates};
 use re_viewer_context::{
-    gpu_bridge, HoveredSpace, SpaceViewSystemExecutionError, ViewContextCollection,
-    ViewPartCollection, ViewQuery, ViewerContext,
+    gpu_bridge, HoveredSpace, SpaceViewSystemExecutionError, SystemExecutionOutput, ViewQuery,
+    ViewerContext,
 };
 
 use super::{
@@ -228,13 +228,17 @@ pub fn view_2d(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
     state: &mut SpatialSpaceViewState,
-    view_ctx: &ViewContextCollection,
-    parts: &ViewPartCollection,
     query: &ViewQuery<'_>,
-    mut draw_data: Vec<re_renderer::QueueableDrawData>,
+    system_output: re_viewer_context::SystemExecutionOutput,
 ) -> Result<(), SpaceViewSystemExecutionError> {
     re_tracing::profile_function!();
 
+    let SystemExecutionOutput {
+        view_systems: parts,
+        context_systems: view_ctx,
+        draw_data,
+    } = system_output;
+
     // Save off the available_size since this is used for some of the layout updates later
     let available_size = ui.available_size();
     let store = ctx.store_db.store();
@@ -312,11 +316,11 @@ pub fn view_2d(
 
         // Create labels now since their shapes participate are added to scene.ui for picking.
         let (label_shapes, ui_rects) = create_labels(
-            &collect_ui_labels(parts),
+            &collect_ui_labels(&parts),
             ui_from_canvas,
             &eye,
             ui,
-            query.highlights,
+            &query.highlights,
             SpatialSpaceViewKind::TwoD,
         );
 
@@ -330,15 +334,15 @@ pub fn view_2d(
                 eye,
                 &mut view_builder,
                 state,
-                view_ctx,
-                parts,
+                &view_ctx,
+                &parts,
                 &ui_rects,
                 query,
                 SpatialSpaceViewKind::TwoD,
             )?;
         }
 
-        for draw_data in draw_data.drain(..) {
+        for draw_data in draw_data {
             view_builder.queue_draw(draw_data);
         }
         if let Ok(shared_render_builders) = view_ctx.get::<SharedRenderBuilders>() {

diff --git a/crates/re_space_view_spatial/src/ui_3d.rs b/crates/re_space_view_spatial/src/ui_3d.rs
@@ -13,8 +13,8 @@ use re_space_view::controls::{
 };
 use re_types::components::ViewCoordinates;
 use re_viewer_context::{
-    gpu_bridge, HoveredSpace, Item, SpaceViewSystemExecutionError, ViewContextCollection,
-    ViewPartCollection, ViewQuery, ViewerContext,
+    gpu_bridge, HoveredSpace, Item, SpaceViewSystemExecutionError, SystemExecutionOutput,
+    ViewQuery, ViewerContext,
 };
 
 use crate::{
@@ -307,19 +307,22 @@ pub fn help_text(re_ui: &re_ui::ReUi) -> egui::WidgetText {
     layout.layout_job.into()
 }
 
-/// TODO(andreas): Split into smaller parts, more re-use with `ui_2d`
 pub fn view_3d(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
     state: &mut SpatialSpaceViewState,
-    view_ctx: &ViewContextCollection,
-    parts: &ViewPartCollection,
     query: &ViewQuery<'_>,
-    mut draw_data: Vec<re_renderer::QueueableDrawData>,
+    system_output: re_viewer_context::SystemExecutionOutput,
 ) -> Result<(), SpaceViewSystemExecutionError> {
     re_tracing::profile_function!();
 
-    let highlights = query.highlights;
+    let SystemExecutionOutput {
+        view_systems: parts,
+        context_systems: view_ctx,
+        draw_data,
+    } = system_output;
+
+    let highlights = &query.highlights;
     let space_cameras = &parts.get::<CamerasPart>()?.space_cameras;
     let view_coordinates = ctx
         .store_db
@@ -412,7 +415,7 @@ pub fn view_3d(
 
     // Create labels now since their shapes participate are added to scene.ui for picking.
     let (label_shapes, ui_rects) = create_labels(
-        &collect_ui_labels(parts),
+        &collect_ui_labels(&parts),
         RectTransform::from_to(rect, rect),
         &eye,
         ui,
@@ -430,8 +433,8 @@ pub fn view_3d(
             eye,
             &mut view_builder,
             state,
-            view_ctx,
-            parts,
+            &view_ctx,
+            &parts,
             &ui_rects,
             query,
             SpatialSpaceViewKind::ThreeD,
@@ -586,7 +589,7 @@ pub fn view_3d(
         }
     }
 
-    for draw_data in draw_data.drain(..) {
+    for draw_data in draw_data {
         view_builder.queue_draw(draw_data);
     }
     if let Ok(shared_render_builders) = view_ctx.get::<SharedRenderBuilders>() {

diff --git a/crates/re_space_view_tensor/src/space_view_class.rs b/crates/re_space_view_tensor/src/space_view_class.rs
@@ -15,7 +15,7 @@ use re_types::{
 use re_viewer_context::{
     gpu_bridge, gpu_bridge::colormap_dropdown_button_ui, SpaceViewClass,
     SpaceViewClassRegistryError, SpaceViewId, SpaceViewState, SpaceViewSystemExecutionError,
-    TensorStatsCache, ViewContextCollection, ViewPartCollection, ViewQuery, ViewerContext,
+    TensorStatsCache, ViewQuery, ViewerContext,
 };
 
 use crate::{tensor_dimension_mapper::dimension_mapping_ui, view_part_system::TensorSystem};
@@ -184,14 +184,12 @@ impl SpaceViewClass for TensorSpaceView {
         ui: &mut egui::Ui,
         state: &mut Self::State,
         _root_entity_properties: &EntityProperties,
-        _view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         _query: &ViewQuery<'_>,
-        _draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
         re_tracing::profile_function!();
 
-        let tensors = &parts.get::<TensorSystem>()?.tensors;
+        let tensors = &system_output.view_systems.get::<TensorSystem>()?.tensors;
 
         if tensors.is_empty() {
             ui.centered_and_justified(|ui| ui.label("(empty)"));

diff --git a/crates/re_space_view_text_document/src/space_view_class.rs b/crates/re_space_view_text_document/src/space_view_class.rs
@@ -3,8 +3,7 @@ use egui::Label;
 use re_viewer_context::external::re_data_store::EntityProperties;
 use re_viewer_context::{
     external::re_log_types::EntityPath, SpaceViewClass, SpaceViewClassRegistryError, SpaceViewId,
-    SpaceViewState, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartCollection,
-    ViewQuery, ViewerContext,
+    SpaceViewState, SpaceViewSystemExecutionError, ViewQuery, ViewerContext,
 };
 
 use crate::view_part_system::TextDocumentEntry;
@@ -99,12 +98,10 @@ impl SpaceViewClass for TextDocumentSpaceView {
         ui: &mut egui::Ui,
         state: &mut Self::State,
         _root_entity_properties: &EntityProperties,
-        _view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         _query: &ViewQuery<'_>,
-        _draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
-        let text_document = parts.get::<TextDocumentSystem>()?;
+        let text_document = system_output.view_systems.get::<TextDocumentSystem>()?;
 
         egui::Frame {
             inner_margin: re_ui::ReUi::view_padding().into(),

diff --git a/crates/re_space_view_text_log/src/space_view_class.rs b/crates/re_space_view_text_log/src/space_view_class.rs
@@ -7,7 +7,7 @@ use re_types::components::TextLogLevel;
 use re_viewer_context::{
     level_to_rich_text, AutoSpawnHeuristic, PerSystemEntities, SpaceViewClass,
     SpaceViewClassRegistryError, SpaceViewId, SpaceViewState, SpaceViewSystemExecutionError,
-    ViewContextCollection, ViewPartCollection, ViewQuery, ViewerContext,
+    ViewQuery, ViewerContext,
 };
 
 use super::view_part_system::{Entry, TextLogSystem};
@@ -136,13 +136,11 @@ impl SpaceViewClass for TextSpaceView {
         ui: &mut egui::Ui,
         state: &mut Self::State,
         _root_entity_properties: &EntityProperties,
-        _view_ctx: &ViewContextCollection,
-        parts: &ViewPartCollection,
         _query: &ViewQuery<'_>,
-        _draw_data: Vec<re_renderer::QueueableDrawData>,
+        system_output: re_viewer_context::SystemExecutionOutput,
     ) -> Result<(), SpaceViewSystemExecutionError> {
         re_tracing::profile_function!();
-        let text = parts.get::<TextLogSystem>()?;
+        let text = system_output.view_systems.get::<TextLogSystem>()?;
 
         // TODO(andreas): Should filter text entries in the part-system instead.
         // this likely requires a way to pass state into a context.