Skip to content

Commit 2148518

Browse files
james7132joseph-gioalice-i-cecile
authored
Override QueryIter::fold to port Query::for_each perf gains to select Iterator combinators (#6773)
# Objective After #6547, `Query::for_each` has been capable of automatic vectorization on certain queries, which is seeing a notable (>50% CPU time improvements) for iteration. However, `Query::for_each` isn't idiomatic Rust, and lacks the flexibility of iterator combinators. Ideally, `Query::iter` and friends should be able to achieve the same results. However, this does seem to blocked upstream (rust-lang/rust#104914) by Rust's loop optimizations. ## Solution This is an intermediate solution and refactor. This moves the `Query::for_each` implementation onto the `Iterator::fold` implementation for `QueryIter` instead. This should result in the same automatic vectorization optimization on all `Iterator` functions that internally use fold, including `Iterator::for_each`, `Iterator::count`, etc. With this, it should close the gap between the two completely. Internally, this PR changes `Query::for_each` to use `query.iter().for_each(..)` instead of the duplicated implementation. Separately, the duplicate implementations of internal iteration (i.e. `Query::par_for_each`) now use portions of the current `Query::for_each` implementation factored out into their own functions. This also massively cleans up our internal fragmentation of internal iteration options, deduplicating the iteration code used in `for_each` and `par_iter().for_each()`. --- ## Changelog Changed: `Query::for_each`, `Query::for_each_mut`, `Query::for_each`, and `Query::for_each_mut` have been moved to `QueryIter`'s `Iterator::for_each` implementation, and still retains their performance improvements over normal iteration. These APIs are deprecated in 0.13 and will be removed in 0.14. --------- Co-authored-by: JoJoJet <[email protected]> Co-authored-by: Alice Cecile <[email protected]>
1 parent e581d74 commit 2148518

20 files changed

+298
-199
lines changed

‎benches/benches/bevy_ecs/iteration/iter_frag_foreach.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ impl<'w> Benchmark<'w> {
2929

3030
#[inline(never)]
3131
pub fn run(&mut self) {
32-
self.1.for_each_mut(&mut self.0, |mut data| {
32+
self.1.iter_mut(&mut self.0).for_each(|mut data| {
3333
data.0 *= 2.0;
3434
});
3535
}

‎benches/benches/bevy_ecs/iteration/iter_frag_foreach_sparse.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ impl<'w> Benchmark<'w> {
4040

4141
#[inline(never)]
4242
pub fn run(&mut self) {
43-
self.1.for_each_mut(&mut self.0, |mut data| {
43+
self.1.iter_mut(&mut self.0).for_each(|mut data| {
4444
data.0 *= 2.0;
4545
});
4646
}

‎benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ impl<'w> Benchmark<'w> {
5757

5858
#[inline(never)]
5959
pub fn run(&mut self) {
60-
self.1.for_each_mut(&mut self.0, |mut data| {
60+
self.1.iter_mut(&mut self.0).for_each(|mut data| {
6161
data.0 .0 *= 2.0;
6262
data.1 .0 *= 2.0;
6363
data.2 .0 *= 2.0;

‎benches/benches/bevy_ecs/iteration/iter_frag_foreach_wide_sparse.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ impl<'w> Benchmark<'w> {
6767

6868
#[inline(never)]
6969
pub fn run(&mut self) {
70-
self.1.for_each_mut(&mut self.0, |mut data| {
70+
self.1.iter_mut(&mut self.0).for_each(|mut data| {
7171
data.0 .0 *= 2.0;
7272
data.1 .0 *= 2.0;
7373
data.2 .0 *= 2.0;

‎benches/benches/bevy_ecs/iteration/iter_simple_foreach.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@ impl<'w> Benchmark<'w> {
3636
#[inline(never)]
3737
pub fn run(&mut self) {
3838
self.1
39-
.for_each_mut(&mut self.0, |(velocity, mut position)| {
39+
.iter_mut(&mut self.0)
40+
.for_each(|(velocity, mut position)| {
4041
position.0 += velocity.0;
4142
});
4243
}

‎benches/benches/bevy_ecs/iteration/iter_simple_foreach_sparse_set.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ impl<'w> Benchmark<'w> {
3838
#[inline(never)]
3939
pub fn run(&mut self) {
4040
self.1
41-
.for_each_mut(&mut self.0, |(velocity, mut position)| {
41+
.iter_mut(&mut self.0)
42+
.for_each(|(velocity, mut position)| {
4243
position.0 += velocity.0;
4344
});
4445
}

‎benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ impl<'w> Benchmark<'w> {
5757

5858
#[inline(never)]
5959
pub fn run(&mut self) {
60-
self.1.for_each_mut(&mut self.0, |mut item| {
60+
self.1.iter_mut(&mut self.0).for_each(|mut item| {
6161
item.1 .0 += item.0 .0;
6262
item.3 .0 += item.2 .0;
6363
item.5 .0 += item.4 .0;

‎benches/benches/bevy_ecs/iteration/iter_simple_foreach_wide_sparse_set.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ impl<'w> Benchmark<'w> {
5959

6060
#[inline(never)]
6161
pub fn run(&mut self) {
62-
self.1.for_each_mut(&mut self.0, |mut item| {
62+
self.1.iter_mut(&mut self.0).for_each(|mut item| {
6363
item.1 .0 += item.0 .0;
6464
item.3 .0 += item.2 .0;
6565
item.5 .0 += item.4 .0;

‎benches/benches/bevy_ecs/scheduling/running_systems.rs

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,17 +49,17 @@ pub fn empty_systems(criterion: &mut Criterion) {
4949

5050
pub fn busy_systems(criterion: &mut Criterion) {
5151
fn ab(mut q: Query<(&mut A, &mut B)>) {
52-
q.for_each_mut(|(mut a, mut b)| {
52+
q.iter_mut().for_each(|(mut a, mut b)| {
5353
std::mem::swap(&mut a.0, &mut b.0);
5454
});
5555
}
5656
fn cd(mut q: Query<(&mut C, &mut D)>) {
57-
q.for_each_mut(|(mut c, mut d)| {
57+
q.iter_mut().for_each(|(mut c, mut d)| {
5858
std::mem::swap(&mut c.0, &mut d.0);
5959
});
6060
}
6161
fn ce(mut q: Query<(&mut C, &mut E)>) {
62-
q.for_each_mut(|(mut c, mut e)| {
62+
q.iter_mut().for_each(|(mut c, mut e)| {
6363
std::mem::swap(&mut c.0, &mut e.0);
6464
});
6565
}
@@ -98,20 +98,20 @@ pub fn busy_systems(criterion: &mut Criterion) {
9898

9999
pub fn contrived(criterion: &mut Criterion) {
100100
fn s_0(mut q_0: Query<(&mut A, &mut B)>) {
101-
q_0.for_each_mut(|(mut c_0, mut c_1)| {
101+
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
102102
std::mem::swap(&mut c_0.0, &mut c_1.0);
103103
});
104104
}
105105
fn s_1(mut q_0: Query<(&mut A, &mut C)>, mut q_1: Query<(&mut B, &mut D)>) {
106-
q_0.for_each_mut(|(mut c_0, mut c_1)| {
106+
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
107107
std::mem::swap(&mut c_0.0, &mut c_1.0);
108108
});
109-
q_1.for_each_mut(|(mut c_0, mut c_1)| {
109+
q_1.iter_mut().for_each(|(mut c_0, mut c_1)| {
110110
std::mem::swap(&mut c_0.0, &mut c_1.0);
111111
});
112112
}
113113
fn s_2(mut q_0: Query<(&mut C, &mut D)>) {
114-
q_0.for_each_mut(|(mut c_0, mut c_1)| {
114+
q_0.iter_mut().for_each(|(mut c_0, mut c_1)| {
115115
std::mem::swap(&mut c_0.0, &mut c_1.0);
116116
});
117117
}

‎benches/benches/bevy_ecs/scheduling/schedule.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,19 @@ pub fn schedule(c: &mut Criterion) {
1515
struct E(f32);
1616

1717
fn ab(mut query: Query<(&mut A, &mut B)>) {
18-
query.for_each_mut(|(mut a, mut b)| {
18+
query.iter_mut().for_each(|(mut a, mut b)| {
1919
std::mem::swap(&mut a.0, &mut b.0);
2020
});
2121
}
2222

2323
fn cd(mut query: Query<(&mut C, &mut D)>) {
24-
query.for_each_mut(|(mut c, mut d)| {
24+
query.iter_mut().for_each(|(mut c, mut d)| {
2525
std::mem::swap(&mut c.0, &mut d.0);
2626
});
2727
}
2828

2929
fn ce(mut query: Query<(&mut C, &mut E)>) {
30-
query.for_each_mut(|(mut c, mut e)| {
30+
query.iter_mut().for_each(|(mut c, mut e)| {
3131
std::mem::swap(&mut c.0, &mut e.0);
3232
});
3333
}

‎benches/benches/bevy_ecs/world/world_get.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,7 @@ pub fn world_query_for_each(criterion: &mut Criterion) {
230230

231231
bencher.iter(|| {
232232
let mut count = 0;
233-
query.for_each(&world, |comp| {
233+
query.iter(&world).for_each(|comp| {
234234
black_box(comp);
235235
count += 1;
236236
black_box(count);
@@ -244,7 +244,7 @@ pub fn world_query_for_each(criterion: &mut Criterion) {
244244

245245
bencher.iter(|| {
246246
let mut count = 0;
247-
query.for_each(&world, |comp| {
247+
query.iter(&world).for_each(|comp| {
248248
black_box(comp);
249249
count += 1;
250250
black_box(count);

‎crates/bevy_ecs/src/lib.rs

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,8 @@ mod tests {
343343
let mut results = Vec::new();
344344
world
345345
.query::<(Entity, &A, &TableStored)>()
346-
.for_each(&world, |(e, &i, &s)| results.push((e, i, s)));
346+
.iter(&world)
347+
.for_each(|(e, &i, &s)| results.push((e, i, s)));
347348
assert_eq!(
348349
results,
349350
&[
@@ -388,7 +389,8 @@ mod tests {
388389
let mut results = Vec::new();
389390
world
390391
.query::<(Entity, &A)>()
391-
.for_each(&world, |(e, &i)| results.push((e, i)));
392+
.iter(&world)
393+
.for_each(|(e, &i)| results.push((e, i)));
392394
assert_eq!(results, &[(e, A(123)), (f, A(456))]);
393395
}
394396

@@ -479,7 +481,8 @@ mod tests {
479481
let mut results = Vec::new();
480482
world
481483
.query_filtered::<&A, With<B>>()
482-
.for_each(&world, |i| results.push(*i));
484+
.iter(&world)
485+
.for_each(|i| results.push(*i));
483486
assert_eq!(results, vec![A(123)]);
484487
}
485488

@@ -506,7 +509,8 @@ mod tests {
506509
let mut results = Vec::new();
507510
world
508511
.query_filtered::<&A, With<SparseStored>>()
509-
.for_each(&world, |i| results.push(*i));
512+
.iter(&world)
513+
.for_each(|i| results.push(*i));
510514
assert_eq!(results, vec![A(123)]);
511515
}
512516

@@ -1395,8 +1399,8 @@ mod tests {
13951399
let mut world_a = World::new();
13961400
let world_b = World::new();
13971401
let mut query = world_a.query::<&A>();
1398-
query.for_each(&world_a, |_| {});
1399-
query.for_each(&world_b, |_| {});
1402+
query.iter(&world_a).for_each(|_| {});
1403+
query.iter(&world_b).for_each(|_| {});
14001404
}
14011405

14021406
#[test]

0 commit comments

Comments
 (0)