Tracking Issue for `Iterator::next_chunk`

[rossmacarthur]

View all comments

Feature gate: #![feature(iter_next_chunk)]

This is a tracking issue for the .next_chunk() and .next_chunk_back() method on iterators which allows you to advance the iterator N elements and return an array [T; N]

Public API

use core::array;

trait Iterator {
    type Item;
    
    fn next_chunk<const N: usize>(&mut self) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>>
    where
        Self: Sized;
}

trait DoubleEndedIterator {
    fn next_chunk_back<const N: usize>(&mut self) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>>
    where
        Self: Sized;
}

Steps / History

• Implementation:
  • Add Iterator::next_chunk #93700
  • Optimized vec::IntoIter::next_chunk impl #98553
  • Optimize slice_iter.copied().next_chunk() #103166
  • optimize next_chunk impls for Filter and FilterMap #111781
  • fix Drop items getting leaked in Filter::next_chunk #126879
  • Implement DoubleEndedIterator::next_chunk_back #156737
• Final comment period (FCP)
• Stabilization PR

Unresolved Questions

• Naming: other options include next_array() or next_array_chunk().
• How should we handle N = 0?

[ChayimFriedman2]

Should this method support N=0? itertools does not, and it sounds plausible; however, if we want to do the same, we will have the same concern as <[T]>::array_chunks() that cannot be stabilized until generic_const_exprs or some subset will be stabilized.

[the8472]

I'm wondering about the intended uses of this API.

If it is for performance (autovectorization) then it seems fairly brittle. I just tried doing this benchmark:

#[bench]
fn bench_next_chunk(b: &mut Bencher) {
    let v = vec![13u8; 2048];

    b.iter(|| {
        const CHUNK: usize = 8;

        let mut sum = [0u32; CHUNK];
        let mut iter = black_box(v.clone()).into_iter();

        while let Ok(chunk) = iter.next_chunk::<CHUNK>() {
            for i in 0..CHUNK {
                sum[i] += chunk[i] as u32;
            }
        }

        sum
    })
}

and it took 17,272 ns/iter when compiled with target-cpu=znver2. With a handrolled implementation I got 211 ns/iter. I encountered some 10x slowdowns in a few other combinations. Sometimes codegen-units=1 produced even worse assembly because vectorizer did silly things such as loading one byte at a time into a simd register with vpinsrb.

And that's with direct, linear memory access on a Vec::IntoIter which should be relatively straight-forward to turn into masked simd loads. Many iterators yield references. A [&T; N] would need a gather-load which is expensive and only available on newer CPUs, this is unlike slice::array_chunks which yields &[T; N]. We could try specializing slice.iter().next_chunk() but then you might as well use array_chunks. And of course it'll get more complex with additional adapters.

Another limitation is that it can't be used in a for-in loop or chained further with other iterator adapters, an issue #92393 didn't have.

[Lokathor]

This is probably an example of a case where you should first use copied(), and then chunk it, and then you turn that array into a simd type for the work, before possibly turning it back at the end of your loop.

And the change into and out of simd form will have a cost, so you'll want to ensure you're doing enough work inside each loop pass or it might exceed the cost of the scalar only code.

[the8472]

yeah that was just an example to illustrate that the default implementation is suboptimal and we'll need optimized implementations on most linear-memory sources and length-preserving adapters.

[Lokathor]

Mm, yeah.

Myself I would imagine using this most on iterators that aren't linear memory sources, and "the point" would be to have a clean and unified way to turn some crazy data source into simple SIMD chunks.

[rossmacarthur]

Might be good to align the naming of this function with Iterator::array_chunks, i.e. either

• Iterator::chunks
• Iterator::next_chunk

Or

• Iterator::array_chunks
• Iterator::next_array_chunk

[mark-i-m]

So I'm finding myself caring less and less about the actual naming. I've wanted this feature at least 3 times in the last year and had to write my own version to get around. Can we just pick something and get the feature stabilized?

[ZhennanWu]

I believe the current implementation also suffers from this compiler optimization issue: #110734

Godbolt: https://godbolt.org/z/Te5Wo9eh7

[the8472]

I'm not sure what that issue has to do with next_chunk specifically. It seems to be mostly about the codegen for it.collect::<Vec<_>>(), which I think is easily disturbed by anything mutating the iterator.

[pitaj]

Let's get this stabilized. Here are my answers to the Unresolved Questions:

Naming: other options include next_array() or next_array_chunk().

next_chunk seems like the best option to me.

Should we also add next_chunk_back to DoubleEndedIterator?

Probably, but that should be under a separate feature and not block this.

How should we handle N = 0?

It should be allowed and return a zero-length array without advancing the iterator. No good reason to forbid it.

[the8472]

It should be allowed and return a zero-length array without advancing the iterator. No good reason to forbid it.

For next_chunk it might not be a problem but for the chunks iterators on slices and for iterator.array_chunks() it is. So we might as well do something about it here

[pitaj]

I disagree. For those, it's necessary by definition. There's no reason to apply the same logic here besides, in my opinion, a misguided desire for consistency. The behavior of next_chunk is totally reasonable for array length 0, unlike the others you mentioned.

"Doing something about it" will pretty much require const expressions just like all the other ones, needlessly delaying the stabilization.

[the8472]

Well, it would prevent implementations of next_chunk using those things internally