Parallelized transform propagation (bevyengine#4775)

james7132 · ItsDoot · commit 9e4ca8641835 · 2023-01-31T20:00:09.000-06:00
# Objective Fixes bevyengine#4697. Hierarchical propagation of properties, currently only Transform -> GlobalTransform, can be a very expensive operation. Transform propagation is a strict dependency for anything positioned in world-space. In large worlds, this can take quite a bit of time, so limiting it to a single thread can result in poor CPU utilization as it bottlenecks the rest of the frame's systems. ## Solution - Move transforms without a parent or a child (free-floating (Global)Transform) entities into a separate parallel system. - Chunk the hierarchy based on the root entities and process it in parallel with `Query::par_for_each_mut`. - Utilize the hierarchy's specific properties introduced in bevyengine#4717 to allow for safe use of `Query::get_unchecked` on multiple threads. Assuming each child is unique in the hierarchy, it is impossible to have an aliased `&mut GlobalTransform` so long as we verify that the parent for a child is the same one propagated from. --- ## Changelog Removed: `transform_propagate_system` is no longer `pub`.
diff --git a/crates/bevy_transform/Cargo.toml b/crates/bevy_transform/Cargo.toml
@@ -17,5 +17,8 @@ bevy_math = { path = "../bevy_math", version = "0.9.0" }
 bevy_reflect = { path = "../bevy_reflect", version = "0.9.0", features = ["bevy"] }
 serde = { version = "1", features = ["derive"], optional = true }
 
+[dev_dependencies]
+bevy_tasks = { path = "../bevy_tasks", version = "0.9.0-dev" }
+
 [features]
 serialize = ["dep:serde", "bevy_math/serialize"]
diff --git a/crates/bevy_transform/src/components/global_transform.rs b/crates/bevy_transform/src/components/global_transform.rs
@@ -19,8 +19,8 @@ use bevy_reflect::{std_traits::ReflectDefault, FromReflect, Reflect};
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system).
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labeled
+/// [`TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the [`Transform`] of an entity in this stage or after, you will notice a 1 frame lag
diff --git a/crates/bevy_transform/src/components/transform.rs b/crates/bevy_transform/src/components/transform.rs
@@ -20,8 +20,8 @@ use std::ops::Mul;
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system).
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labeled
+/// [`TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the [`Transform`] of an entity in this stage or after, you will notice a 1 frame lag
diff --git a/crates/bevy_transform/src/lib.rs b/crates/bevy_transform/src/lib.rs
@@ -1,10 +1,10 @@
 #![warn(missing_docs)]
+#![warn(clippy::undocumented_unsafe_blocks)]
 #![doc = include_str!("../README.md")]
 
 /// The basic components of the transform crate
 pub mod components;
 mod systems;
-pub use crate::systems::transform_propagate_system;
 
 #[doc(hidden)]
 pub mod prelude {
@@ -32,8 +32,8 @@ use prelude::{GlobalTransform, Transform};
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`].
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labeled
+/// [`TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the [`Transform`] of an entity in this stage or after, you will notice a 1 frame lag
@@ -91,11 +91,19 @@ impl Plugin for TransformPlugin {
             // add transform systems to startup so the first update is "correct"
             .add_startup_system_to_stage(
                 StartupStage::PostStartup,
-                systems::transform_propagate_system.label(TransformSystem::TransformPropagate),
+                systems::sync_simple_transforms.label(TransformSystem::TransformPropagate),
+            )
+            .add_startup_system_to_stage(
+                StartupStage::PostStartup,
+                systems::propagate_transforms.label(TransformSystem::TransformPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                systems::sync_simple_transforms.label(TransformSystem::TransformPropagate),
             )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
-                systems::transform_propagate_system.label(TransformSystem::TransformPropagate),
+                systems::propagate_transforms.label(TransformSystem::TransformPropagate),
             );
     }
 }
diff --git a/crates/bevy_transform/src/systems.rs b/crates/bevy_transform/src/systems.rs
@@ -2,95 +2,140 @@ use crate::components::{GlobalTransform, Transform};
 use bevy_ecs::prelude::{Changed, Entity, Query, With, Without};
 use bevy_hierarchy::{Children, Parent};
 
+/// Update [`GlobalTransform`] component of entities that aren't in the hierarchy
+pub fn sync_simple_transforms(
+    mut query: Query<
+        (&Transform, &mut GlobalTransform),
+        (Changed<Transform>, Without<Parent>, Without<Children>),
+    >,
+) {
+    query.par_for_each_mut(1024, |(transform, mut global_transform)| {
+        *global_transform = GlobalTransform::from(*transform);
+    });
+}
+
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
 /// [`Transform`] component.
-pub fn transform_propagate_system(
+pub fn propagate_transforms(
     mut root_query: Query<
         (
-            Option<(&Children, Changed<Children>)>,
+            Entity,
+            &Children,
             &Transform,
             Changed<Transform>,
+            Changed<Children>,
             &mut GlobalTransform,
-            Entity,
         ),
         Without<Parent>,
     >,
-    mut transform_query: Query<(
-        &Transform,
-        Changed<Transform>,
-        &mut GlobalTransform,
-        &Parent,
-    )>,
+    transform_query: Query<(&Transform, Changed<Transform>, &mut GlobalTransform), With<Parent>>,
+    parent_query: Query<&Parent>,
     children_query: Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
 ) {
-    for (children, transform, transform_changed, mut global_transform, entity) in
-        root_query.iter_mut()
-    {
-        let mut changed = transform_changed;
-        if transform_changed {
-            *global_transform = GlobalTransform::from(*transform);
-        }
+    root_query.par_for_each_mut(
+        // The differing depths and sizes of hierarchy trees causes the work for each root to be
+        // different. A batch size of 1 ensures that each tree gets it's own task and multiple
+        // large trees are not clumped together.
+        1,
+        |(entity, children, transform, mut changed, children_changed, mut global_transform)| {
+            if changed {
+                *global_transform = GlobalTransform::from(*transform);
+            }
 
-        if let Some((children, changed_children)) = children {
             // If our `Children` has changed, we need to recalculate everything below us
-            changed |= changed_children;
-            for child in children {
-                let _ = propagate_recursive(
+            changed |= children_changed;
+
+            for child in children.iter() {
+                propagate_recursive(
                     &global_transform,
-                    &mut transform_query,
+                    &transform_query,
+                    &parent_query,
                     &children_query,
-                    *child,
                     entity,
+                    *child,
                     changed,
                 );
             }
-        }
-    }
+        },
+    );
 }
 
 fn propagate_recursive(
     parent: &GlobalTransform,
-    transform_query: &mut Query<(
-        &Transform,
-        Changed<Transform>,
-        &mut GlobalTransform,
-        &Parent,
-    )>,
+    unsafe_transform_query: &Query<
+        (&Transform, Changed<Transform>, &mut GlobalTransform),
+        With<Parent>,
+    >,
+    parent_query: &Query<&Parent>,
     children_query: &Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
-    entity: Entity,
     expected_parent: Entity,
+    entity: Entity,
     mut changed: bool,
     // We use a result here to use the `?` operator. Ideally we'd use a try block instead
-) -> Result<(), ()> {
+) {
+    let Ok(actual_parent) = parent_query.get(entity) else {
+        panic!("Propagated child for {:?} has no Parent component!", entity);
+    };
+    assert_eq!(
+        actual_parent.get(), expected_parent,
+        "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
+    );
+
     let global_matrix = {
-        let (transform, transform_changed, mut global_transform, child_parent) =
-            transform_query.get_mut(entity).map_err(drop)?;
-        // Note that for parallelising, this check cannot occur here, since there is an `&mut GlobalTransform` (in global_transform)
-        assert_eq!(
-            child_parent.get(), expected_parent,
-            "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
-        );
+        let Ok((transform, transform_changed, mut global_transform)) =
+            // SAFETY: This call cannot create aliased mutable references.
+            //   - The top level iteration parallelizes on the roots of the hierarchy.
+            //   - The above assertion ensures that each child has one and only one unique parent throughout the entire
+            //     hierarchy.
+            //
+            // For example, consider the following malformed hierarchy:
+            //
+            //     A
+            //   /   \
+            //  B     C
+            //   \   /
+            //     D
+            //
+            // D has two parents, B and C. If the propagation passes through C, but the Parent component on D points to B,
+            // the above check will panic as the origin parent does match the recorded parent.
+            //
+            // Also consider the following case, where A and B are roots:
+            //
+            //  A       B
+            //   \     /
+            //    C   D
+            //     \ /
+            //      E
+            //
+            // Even if these A and B start two separate tasks running in parallel, one of them will panic before attempting
+            // to mutably access E.
+            (unsafe { unsafe_transform_query.get_unchecked(entity) }) else {
+                return;
+            };
+
         changed |= transform_changed;
         if changed {
             *global_transform = parent.mul_transform(*transform);
         }
         *global_transform
     };
 
-    let (children, changed_children) = children_query.get(entity).map_err(drop)?;
+    let Ok((children, changed_children)) = children_query.get(entity) else {
+        return
+    };
     // If our `Children` has changed, we need to recalculate everything below us
     changed |= changed_children;
     for child in children {
-        let _ = propagate_recursive(
+        propagate_recursive(
             &global_matrix,
-            transform_query,
+            unsafe_transform_query,
+            parent_query,
             children_query,
-            *child,
             entity,
+            *child,
             changed,
         );
     }
-    Ok(())
 }
 
 #[cfg(test)]
@@ -99,9 +144,10 @@ mod test {
     use bevy_ecs::prelude::*;
     use bevy_ecs::system::CommandQueue;
     use bevy_math::vec3;
+    use bevy_tasks::{ComputeTaskPool, TaskPool};
 
     use crate::components::{GlobalTransform, Transform};
-    use crate::systems::transform_propagate_system;
+    use crate::systems::*;
     use crate::TransformBundle;
     use bevy_hierarchy::{BuildChildren, BuildWorldChildren, Children, Parent};
 
@@ -110,10 +156,12 @@ mod test {
 
     #[test]
     fn did_propagate() {
+        ComputeTaskPool::init(TaskPool::default);
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -152,8 +200,10 @@ mod test {
     #[test]
     fn did_propagate_command_buffer() {
         let mut world = World::default();
+
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -192,10 +242,12 @@ mod test {
 
     #[test]
     fn correct_children() {
+        ComputeTaskPool::init(TaskPool::default);
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -272,8 +324,10 @@ mod test {
     #[test]
     fn correct_transforms_when_no_children() {
         let mut app = App::new();
+        ComputeTaskPool::init(TaskPool::default);
 
-        app.add_system(transform_propagate_system);
+        app.add_system(sync_simple_transforms);
+        app.add_system(propagate_transforms);
 
         let translation = vec3(1.0, 0.0, 0.0);
 
@@ -313,12 +367,14 @@ mod test {
     #[test]
     #[should_panic]
     fn panic_when_hierarchy_cycle() {
+        ComputeTaskPool::init(TaskPool::default);
         // We cannot directly edit Parent and Children, so we use a temp world to break
         // the hierarchy's invariants.
         let mut temp = World::new();
         let mut app = App::new();
 
-        app.add_system(transform_propagate_system);
+        app.add_system(propagate_transforms)
+            .add_system(sync_simple_transforms);
 
         fn setup_world(world: &mut World) -> (Entity, Entity) {
             let mut grandchild = Entity::from_raw(0);
