Implement parallel contiguous iteration.

benches/benches/bevy_ecs/iteration/contiguous_par_iter_simple.rs

@@ -0,0 +1,78 @@
+use bevy_ecs::prelude::*;
+use bevy_tasks::{ComputeTaskPool, TaskPool};
+use glam::*;
+
+#[derive(Component, Copy, Clone)]
+struct Transform(Mat4);
+
+#[derive(Component, Copy, Clone)]
+struct Position(Vec3);
+
+#[derive(Component, Copy, Clone)]
+struct Rotation(Vec3);
+
+#[derive(Component, Copy, Clone)]
+struct Velocity(Vec3);
+
+#[derive(Component, Copy, Clone, Default)]
+struct Data<const X: u16>(f32);
+pub struct Benchmark<'w>(World, QueryState<(&'w Velocity, &'w mut Position)>);
+
+fn insert_if_bit_enabled<const B: u16>(entity: &mut EntityWorldMut, i: u16) {
+    if i & (1 << B) != 0 {
+        entity.insert(Data::<B>(1.0));
+    }
+}
+
+impl<'w> Benchmark<'w> {
+    pub fn new(fragment: u16) -> Self {
+        ComputeTaskPool::get_or_init(TaskPool::default);
+
+        let mut world = World::new();
+
+        let iter = world.spawn_batch(core::iter::repeat_n(
+            (
+                Transform(Mat4::from_scale(Vec3::ONE)),
+                Position(Vec3::X),
+                Rotation(Vec3::X),
+                Velocity(Vec3::X),
+            ),
+            100_000,
+        ));
+        let entities = iter.into_iter().collect::<Vec<Entity>>();
+        for i in 0..fragment {
+            let mut e = world.entity_mut(entities[i as usize]);
+            insert_if_bit_enabled::<0>(&mut e, i);
+            insert_if_bit_enabled::<1>(&mut e, i);
+            insert_if_bit_enabled::<2>(&mut e, i);
+            insert_if_bit_enabled::<3>(&mut e, i);
+            insert_if_bit_enabled::<4>(&mut e, i);
+            insert_if_bit_enabled::<5>(&mut e, i);
+            insert_if_bit_enabled::<6>(&mut e, i);
+            insert_if_bit_enabled::<7>(&mut e, i);
+            insert_if_bit_enabled::<8>(&mut e, i);
+            insert_if_bit_enabled::<9>(&mut e, i);
+            insert_if_bit_enabled::<10>(&mut e, i);
+            insert_if_bit_enabled::<11>(&mut e, i);
+            insert_if_bit_enabled::<12>(&mut e, i);
+            insert_if_bit_enabled::<13>(&mut e, i);
+            insert_if_bit_enabled::<14>(&mut e, i);
+            insert_if_bit_enabled::<15>(&mut e, i);
+        }
+
+        let query = world.query::<(&Velocity, &mut Position)>();
+        Self(world, query)
+    }
+
+    #[inline(never)]
+    pub fn run(&mut self) {
+        self.1
+            .contiguous_par_iter_mut(&mut self.0)
+            .unwrap()
+            .for_each(|(vs, mut ps)| {
+                for (v, p) in vs.iter().zip(ps.iter_mut()) {
+                    p.0 += v.0;
+                }
+            });
+    }
+}

benches/benches/bevy_ecs/iteration/mod.rs

@@ -1,3 +1,4 @@
+mod contiguous_par_iter_simple;
 mod heavy_compute;
 mod iter_frag;
 mod iter_frag_foreach;
@@ -166,3 +167,19 @@ fn par_iter_simple(c: &mut Criterion) {
         b.iter(move || bench.run());
     });
 }
+
+fn contiguous_par_iter_simple(c: &mut Criterion) {
+    let mut group = c.benchmark_group("contiguous_par_iter_simple");
+    group.warm_up_time(core::time::Duration::from_millis(500));
+    group.measurement_time(core::time::Duration::from_secs(4));
+    for f in [0, 10, 100, 1000] {
+        group.bench_function(format!("with_{f}_fragment"), |b| {
+            let mut bench = par_iter_simple::Benchmark::new(f);
+            b.iter(move || bench.run());
+        });
+    }
+    group.bench_function("hybrid".to_string(), |b| {
+        let mut bench = par_iter_simple_foreach_hybrid::Benchmark::new();
+        b.iter(move || bench.run());
+    });
+}

crates/bevy_ecs/src/query/fetch.rs

@@ -4247,9 +4247,11 @@ impl<C: Component, T: Copy, S: Copy> Copy for StorageSwitch<C, T, S> {}
 
 #[cfg(test)]
 mod tests {
+    use core::sync::atomic::{AtomicUsize, Ordering};
+
     use super::*;
     use crate::change_detection::DetectChanges;
-    use crate::query::Without;
+    use crate::query::{QueryNotDenseError, Without};
     use crate::system::{assert_is_system, Query};
     use bevy_ecs::prelude::Schedule;
     use bevy_ecs_macros::QueryData;
@@ -4661,4 +4663,81 @@ mod tests {
 
         assert_eq!(present, [true; 3]);
     }
+
+    // Tests that contiguous parallel iteration can correctly mutate all
+    // instances of a component in the world.
+    #[test]
+    fn contiguous_par_iter_success_test() {
+        // Declare a couple of components.
+
+        #[derive(Component, PartialEq, Eq, Debug)]
+        pub struct C {
+            id: i32,
+            found: bool,
+        }
+
+        #[derive(Component, PartialEq, Eq, Debug)]
+        pub struct D(i32);
+
+        // Build a task pool.
+        bevy_tasks::ComputeTaskPool::get_or_init(bevy_tasks::TaskPool::new);
+
+        // Spawn a world with a couple of tables.
+        let mut world = World::new();
+        for id in 0..100 {
+            world.spawn(C { id, found: false });
+        }
+        for id in 100..150 {
+            world.spawn((C { id, found: false }, D(id)));
+        }
+
+        // Update every row, and check that the correct number of rows were
+        // matched contiguously.
+        let total_found = AtomicUsize::new(0);
+        let mut contiguous_query = world.query::<&mut C>();
+        contiguous_query
+            .contiguous_par_iter_mut(&mut world)
+            .unwrap()
+            .for_each(|cs| {
+                for c in cs {
+                    c.found = true;
+                    total_found.fetch_add(1, Ordering::Relaxed);
+                }
+            });
+        assert_eq!(total_found.load(Ordering::Relaxed), 150);
+
+        // Check that the query updated every row.
+        let mut check_query = world.query::<&C>();
+        assert!(check_query.iter(&world).all(|c| c.found));
+    }
+
+    // Tests that attempting to contiguously iterate in parallel over a query
+    // that contains sparse sets fails (as the query isn't dense).
+    #[test]
+    fn contiguous_par_iter_failure_test() {
+        // Declare a couple of components, one of which is a sparse set.
+
+        #[derive(Component, Clone, Copy)]
+        struct C;
+
+        #[derive(Component, Clone, Copy)]
+        #[component(storage = "SparseSet")]
+        struct S;
+
+        // Build a task pool.
+        bevy_tasks::ComputeTaskPool::get_or_init(bevy_tasks::TaskPool::new);
+
+        // Spawn a world with those two components.
+        let mut world = World::new();
+        for _ in 0..100 {
+            world.spawn((C, S));
+        }
+
+        // This query should fail, as queries over sparse sets aren't dense.
+        let mut sparse_query = world.query::<(&C, &S)>();
+        assert!(matches!(
+            sparse_query.contiguous_par_iter(&world),
+            Err(QueryNotDenseError(_))
+        ));
+    }
 }

crates/bevy_ecs/src/query/par_iter.rs

@@ -2,6 +2,7 @@ use crate::{
     batching::BatchingStrategy,
     change_detection::Tick,
     entity::{EntityEquivalent, UniqueEntityEquivalentVec},
+    query::{ArchetypeFilter, ContiguousQueryData, QueryContiguousIter},
     world::unsafe_world_cell::UnsafeWorldCell,
 };
 
@@ -155,6 +156,182 @@ impl<'w, 's, D: IterQueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> {
     }
 }
 
+/// A parallel iterator over contiguous query results on a
+/// [`Query`](crate::system::Query).
+///
+/// The
+/// [`Query::contiguous_par_iter`](crate::system::Query::contiguous_par_iter)
+/// and
+/// [`Query::contiguous_par_iter_mut`](crate::system::Query::contiguous_par_iter_mut)
+/// methods create instances of this structure.
+pub struct QueryContiguousParIter<'w, 's, D, F>
+where
+    D: ContiguousQueryData,
+    F: ArchetypeFilter,
+{
+    /// A reference to the world that contains the components that this query
+    /// iterates over.
+    pub(crate) world: UnsafeWorldCell<'w>,
+    /// Scoped access to the world state.
+    pub(crate) state: &'s QueryState<D, F>,
+    /// The tick that corresponds to the previous time this query ran.
+    pub(crate) last_run: Tick,
+    /// The tick that corresponds to the current run of the query.
+    pub(crate) this_run: Tick,
+    /// How matched rows are to be divided among worker threads.
+    pub(crate) batching_strategy: BatchingStrategy,
+}
+
+impl<'w, 's, D, F> QueryContiguousParIter<'w, 's, D, F>
+where
+    D: ContiguousQueryData,
+    F: ArchetypeFilter,
+{
+    /// Returns `None` if `query_state` is not dense, and hence not contiguously iterable.
+    pub(crate) fn new(
+        world: UnsafeWorldCell<'w>,
+        state: &'s QueryState<D, F>,
+        last_run: Tick,
+        this_run: Tick,
+    ) -> Option<Self> {
+        state.is_dense.then(|| Self {
+            world,
+            state,
+            last_run,
+            this_run,
+            batching_strategy: BatchingStrategy::new(),
+        })
+    }
+
+    /// Changes the batching strategy used when iterating.
+    ///
+    /// For more information on how this affects the resultant iteration, see
+    /// [`BatchingStrategy`].
+    pub fn batching_strategy(mut self, strategy: BatchingStrategy) -> Self {
+        self.batching_strategy = strategy;
+        self
+    }
+
+    /// Runs `func` on each contiguous chunk of query results in parallel.
+    ///
+    /// # Panics
+    /// If the [`ComputeTaskPool`] is not initialized. If using this from a
+    /// query that is being initialized and run from the ECS scheduler, this
+    /// should never panic.
+    ///
+    /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool
+    #[inline]
+    pub fn for_each(self, func: impl Fn(D::Contiguous<'w, 's>) + Send + Sync + Clone) {
+        self.for_each_init(|| {}, |_, item| func(item));
+    }
+
+    /// Runs `func` on each query result in parallel on a value returned by `init`.
+    ///
+    /// `init` may be called multiple times per thread, and the values returned may be discarded between tasks on any given thread.
+    /// Callers should avoid using this function as if it were a parallel version
+    /// of [`Iterator::fold`].
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use bevy_utils::Parallel;
+    /// use crate::{bevy_ecs::prelude::Component, bevy_ecs::system::Query};
+    /// #[derive(Component)]
+    /// struct T;
+    /// fn system(query: Query<&T>){
+    ///     let mut queue: Parallel<usize> = Parallel::default();
+    ///     // queue.borrow_local_mut() will get or create a thread_local queue for each task/thread;
+    ///     query.contiguous_par_iter().for_each_init(|| queue.borrow_local_mut(),|local_queue, items| {
+    ///         for _ in items {
+    ///             **local_queue += 1;
+    ///         }
+    ///      });
+    ///
+    ///     // collect value from every thread
+    ///     let entity_count: usize = queue.iter_mut().map(|v| *v).sum();
+    /// }
+    /// ```
+    ///
+    /// # Panics
+    /// If the [`ComputeTaskPool`] is not initialized. If using this from a
+    /// query that is being initialized and run from the ECS scheduler, this
+    /// should never panic.
+    ///
+    /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool
+    pub fn for_each_init<T>(
+        self,
+        init: impl Fn() -> T + Sync + Send + Clone,
+        func: impl Fn(&mut T, D::Contiguous<'w, 's>) + Send + Sync + Clone,
+    ) {
+        let func = |mut init, item| {
+            func(&mut init, item);
+            init
+        };
+
+        #[cfg(any(target_arch = "wasm32", not(feature = "multi_threaded")))]
+        {
+            unsafe {
+                QueryContiguousIter::new(self.world, self.state, self.last_run, self.this_run)
+                    .unwrap()
+                    .fold(init(), func);
+            }
+        }
+
+        #[cfg(all(not(target_arch = "wasm32"), feature = "multi_threaded"))]
+        {
+            let thread_count = bevy_tasks::ComputeTaskPool::get().thread_num();
+            // SAFETY: This method can only be called once per instance of
+            // `QueryContiguousParIter`, which ensures that mutable queries
+            // cannot be executed multiple times at once.  Mutable instances of
+            // `QueryContiguousParIter` can only be created via an exclusive
+            // borrow of a `Query` or a `World`, which ensures that multiple
+            // aliasing `QueryContiguousParIter`s cannot exist at the same time.
+            unsafe {
+                if thread_count <= 1 {
+                    // Just run sequentially.
+                    QueryContiguousIter::new(self.world, self.state, self.last_run, self.this_run)
+                        .unwrap()
+                        .fold(init(), func);
+                    return;
+                }
+
+                // Dispatch to `contiguous_par_fold_init_unchecked_manual` for
+                // parallel iteration.
+                let batch_size = self.get_batch_size(thread_count).max(1);
+                self.state.contiguous_par_fold_init_unchecked_manual(
+                    init,
+                    self.world,
+                    batch_size,
+                    func,
+                    self.last_run,
+                    self.this_run,
+                );
+            }
+        }
+    }
+
+    /// Returns the size of each batch in rows, given a thread count and the
+    /// current batching strategy.
+    fn get_batch_size(&self, thread_count: usize) -> u32 {
+        let max_items = || {
+            let id_iter = self.state.matched_storage_ids.iter();
+            // SAFETY: We only access table metadata.
+            let tables = unsafe { &self.world.storages().tables };
+            id_iter
+                .map(|id| {
+                    // SAFETY: Contiguous iteration can only process tables, so
+                    // we must have a table here.
+                    unsafe { tables[id.table_id].entity_count() }
+                })
+                .max()
+                .map(|v| v as usize)
+                .unwrap_or(0)
+        };
+        self.batching_strategy
+            .calc_batch_size(max_items, thread_count) as u32
+    }
+}
+
 /// A parallel iterator over the unique query items generated from an [`Entity`] list.
 ///
 /// This struct is created by the [`Query::par_iter_many`] method.