Run individual SystemSets in Schedules

benches/benches/bevy_ecs/scheduling/mod.rs

@@ -19,4 +19,5 @@ criterion_group!(
     schedule,
     build_schedule,
     empty_schedule_run,
+    run_schedule,
 );

benches/benches/bevy_ecs/scheduling/schedule.rs

@@ -138,3 +138,123 @@ pub fn empty_schedule_run(criterion: &mut Criterion) {
 
     group.finish();
 }
+
+pub fn run_schedule(criterion: &mut Criterion) {
+    #[derive(Component)]
+    struct A(f32);
+    #[derive(Component)]
+    struct B(f32);
+    #[derive(Component)]
+    struct C(f32);
+
+    #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    struct SetA;
+    #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    struct SetB;
+    #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    struct SetC;
+
+    fn system_a(mut query: Query<&mut A>) {
+        query.iter_mut().for_each(|mut a| {
+            a.0 += 1.0;
+        });
+    }
+
+    fn system_b(mut query: Query<&mut B>) {
+        query.iter_mut().for_each(|mut b| {
+            b.0 += 1.0;
+        });
+    }
+
+    fn system_c(mut query: Query<&mut C>) {
+        query.iter_mut().for_each(|mut c| {
+            c.0 += 1.0;
+        });
+    }
+
+    let mut group = criterion.benchmark_group("run_schedule");
+    group.warm_up_time(core::time::Duration::from_millis(500));
+    group.measurement_time(core::time::Duration::from_secs(4));
+
+    group.bench_function("full/single_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::SingleThreaded);
+        schedule.add_systems(system_a.in_set(SetA));
+        schedule.add_systems(system_b.in_set(SetB));
+        schedule.add_systems(system_c.in_set(SetC));
+        schedule.run(&mut world);
+
+        bencher.iter(|| schedule.run(&mut world));
+    });
+
+    group.bench_function("full/multi_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::MultiThreaded);
+        schedule.add_systems(system_a.in_set(SetA));
+        schedule.add_systems(system_b.in_set(SetB));
+        schedule.add_systems(system_c.in_set(SetC));
+        schedule.run(&mut world);
+
+        bencher.iter(|| schedule.run(&mut world));
+    });
+
+    group.bench_function("single_system_in_set/single_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::SingleThreaded);
+        schedule.add_systems(system_a.in_set(SetA));
+        schedule.add_systems(system_b.in_set(SetB));
+        schedule.add_systems(system_c.in_set(SetC));
+        schedule.run_system_set(&mut world, SetB);
+
+        bencher.iter(|| schedule.run_system_set(&mut world, SetB));
+    });
+
+    group.bench_function("single_system_in_set/multi_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::MultiThreaded);
+        schedule.add_systems(system_a.in_set(SetA));
+        schedule.add_systems(system_b.in_set(SetB));
+        schedule.add_systems(system_c.in_set(SetC));
+        schedule.run_system_set(&mut world, SetB);
+
+        bencher.iter(|| schedule.run_system_set(&mut world, SetB));
+    });
+
+    group.bench_function("multiple_systems_in_set/single_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::SingleThreaded);
+        schedule.add_systems((system_a, system_b, system_c).in_set(SetA));
+        schedule.run_system_set(&mut world, SetA);
+
+        bencher.iter(|| schedule.run_system_set(&mut world, SetA));
+    });
+
+    group.bench_function("multiple_systems_in_set/multi_threaded", |bencher| {
+        let mut world = World::default();
+        world.spawn_batch((0..10000).map(|_| (A(0.0), B(0.0), C(0.0))));
+
+        let mut schedule = Schedule::default();
+        schedule.set_executor_kind(bevy_ecs::schedule::ExecutorKind::MultiThreaded);
+        schedule.add_systems((system_a, system_b, system_c).in_set(SetA));
+        schedule.run_system_set(&mut world, SetA);
+
+        bencher.iter(|| schedule.run_system_set(&mut world, SetA));
+    });
+
+    group.finish();
+}

crates/bevy_ecs/src/schedule/executor/mod.rs

@@ -3,6 +3,7 @@ mod multi_threaded;
 mod single_threaded;
 
 use alloc::{vec, vec::Vec};
+use bevy_platform::collections::HashMap;
 use bevy_utils::prelude::DebugName;
 use core::any::TypeId;
 
@@ -34,6 +35,7 @@ pub(super) trait SystemExecutor: Send + Sync {
         &mut self,
         schedule: &mut SystemSchedule,
         world: &mut World,
+        subgraph: Option<SystemSetKey>,
         skip_systems: Option<&FixedBitSet>,
         error_handler: fn(BevyError, ErrorContext),
     );
@@ -105,11 +107,14 @@ pub struct SystemSchedule {
     ///
     /// If a set doesn't run because of its conditions, this is used to skip all systems in it.
     pub(super) systems_in_sets_with_conditions: Vec<FixedBitSet>,
+    /// Sparse mapping of system set node id to systems in the set, used for running
+    /// subgraphs. This is filled lazily when a subgraph is run for the first time.
+    pub(super) systems_in_sets: HashMap<SystemSetKey, FixedBitSet>,
 }
 
 impl SystemSchedule {
     /// Creates an empty [`SystemSchedule`].
-    pub const fn new() -> Self {
+    pub fn new() -> Self {
         Self {
             systems: Vec::new(),
             system_conditions: Vec::new(),
@@ -120,6 +125,7 @@ impl SystemSchedule {
             system_dependents: Vec::new(),
             sets_with_conditions_of_systems: Vec::new(),
             systems_in_sets_with_conditions: Vec::new(),
+            systems_in_sets: HashMap::default(),
         }
     }
 }

crates/bevy_ecs/src/schedule/executor/multi_threaded.rs

@@ -17,7 +17,7 @@ use crate::{
     prelude::Resource,
     schedule::{
         is_apply_deferred, ConditionWithAccess, ExecutorKind, SystemExecutor, SystemSchedule,
-        SystemWithAccess,
+        SystemSetKey, SystemWithAccess,
     },
     system::{RunSystemError, ScheduleSystem},
     world::{unsafe_world_cell::UnsafeWorldCell, World},
@@ -239,6 +239,7 @@ impl SystemExecutor for MultiThreadedExecutor {
         &mut self,
         schedule: &mut SystemSchedule,
         world: &mut World,
+        subgraph: Option<SystemSetKey>,
         _skip_systems: Option<&FixedBitSet>,
         error_handler: ErrorHandler,
     ) {
@@ -253,6 +254,23 @@ impl SystemExecutor for MultiThreadedExecutor {
             .clone_from(&schedule.system_dependencies);
         state.ready_systems.clone_from(&self.starting_systems);
 
+        if let Some(set) = subgraph {
+            // Get the systems in the set
+            let systems_in_set = schedule
+                .systems_in_sets
+                .get(&set)
+                .expect("System set not found in schedule.");
+
+            // Mark all systems in the set as completed (waiting to be flipped)
+            state.completed_systems = systems_in_set.clone();
+            // Flip all bits to get the systems not in the set
+            state.completed_systems.toggle_range(..);
+            // Signal dependents of systems not in the set, as though they had run
+            state.signal_all_dependents();
+            // Only mark systems in the set as ready to run
+            state.ready_systems.intersect_with(systems_in_set);
+        }
+
         // If stepping is enabled, make sure we skip those systems that should
         // not be run.
         #[cfg(feature = "bevy_debug_stepping")]
@@ -795,6 +813,19 @@ impl ExecutorState {
             }
         }
     }
+
+    fn signal_all_dependents(&mut self) {
+        for system_index in self.completed_systems.ones() {
+            for &dep_idx in &self.system_task_metadata[system_index].dependents {
+                let remaining = &mut self.num_dependencies_remaining[dep_idx];
+                debug_assert!(*remaining >= 1);
+                *remaining -= 1;
+                if *remaining == 0 && !self.completed_systems.contains(dep_idx) {
+                    self.ready_systems.insert(dep_idx);
+                }
+            }
+        }
+    }
 }
 
 fn apply_deferred(

crates/bevy_ecs/src/schedule/executor/single_threaded.rs

@@ -13,6 +13,7 @@ use crate::{
     error::{ErrorContext, ErrorHandler},
     schedule::{
         is_apply_deferred, ConditionWithAccess, ExecutorKind, SystemExecutor, SystemSchedule,
+        SystemSetKey,
     },
     system::{RunSystemError, ScheduleSystem},
     world::World,
@@ -57,9 +58,22 @@ impl SystemExecutor for SingleThreadedExecutor {
         &mut self,
         schedule: &mut SystemSchedule,
         world: &mut World,
+        subgraph: Option<SystemSetKey>,
         _skip_systems: Option<&FixedBitSet>,
         error_handler: ErrorHandler,
     ) {
+        if let Some(set) = subgraph {
+            // Get the systems in the set
+            let systems_in_set = schedule
+                .systems_in_sets
+                .get(&set)
+                .expect("System set not found in schedule.");
+            // Mark all systems in the set as completed (waiting to be flipped)
+            self.completed_systems = systems_in_set.clone();
+            // Flip all bits to get the systems not in the set
+            self.completed_systems.toggle_range(..);
+        }
+
         // If stepping is enabled, make sure we skip those systems that should
         // not be run.
         #[cfg(feature = "bevy_debug_stepping")]