Batched ECS Query

[farnoy]

What problem does this solve or what need does it fill?

More optimizations could be done in systems that could iterate over queries in a batched, packed way rather than individually. A clear example would be SIMD.

What solution would you like?

I'd like to be able to use queries in the following way:

world
    .query_filtered::<Batched<(&mut Position, &Velocity)>, Without<FrozenInPlace>>()
    .par_for_each_mut(
        &task_pool,
        4,
        |(mut positions, velocities): (impl DerefMut<&mut [Position]>, impl Deref<&[Velocity]>)| { 
           use core::intrinsics::{assume, likely};
           if unsafe { likely(positions.len() == 4) } {
                unsafe { assume(velocities.len() == 4); } // hopefully done by Bevy
                // some batch fast path
           } else {
                for (mut pos, vel) in positions.iter_mut().zip(velocities.iter()) {
                       // fallback scalar slow path
                }
           }
        }
    );

I'm using Deref & DerefMut informally in this example. I imagine it would be some wrapper like Res, but I don't want to suggest any names.

What alternative(s) have you considered?

Manual buffering to reconstruct this information, but it's hard to implement especially for parallel iteration.

Additional context

For densely stored archetypes, I would like it to return however many spatially contiguous items as it's able to, up to the batch_size setting of course. If the total number of compatible rows for a batched query is equal to batch_size, but it's spread across different archetypes/tables and therefore not contiguous, I would prefer to get N separate batches under batch_size. That is what the &[Component] API would require anyway.

It would also be great if Bevy could prove to the compiler that all of these component-slices are of the same length. That would perhaps improve optimization.

[alice-i-cecile]

This is compelling, and I could see this having sizable performance impacts. Implementation will be fairly complex, but I don't think this will be a controversial feature.

[InBetweenNames]

#1822

[InBetweenNames]

So, thinking about this a little more, to truly get the benefits of an SoA representation for calculations such as position = velocity*time, we'd actually need each element of these vectors to be stored as separate components, e.g., PositionX, PositionY, PositionZ... maybe Bundles might make that a bit more convenient. Otherwise, they get laid out as [x, y, z, x, y, z...] which would make things complicated from a batching sense. You'd need to do some kind of rearrangement to get the data in a form amenable for vectorization, and then rearrange it back when you're done.

Ideally, you'd have the X components stored as [x,x,x,...], Y as [y,y,y...]... and so on. This makes the vectorization straightforward, because you simply load chunks at the native vector width (as OP suggested), and do the math in each SIMD lane.

For example, for position = velocity*time, you could load 8 entities PositionX components in one go (for e.g., AVX):

Px = [x0, x1, ... x7]
Py = [y0, y1, ..., y7]
Pz = [z0, z1, ..., z7]

t is a scalar

Velocity (Vx, Vy, Vz) follows same form as position.

Now, Px' = Px + tVx
Py' = Py + tVy
Pz' = Pz + t*Vz

These are straightforward assembly instructions and bonus, the data for Px', Py', Pz' are already in the correct form to just replace directly in the storage as-is.

References:

https://www.intel.com/content/www/us/en/developer/articles/technical/memory-layout-transformations.html (AoS and SoA in general)
https://www.intel.com/content/www/us/en/developer/articles/technical/how-to-manipulate-data-structure-to-optimize-memory-use-on-32-bit-intel-architecture.html (this specific technique)

[InBetweenNames]

Seems I just can't help myself! Digging around the code a bit, to facilitate this kind of change, we might need to ensure the allocator for the tables are correctly aligned for SIMD if this isn't already done. (I'm not sure what the Rust default is, but if it's just using the system malloc, then usually it's 16 bytes aligned for free). Normally this kind of change is pretty simple to implement and due to the way the ECS is stored, even if you didn't use SIMD, you'd only incur an overhead of a few bytes to achieve the desired "overalignment". For example, AVX2 requires a 32-byte alignment.

The next observation I've made is that if the alignment is good, then this could probably be implemented without too much trouble (of course now that I've said that... :) ). It only really makes sense for Components using Dense storage. I would imagine the Sparse storage implementation would just have an "adapter" to allow it to work, and possibly kick out a warning that what you're doing doesn't really make sense. Or, just panic, because that would be a leaky abstraction? Not sure what the Bevy policy is on that one. Maybe there's some compile-time magic that could make it so doing a "batched" query over a Sparsely backed component would fail to compile. The only reason you'd use this API would be for performance, so this is unlikely to affect a beginner.

It seems, due to the nice design of the Bevy ECS, that most of the work could be accomplished directly within the Query code. From what I can tell, aside from the alignment concerns I brought up earlier, no changes need to be made to the fundamental parts of the ECS to support this.

Regarding the Query API additions, there is a question of what the best way to expose the batches might be. Ideally, since the data is already aligned, we could return a "Batch" type that simply references that data, e.g, Batch<f32, 8>. This type would, provided certain constraints are met, be zero-copy convertable to some kind of SIMD type that would support your vector operations. (I'm not sure what the preferred way of doing SIMD is in Rust, but in C++ land, the xsimd library is really nice to work with and is general over different vector sizes). For a start, we might consider just returning glam vectors (which have the vector size hardcoded unfortunately, so would only work on batch sizes of 2 or 4). Again these are just ideas.

For types that don't have a "direct" SIMD representation, we could leave it to the user to define their own rearrangement.

Example usage, adapted from example in the OP:

//These are just markers for dimensions
struct X;
struct Y;
struct Z;

#[derive(Component)]
struct Position<Dimension>(f32);

#[derive(Component)]
struct Velocity<Dimension>(f32);

// For convenience, something like:
//type PositionBundle = (Position<X>, Position<Y>, Position<Z>) 

fn update_positions<Dimension>(ents: Query<(&mut Position<Dimension>, &Velocity<Dimension>)>, time: Res<Time>)
{
  // This "8"  for the batch size should be provided at compile time, maybe with a bevy_simd::NATIVE_F32_WIDTH constant?
   ents.for_each_mut_batch::<8>::(|mut positions, velocities| {
           if unsafe { likely(positions.len() == 8) } { // positions.len() == velocities.len()

                // I'm being vague with the "simd" type here --
                // the only thing that is important is the support for + and scalar multiplication in this case
                
                //This updates the positions in place.  The "Batch" type maps the result back to the storage,
                // and ideally *is* the storage
                positions = (positions.to_simd() + time*velocities.to_simd()).from_simd();
                
                // Note that .to_simd() would only be available for "SIMD" types, e.g., i32, f32... if you had a Vec3 in your batch,
                // you'd need to do this conversion yourself.  This is the 'rearrangement' I mentioned earlier.  And you'd
                // need to 'rearrange' this data back, too.
                
           } else {
           
                // Necessary, because there's no guarantee that the query results will be
                // a multiple of the batch size.  I think this could happen if there's a query filter, or
                // if we're at the end of the iteration and the number of remaining elements isn't a multiple
                // of the batch size.
                // This is a topic that could be explored a bit more later.
                for (mut pos, vel) in positions.iter_mut().zip(velocities.iter()) {
                       
                       pos = pos + t*vel; //Just scalar code

                }
           }
   });

}

// In main
app.add_system(update_positions<X>);
app.add_system(update_positions<Y>);
app.add_system(update_positions<Z>);

The reason in this example each coordinate is processed separately is that if you process them together, you're going to incur a ton of cache misses from the data access pattern. It's better to "keep on going" down each coordinate than switch from X to Y to Z, which live in different storages and could be quite far in memory from one another.

Obviously this is a lot more complicated than the simple cases we're used to... but again, this is for cases where performance is a concern and we need to bring out the big guns. People should ideally prefer to not write things this way by default. The sample code only really makes sense when you start considering large numbers of entities, etc, that are being iterated over. Iteration could be combined with something like rayon even. I could see simulations being a particularly good use case for that.

On a final note, I would recommend allowing users to specify the batch size they want to work with, as sometimes you get some benefits by using a multiple of the native vector width (e.g., using a batch size of 16 instead of 8 for AVX) due to pipelining.

Sorry about the ramble -- I just find this really interesting, and I think Bevy is in a unique position to capitalize on SIMD this way.

EDIT: one last thing, notice that these are now 3 separate systems in the example. So even without something like rayon, there is opportunity to parallelize these systems as they work on completely different data.

[InBetweenNames]

Just curious, is there interest in a prototype implementation of this? Or benchmarks to show potential speedup under (good) circumstances? I believe I could mock up a benchmark with the existing Bevy. I'd be willing to take this on as a fun little project.

[alice-i-cecile]

I'd be interested in seeing this! I'd try to steer clear of splitting Transform in the proof of concept, just to avoid derailing.

[InBetweenNames]

Agreed! I'll keep the PoC focused on just a hypothetical custom "Position" and "Vector" component :)