Switching par_iter makes memory usage explode (40x)

[BloodStainedCrow]

rayon 1.10.0

I wrote a small Mandelbrot visualizer including this snippet:

let data: Vec<u8> = num_iterations
        .par_iter()
        .map(|num_iterations| (counts[*num_iterations as usize], num_iterations))
        .map(|(counts_cum, num_iterations)| {
            (
                1.0 - f64::from(counts_cum) / (WIDTH * HEIGHT) as f64,
                num_iterations,
            )
        })
        .map(|(brightness, num_iterations)| (brightness.sqrt(), num_iterations))
        .flat_map(|(mut brightness, num_iterations)| {
            if *num_iterations == MAX_ITER {
                brightness = 0.0;
            }

            [
                (brightness * 255.0) as u8,
                (brightness * 255.0) as u8,
                (brightness * 255.0) as u8,
                255u8,
            ]
        })
        .collect();

num_iterations is a Vec of length 10_000 * 10_000, counts a Vec of length 10_000. The largest memory usage I observed was around 380MB which seems to fit the math quite well: 10_000 * 10_000 * 4 bytes ~= 400MB.

I am on a 6/12 core processor, so the increase is not linear with the number of cores as I first suspected.

When I switch to using par_iter The usage spikes to over 16GB (increase of 40x) after a couple of seconds. The runtime is also much slower than single threaded. (I was unable to time it, since it took too long on windows and kept getting OOM killed on Linux).

Incase it is useful the repo with the full code is here.

[cuviper]

Unfortunately, flat_map is kind of a performance footgun, especially to collect. It causes another full round of parallel splitting every time, which is quite wasteful for just your 4 elements. It also prevents this from being an IndexedParallelIterator, so collect can't write results directly into a final Vec -- instead it has to create a bunch of temporary Vec and then merge them later. In theory that would only result in about double the overall memory usage, but the allocator might not be releasing the memory back to the OS for so many differently-sized allocations.

You could use flat_map_iter for the first problem, so that part will run sequentially in each parallel job, but that doesn't totally solve the indexing / temporary-Vec problem. Since we can see ourselves that the result will be exactly 4 * len, you could set that up a bit manually instead, something like:

let mut data = vec![0u8; num_iterations.len() * 4];
data.par_chunks_mut(4).zip(num_iterations).for_each(|(data, num_iterations)| {
    // fill in this chunk
});

[BloodStainedCrow]

Hmm, that makes sense. Thanks for the explanation. This seems quite unfortunate, especially since being a drop-in replacement for iter is the main selling point of par_iter and flat_map is (in my experience) often used in this manner. Especially since this performance footgun is this large.

But without specialization, I doubt there is anything that can be done to address this, without also impacting using flat_map with large (and efficiently parallelizable) Iterators. Unfortunate :(

[cuviper]

The naive flat_map inner splitting comes from this line:

rayon/src/iter/flat_map.rs

Line 127 in 7b3ef64

let result = par_iter.drive_unindexed(consumer);

That could possibly be improved by keeping the same plumbing::Splitter instance from the outer to inner part, but I don't see any way to do that with any generic ParallelIterator. Even if that worked, it would at most achieve the same thing as flat_map_iter, still not as good as you can do with direct data writes.

It would be really cool if FlatMap could tell when the inner iterators have statically-known length, and then it could be a full IndexedParallelIterator. The standard library has an internal trait ConstSizeIntoIterator for this, used with real specialization, but maybe we can do similar even with our opt_len pseudo-specialization. I'm not sure though, because that would also need to be able to split at arbitrary points, including the middle of flattened iterators.

These are hard problems to make the user side easy. :) Did the workaround that I suggested fare better for you?

[BloodStainedCrow]

Yes. I switched from flat_map to flat_map_iter. Determining the used amount of memory was suddenly hard, since the program finished this step of the algorithm in fractions of a second (instead of minutes), but I think it is somewhere in the neighborhood of 700MB at most, which seems in line with your predictions. After that worked so well, I did not try preallocating the Vec and passing mutable chunks.

[BloodStainedCrow]

If I understand it correctly, opt_len does not help, for the same reason ExactSizeIterator does not help. It guarantees that we know the size of any one iterator, but unlike ConstSizeIntoIterator does not guarantee ALL of them to be of this size.

[BloodStainedCrow]

It might be possible by adding a way to (if possible) get a len from the type directly (not an instance). For example adding an associated function IndexedParallelIterator::const_len() -> Option(usize)
Since this function would return the same length for all iterators (if possible).
So a parallel Iterator for an [u32; 10] would return Some(10), a parallel Iterator for a slice would return None, since the size of a slice is not known at compile time.

With that it would be possible to do a sort of runtime specialization.

[BloodStainedCrow]

Crap, I fell into the same trap as the original proposal of ConstSizeIterator. The size of an Iterator changes...

This would need to be a function on IntoParallelIterator.

[BloodStainedCrow]

I have experimented a bit with my Idea for a possible const_len function. I have managed to convert flat_map such that it now used the direct write version of collect (special_extend).

This has improved the memory footprint (and therefore also runtime, since most of the runtime seems to be memory allocation).
Where before, running flat_map on iterators with 10_000 elements (and 4 element arrays as sub-iterators) would consistently get the program SIGKILLed on Linux, it now uses approximately the same amount of memory as flat_map_iter, which seems like a step in the right direction.

This does not fix the slow runtime of flat_map, it is 15x slower than flat_map_iter in my test (which is not representative at all), but that would be because the issue of splitting these tiny tasks is of course still a problem.

But I guess 15x slower vs DNF is a start

[cuviper]

If I understand it correctly, opt_len does not help, for the same reason ExactSizeIterator does not help. It guarantees that we know the size of any one iterator, but unlike ConstSizeIntoIterator does not guarantee ALL of them to be of this size.

It doesn't help the inner iterator, but I meant that the outer FlatMap::opt_len can return a proper length if the inner iterator has a constant length, i.e. runtime_len * const_inner_len.

[BloodStainedCrow]

That is exactly what I did: https://github.com/BloodStainedCrow/rayon/blob/main/src/iter/flat_map.rs#L47.

If I understand opt_len, it does not guarantee that the inner len is constant, right?
That's why I added IntoParallelIterator::const_length

[cuviper]

If I understand opt_len, it does not guarantee that the inner len is constant, right?

Definitely not, since it can examine &self for separate instances. Even a new constant method can't really guarantee that, especially as a safe trait to implement, but we can still make that a condition that will likely panic if violated.

That's why I added IntoParallelIterator::const_length

This might be better on ParallelIterator, since we don't need to worry about Iterator::next consumption like the standard library. This way it still works after into_par_iter() and with many adaptors, and FlatMap can still call the associated <PI as IntoParallelIterator>::Iter::const_len() or whatever.

[cuviper]

Crap, I fell into the same trap as the original proposal of ConstSizeIterator. The size of an Iterator changes...

For ParallelIterator, I don't think it does -- did this bite you in some way?

[BloodStainedCrow]

You are right, I thought I had to implement it on IntoIterator for that reason, implementing it on Iterator does make more sense here.

see below

[BloodStainedCrow]

The same assumptions about the length being correct should apply to const_length and to opt_len. If they are wrong, the program may crash, or do other nonsensical behavior, but no memory unsafety.

So as long as all "base" consumers, which use unsafe do safety checks, everything should be fine

[BloodStainedCrow]

The issue with moving it to the ParallelIterator is, that we do not have seperate types for ParallelIterators of e.g. Arrays of different length. So an associated function like const_length cannot distinguish [0; 100].into_par_iter() from [0;10].into_par_iter().

It is both a slice::IterMut or whatever. So the type level information of the length is lost. That was another reason to put it in the IntoParallelIterator trait.

We can distinguish different "Instances" but before we start consuming the iterators we do not have an instance yet. And at that point it is too late