Fix parallelize workload imbalance

[mratsim]

Currently if we take 40 items divided into 12 threads (AMD Ryzen 7800X, Apple M2 Pro or Intel i5-12600) the partitioning will lead to 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 7 = 3*11 + 7 = 40. The “remainder” thread will have 2.33x more work to do.

This can be quite extreme for example if we have 351 items to split on 32 cores, 351/32 = 10.96, rounded to 10 with integer division. 10*31 = 310, the last core needs to process 41 items, 4.1x more than the others.

This ensures whatever the number of items and the number of threads, the workload varies by at most 1.

This probably explains why not all cores are used while benchmarking (benchmarks TODO). See also Amdahl's law.

Unsafe

Note: I'm not familiar with Rust standard library and parallelism development since October 2016 (my last foray into writing Rust programs).
I've looked into:

• https://doc.rust-lang.org/std/primitive.slice.html
• Rayon chunking: https://github.com/rayon-rs/rayon/blob/2c97943/src/slice/chunks.rs

Note of the chunking routines provides balanced partitioning of a range.

• Either the last chunk can be extremely large like with the current usage of chunks_mut
• or it's not returned like in array_chunks and requires a separate remainder call
• or we get more chunks than our number of threads like with as_chunks

I might very well have missed something though.

[han0110]

I think chunks_mut or chunks take input chunk_size instead of number of chunks (try this rust playground for double check), so I think it wouldn't incur load imbalance, tho it might cause us to call more spawn than we expected (ideally we'd want to call spawn as many times as number of threads, right?).

Also I think when the size of input is large enough this PR would be almost same as before, but the approach in this PR indeed makes it closer to what we expected, so I think we can update the comments and remove the unsafe then merge it.

A version without unsafe might looks like:

pub fn parallelize<T: Send, F: Fn(&mut [T], usize) + Send + Sync + Clone>(v: &mut [T], f: F) {
    // Comments...

    let len = v.len();
    let num_threads = multicore::current_num_threads();
    let chunk_size_lo = len / num_threads;
    let chunk_size_hi = chunk_size_lo + 1;
    let len_hi = (len % num_threads) * chunk_size_hi;
    let (v_hi, v_lo) = v.split_at_mut(len_hi);
    let f = &f;

    multicore::scope(|scope| {
        if chunk_size_hi > 0 {
            for (idx, v) in v_hi.chunks_exact_mut(chunk_size_hi).enumerate() {
                scope.spawn(move |_| f(v, idx * chunk_size_hi));
            }
        }
        if chunk_size_lo > 0 {
            for (idx, v) in v_lo.chunks_exact_mut(chunk_size_lo).enumerate() {
                scope.spawn(move |_| f(v, len_hi + idx * chunk_size_lo));
            }
        }
    });
}

[mratsim]

I think chunks_mut or chunks take input chunk_size instead of number of chunks (try this rust playground for double check), so I think it wouldn't incur load imbalance, tho it might cause us to call more spawn than we expected (ideally we'd want to call spawn as many times as number of threads, right?).

Oh indeed good catch.

In that case for that example, one thread will do 6 iterations, 1 thread will do a single iteration, and the rust will do 3. So we can expect a 33% speedup compared to optimal (threads to either 4 or 3 iterations), there is still a noticeable load imbalance.

Also I think when the size of input is large enough this PR would be almost same as before

For example if we have 2²⁷ = 134M items (size of the trusted setup of Filecoin which AFAIK is the largest), a 112-way split (Intel Xeon Platinum top-end from 2019) would be 1198372*112 + 64

fn main() {
    let v = [(); 1 << 27];
    let num_threads = 112;
    let chunk_size = v.len() / num_threads;
    println!("len: {}", v.chunks(chunk_size).count());
    println!("chunks: {:?}", v.chunks(chunk_size).map(<[_]>::len).collect::<Vec<_>>());
}

So indeed a large input will render these negligeable, however parallelize is used in many places (50+) currently, I expect some input size to be around 20-32 for polynomials as thsoe are the degrees we work with (or it might be [2²⁰, 2³²) and in that case, negligeable).

(ideally we'd want to call spawn as many times as number of threads, right?).

This is the difficult question and the answer is "it depends".

Ideally, you want exactly the same number of spawns as you have threads, this maximize the time making progress and minimize overhead spent in the threadpool (principle behind OpenMP static scheduling).

However, in practice, there are other stuff running on the machine, networking, other programs, maybe nested parallelism, unbalanced workload for example you parallelize drawing a scene but part of it is a white wall and part of it has someone with curly hair, so some threads may finish way earlier than others.

In that case you want way more tasks than threads (fine-grained parallelism) so that if one thread finishes earlier, they find something to do. Intel TBB uses a heuristic magic factor 4 (which is perfect for some workload and is totally inadequate for many workloads but that's another issue).

Thing is currently, we neither have exact number of spawns, or 4x more spawns than threads. We're in the always less efficient valley in-between.

A version without unsafe might looks like:

Ah indeed, didn't think of that, updating the PR.

[CPerezz]

Good catch @han0110 and thanks for the PR @mratsim !

Indeed, a more reallistic range considering the curves we use is [2^26-2^28).
In any case, I agree with Han completelly. If that makes it closer to what we look for, let's improve it.

[han0110]

LGTM! Thanks for sharing so much about parallelization, learned a lot!