[Merged by Bors] - Basic adaptive batching for parallel query iteration

benches/benches/bevy_ecs/ecs_bench_suite/heavy_compute.rs

@@ -1,5 +1,5 @@
 use bevy_ecs::prelude::*;
-use bevy_tasks::TaskPool;
+use bevy_tasks::{ComputeTaskPool, TaskPool};
 use glam::*;
 
 #[derive(Component, Copy, Clone)]
@@ -29,8 +29,8 @@ impl Benchmark {
  )
  }));
 
- fn sys(task_pool: Res<TaskPool>, mut query: Query<(&mut Position, &mut Transform)>) {
- query.par_for_each_mut(&task_pool, 128, |(mut pos, mut mat)| {
+ fn sys(mut query: Query<(&mut Position, &mut Transform)>) {
+ query.par_iter_mut().for_each_mut(|(mut pos, mut mat)| {
  for _ in 0..100 {
  mat.0 = mat.0.inverse();
  }
@@ -39,7 +39,7 @@ impl Benchmark {
  });
  }
 
- world.insert_resource(TaskPool::default());
+ world.insert_resource(ComputeTaskPool(TaskPool::default()));
  let mut system = IntoSystem::into_system(sys);
  system.initialize(&mut world);
  system.update_archetype_component_access(&world);

crates/bevy_ecs/src/lib.rs

@@ -59,7 +59,7 @@ mod tests {
  query::{Added, ChangeTrackers, Changed, FilteredAccess, With, Without, WorldQuery},
  world::{Mut, World},
  };
- use bevy_tasks::TaskPool;
+ use bevy_tasks::{ComputeTaskPool, TaskPool};
  use std::{
  any::TypeId,
  sync::{
@@ -376,7 +376,7 @@ mod tests {
  #[test]
  fn par_for_each_dense() {
  let mut world = World::new();
- let task_pool = TaskPool::default();
+ world.insert_resource(ComputeTaskPool(TaskPool::default()));
  let e1 = world.spawn().insert(A(1)).id();
  let e2 = world.spawn().insert(A(2)).id();
  let e3 = world.spawn().insert(A(3)).id();
@@ -385,7 +385,8 @@ mod tests {
  let results = Arc::new(Mutex::new(Vec::new()));
  world
  .query::<(Entity, &A)>()
- .par_for_each(&world, &task_pool, 2, |(e, &A(i))| {
+ .par_iter(&world)
+ .for_each(|(e, &A(i))| {
  results.lock().unwrap().push((e, i));
  });
  results.lock().unwrap().sort();
@@ -398,20 +399,17 @@ mod tests {
  #[test]
  fn par_for_each_sparse() {
  let mut world = World::new();
-
- let task_pool = TaskPool::default();
+ world.insert_resource(ComputeTaskPool(TaskPool::default()));
  let e1 = world.spawn().insert(SparseStored(1)).id();
  let e2 = world.spawn().insert(SparseStored(2)).id();
  let e3 = world.spawn().insert(SparseStored(3)).id();
  let e4 = world.spawn().insert_bundle((SparseStored(4), A(1))).id();
  let e5 = world.spawn().insert_bundle((SparseStored(5), A(1))).id();
  let results = Arc::new(Mutex::new(Vec::new()));
- world.query::<(Entity, &SparseStored)>().par_for_each(
- &world,
- &task_pool,
- 2,
- |(e, &SparseStored(i))| results.lock().unwrap().push((e, i)),
- );
+ world
+ .query::<(Entity, &SparseStored)>()
+ .par_iter(&world)
+ .for_each(|(e, &SparseStored(i))| results.lock().unwrap().push((e, i)));
  results.lock().unwrap().sort();
  assert_eq!(
  &*results.lock().unwrap(),

crates/bevy_ecs/src/query/mod.rs

@@ -2,12 +2,14 @@ mod access;
 mod fetch;
 mod filter;
 mod iter;
+mod par_iter;
 mod state;
 
 pub use access::*;
 pub use fetch::*;
 pub use filter::*;
 pub use iter::*;
+pub use par_iter::*;
 pub use state::*;
 
 #[allow(unreachable_code)]

crates/bevy_ecs/src/query/par_iter.rs

@@ -0,0 +1,135 @@
+use crate::world::World;
+
+use super::{Fetch, QueryFetch, QueryItem, QueryState, ROQueryFetch, ROQueryItem, WorldQuery};
+
+pub struct QueryParIter<'w, 's, Q: WorldQuery, QF: Fetch<'w, State = Q::State>, F: WorldQuery> {
+ pub(crate) world: &'w World,
+ pub(crate) state: &'s QueryState<Q, F>,
+ pub(crate) batch_size: Option<usize>,
+ pub(crate) marker_: std::marker::PhantomData<fn() -> QF>,
+}
+
+impl<'w, 's, Q: WorldQuery, QF, F: WorldQuery> QueryParIter<'w, 's, Q, QF, F>
+where
+ QF: Fetch<'w, State = Q::State>,
+{
+ pub fn batch_size(mut self, batch_size: usize) -> Self {
+ self.batch_size = Some(batch_size);
+ self
+ }
+
+ /// Runs `func` on each query result in parallel.
+ ///
+ /// This can only be called for read-only queries, see [`Self::for_each_mut`] for
+ /// write-queries.
+ ///
+ /// # Panics
+ /// The [`ComputeTaskPool`] resource must be added to the `World` before using this method. 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<FN: Fn(ROQueryItem<'w, Q>) + Send + Sync + Clone>(self, func: FN) {
+ let batch_size = match self.batch_size.or_else(|| self.get_default_batch_size()) {
+ Some(batch_size) => batch_size.max(1),
+ None => return,
+ };
+ // SAFETY: query is read only
+ unsafe {
+ self.state
+ .par_for_each_unchecked_manual::<ROQueryFetch<Q>, FN>(
+ self.world,
+ batch_size,
+ func,
+ self.world.last_change_tick(),
+ self.world.read_change_tick(),
+ );
+ }
+ }
+
+ /// Runs `func` on each query result in parallel.
+ ///
+ /// # Panics
+ /// The [`ComputeTaskPool`] resource must be added to the `World` before using this method. 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_mut<FN: Fn(QueryItem<'w, Q>) + Send + Sync + Clone>(self, func: FN) {
+ let batch_size = match self.batch_size.or_else(|| self.get_default_batch_size()) {
+ Some(batch_size) => batch_size.max(1),
+ None => return,
+ };
+ // SAFETY: query has unique world access
+ unsafe {
+ self.state
+ .par_for_each_unchecked_manual::<QueryFetch<Q>, FN>(
+ self.world,
+ batch_size,
+ func,
+ self.world.last_change_tick(),
+ self.world.read_change_tick(),
+ );
+ }
+ }
+
+ /// Runs `func` on each query result in parallel.
+ ///
+ /// # Panics
+ /// The [`ComputeTaskPool`] resource must be added to the `World` before using this method. If using this from a query
+ /// that is being initialized and run from the ECS scheduler, this should never panic.
+ ///
+ /// # Safety
+ ///
+ /// This does not check for mutable query correctness. To be safe, make sure mutable queries
+ /// have unique access to the components they query.
+ ///
+ /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool
+ #[inline]
+ pub unsafe fn for_each_unchecked<FN: Fn(QueryItem<'w, Q>) + Send + Sync + Clone>(
+ self,
+ func: FN,
+ ) {
+ let batch_size = match self.batch_size.or_else(|| self.get_default_batch_size()) {
+ Some(batch_size) => batch_size.max(1),
+ None => return,
+ };
+ self.state
+ .par_for_each_unchecked_manual::<QueryFetch<Q>, FN>(
+ self.world,
+ batch_size,
+ func,
+ self.world.last_change_tick(),
+ self.world.read_change_tick(),
+ );
+ }
+
+ fn get_default_batch_size(&self) -> Option<usize> {
+ let thread_count = self
+ .state
+ .task_pool
+ .as_ref()
+ .map(|pool| pool.thread_num())
+ .unwrap_or(0);
+ assert!(
+ thread_count > 0,
+ "Attempted to run parallel iteration over a query with an empty TaskPool"
+ );
+ let max_size = if QF::IS_DENSE && <QueryFetch<'static, F>>::IS_DENSE {
+ let tables = &self.world.storages().tables;
+ self.state
+ .matched_table_ids
+ .iter()
+ .map(|id| tables[*id].len())
+ .max()
+ } else {
+ let archetypes = &self.world.archetypes();
+ self.state
+ .matched_archetype_ids
+ .iter()
+ .map(|id| archetypes[*id].len())
+ .max()
+ };
+ max_size.map(|max| max / thread_count)
+ }
+}