Specialize .zip() for efficient slice and slice iteration by bluss · Pull Request #33090 · rust-lang/rust (original) (raw)

The idea is to introduce a private trait TrustedRandomAccess and specialize .zip() for random access iterators into a counted loop.

The implementation in the PR is internal and has no visible effect in the API

Why a counted loop? To have each slice iterator compile to just a pointer, and both pointers are indexed with the same loop counter value in the generated code. When this succeeds, copying loops are readily recognized and replaced with memcpy and addition loops autovectorize well.

TrustedRandomAccess

The TrustedRandomAccess approach works very well on the surface. Microbenchmarks optimize well, following the ideas above, and that is a dramatic improvement of .zip()'s codegen.

// old zip before this PR: bad, byte-for-byte loop // with specialized zip: memcpy pub fn copy_zip(xs: &[u8], ys: &mut [u8]) { for (a, b) in ys.iter_mut().zip(xs) { *a = *b; } }

// old zip before this PR: single addition per iteration // with specialized zip: vectorized pub fn add_zip(xs: &[f32], ys: &mut [f32]) { for (a, b) in ys.iter_mut().zip(xs) { *a += *b; } }

// old zip before this PR: single addition per iteration // with specialized zip: vectorized (!!) pub fn add_zip3(xs: &[f32], ys: &[f32], zs: &mut [f32]) { for ((a, b), c) in zs.iter_mut().zip(xs).zip(ys) { *a += *b * *c; } }

Non-null issues

Yet in more complex situations, the .zip() loop can still fall back to its old behavior where phantom null checks throw in fake premature end of the loop conditionals. Remember that a NULL inside
Option<(&T, &T)> makes it a None value and a premature (in this case)
end of the loop.

So even if we have 1) an explicit Some in the code and 2) the types of the pointers are &T or &mut T which are nonnull, we can still get a phantom null check at that point.

One example that illustrates the difference is copy_zip with slice versus Vec arguments. The involved iterator types are exactly the same, but the Vec version doesn't compile down to memcpy. Investigating into this, the function argument metadata emitted to llvm plays the biggest role. As eddyb summarized, we need nonnull for the loop to autovectorize and noalias for it to replace with memcpy.

There was an experiment to use assume to add a non-null assumption on each of the two elements in the specialized zip iterator, but this only helped in some of the test cases and regressed others. Instead I think the nonnull/noalias metadata issue is something we need to solve separately anyway.

Adaptors

These have conditionally implemented TrustedRandomAccess

These have not implemented it