Learn Rust Series (#70) - Data Parallelism with Rayon and Parallel Iterators

scipio(71)

Published in

StemSocial

Words

3192

Reading

15 min

Listen

Play

Learn Rust Series (#70) - Data Parallelism with Rayon and Parallel Iterators

What will I learn

You will learn what data parallelism is and how rayon turns a sequential iterator parallel;
how par_iter and into_par_iter split work across a thread pool with a one-word change;
how rayon::join expresses fork-join divide-and-conquer parallelism;
how to build the same patterns by hand with thread::scope, so you know what rayon automates;
when parallelism actually pays off, and when the coordination cost makes it a loss.

Requirements

A working modern computer running macOS, Windows or Ubuntu, with Cargo;
An installed Rust toolchain (via rustup, from rustup.rs);
The previous sixty-nine episodes, especially iterators, scoped threads, and slices;
The ambition to learn systems programming from the ground up.

Difficulty

Intermediate

Curriculum (of the `Learn Rust Series`):

Learn Rust Series (#70) - Data Parallelism with Rayon and Parallel Iterators

After two episodes of hand-rolled lock-free structures, here is the good news: most of the time you do not write any of that. When you have a big collection and want to process its elements across all your CPU cores, the rayon crate makes it almost free -- you change iter() to par_iter() and rayon handles the thread pool, the work splitting, and the load balancing for you. It is the friendliest face of Rust concurrency, and (this is the part I really want you to hold on to) it is built on exactly the atomics and scoped threads we have been studying for ten episodes now. Nothing under the hood is magic. So the plan for today is simple: we study rayon's interface, and then we rebuild each pattern in plain std so that when you reach for the crate, you know precisely what work it is doing on your behalf ;-)

Solutions to Episode 69 Exercises

Episode 69 built the Michael-Scott lock-free queue. Here are worked solutions to the three exercises.

Exercise 1 -- an is_empty that reads the dummy's next with a single Acquire load:

use std::sync::atomic::{AtomicPtr, Ordering};
use std::ptr;

struct Node { value: Option<i32>, next: AtomicPtr<Node> }
struct Queue { head: AtomicPtr<Node>, tail: AtomicPtr<Node> }

impl Queue {
    fn new() -> Self {
        let dummy = Box::into_raw(Box::new(Node {
            value: None,
            next: AtomicPtr::new(ptr::null_mut()),
        }));
        Queue { head: AtomicPtr::new(dummy), tail: AtomicPtr::new(dummy) }
    }

    fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::Acquire);
        // If the dummy has no successor, the queue holds no real values.
        unsafe { (*head).next.load(Ordering::Acquire).is_null() }
        // NB: this answer is only a snapshot. The instant it returns, another
        // thread may have enqueued, so `true` means "empty a moment ago",
        // not "empty right now". Under concurrency every observation is stale.
    }
}

fn main() {
    let q = Queue::new();
    println!("{}", q.is_empty()); // true
}

The key insight is that a lock-free reader never gets a "now" -- it gets a "then". The value is correct for the instant of the load and potentially obsolete one nanosecond later, which is exactly why you must never branch on is_empty() and then assume the queue is still empty on the next line.

Exercise 2 -- enqueue from three scoped threads, join them, then drain on the main thread and assert the count. I will lean on the concurrent enqueue we wrote last time and wrap the shared queue in an Arc so the scoped threads can all reach it:

use std::sync::Arc;
use std::thread;

fn main() {
    let queue = Arc::new(MsQueue::new()); // the Michael-Scott queue from ep69
    thread::scope(|s| {
        for id in 0..3 {
            let q = Arc::clone(&queue);
            s.spawn(move || {
                for i in 0..4 {
                    q.enqueue(id * 10 + i); // 4 values per thread, 12 total
                }
            });
        }
    }); // scope join: all three producers have finished here

    let mut drained = 0;
    while queue.dequeue().is_some() {
        drained += 1;
    }
    assert_eq!(drained, 12);
    println!("drained {drained} items"); // drained 12 items
}

The important detail is the barrier: thread::scope (episode 62) does not return until every spawned thread has joined, so by the time we start draining, all twelve enqueues are provably complete and no producer can still be mutating the tail.

Exercise 3 -- why fully concurrent dequeue needs safe reclamation, and a crate that provides it:

// A single-threaded (or single-consumer) drain is safe because only one
// thread ever frees a node. But if MANY threads dequeue at once, thread A
// can Box::from_raw a node while thread B still holds a raw pointer to it
// and is about to dereference it: a use-after-free, and the seed of the
// ABA problem all over again.
//
// The fix is deferred reclamation: do not free a node until every thread
// has provably moved past the point where it could still be looking at it.
// In Rust you do NOT hand-roll this. Reach for `crossbeam-epoch`
// (epoch-based reclamation) or hazard pointers via `haphazard`.

The lesson from episodes 68 and 69 in one sentence: the algorithm is the easy half, safe memory reclamation is the hard half, and the mature move is to delegate the hard half to a vetted crate. Now, on to the crate that makes all of this a non-issue when you do not need a queue at all -- rayon.

par_iter: parallelism with a one-word change

This is rayon's headline, and it is worth pausing on how understated it is. Bring the prelude into scope, change iter() to par_iter(), and the sum is computed across a thread pool. The API deliberately mirrors the standard Iterator you learned back in episode 11, so almost nothing else in your code changes:

// requires the `rayon` crate: shown for illustration, not compiled locally
use rayon::prelude::*;

fn main() {
    let data: Vec<i64> = (0..1_000_000).collect();
    let sum: i64 = data.par_iter().sum(); // split across all cores automatically
    println!("{sum}");
}

Look at what did not change. No thread handles, no Arc<Mutex<...>>, no channels, no join. The type of sum is still i64. If you deleted the par_ prefix it would compile and run identically, just on one core. That is the whole design philosophy of rayon: parallelism as a drop-in, not a rewrite.

Under the hood rayon splits the range into chunks, hands them to a work-stealing thread pool, and combines the partial results. par_iter() borrows (&T items), par_iter_mut() gives &mut T, and into_par_iter() consumes the collection and yields owned T -- exactly parallel to iter, iter_mut, and into_iter. Because the closures you hand these methods run on other threads, rayon requires them to be Send (episode 61); the compiler enforces that for you, so a closure that captures a non-thread-safe value simply will not compile.

What rayon is actually doing: work stealing

The phrase "work-stealing thread pool" is the heart of it, so let me unpack it before we rebuild anything. Rayon starts (lazily, on first use) a pool with one worker thread per logical CPU core. Each worker owns a double-ended queue of tasks. When you call par_iter, rayon does not eagerly cut the data into N equal pieces; instead it splits recursively: it halves the work, and if a worker finds its own queue empty it steals a half-task from the back of a busy worker's queue.

Why bother with stealing instead of a simple equal split? Because equal splits assume every element costs the same amount of work, and real workloads are lumpy. If you split a million items into four fixed quarters and one quarter happens to contain all the expensive elements, three cores finish early and sit idle while the fourth grinds on. Work stealing fixes this automatically: idle cores pull work from busy ones until the whole thing drains. This is dynamic load balancing, and it is the single biggest reason par_iter tends to "just work" on irregular data where a naive by-hand split would leave cores starving.

The same thing by hand

To prove there is no magic, here is a parallel sum in pure std: split the slice into chunks and give each to a scoped thread, then add the partial sums. This is essentially rayon's strategy minus the automatic load balancing -- a fixed, equal split:

use std::thread;

fn parallel_sum(data: &[i64], threads: usize) -> i64 {
    let chunk = (data.len() + threads - 1) / threads; // ceil division
    thread::scope(|s| {
        let handles: Vec<_> = data
            .chunks(chunk)
            .map(|c| s.spawn(move || c.iter().sum::<i64>()))
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).sum()
    })
}

fn main() {
    let data: Vec<i64> = (0..1000).collect();
    println!("{}", parallel_sum(&data, 4)); // 499500
}

Notice how much rayon is saving you. We had to pick the chunk count ourselves, do the ceil-division arithmetic, collect the handles, and join them. And this version is the naive one: it splits equally and hopes the work is even. Rayon's recursive splitting plus stealing would handle the uneven case for free. Having said that, the shape is identical, which is the point -- par_iter is this pattern, generalised and load-balanced.

Parallel map, filter, and collect

The full iterator vocabulary has parallel versions, and they compose exactly like the sequential ones. A filter then map then collect runs each stage across the pool, and crucially collect reassembles the results in the original order even though they were produced out of order on different threads:

// requires the `rayon` crate: shown for illustration, not compiled locally
use rayon::prelude::*;

fn main() {
    let result: Vec<i64> = (0..100i64)
        .into_par_iter()
        .filter(|n| n % 2 == 0)
        .map(|n| n * n)
        .collect();
    println!("{}", result.len()); // 50
}

The std equivalent maps each chunk in parallel and stitches the pieces back together in order. The ordering guarantee that rayon gives you for free is something we have to arrange manually here, by concatenating the per-chunk results in sequence:

use std::thread;

fn parallel_map_square(data: &[i64], threads: usize) -> Vec<i64> {
    let chunk = (data.len() + threads - 1) / threads;
    let parts: Vec<Vec<i64>> = thread::scope(|s| {
        let handles: Vec<_> = data
            .chunks(chunk)
            .map(|c| s.spawn(move || c.iter().map(|x| x * x).collect::<Vec<i64>>()))
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });
    parts.concat() // reassemble chunk results in order
}

fn main() {
    println!("{:?}", parallel_map_square(&[1, 2, 3, 4], 2)); // [1, 4, 9, 16]
}

The parts.concat() at the end is doing the ordering work. Because chunks yields the slices left-to-right and we collect the handles in that same order, joining them in order and concatenating preserves the original sequence. Get that ordering wrong by hand -- say, by using a channel and pushing results as they finish -- and you get a correct set of answers in a scrambled order. Rayon just does not let you make that mistake.

Reductions done right: reduce, sum, and fold

sum() is a special case of a more general operation: a reduction, where you fold many values into one. Rayon parallelises reductions by having each worker reduce its own chunk to a partial result, then combining the partials. For that to be correct the combine step must be associative -- (a op b) op c must equal a op (b op c) -- because rayon decides the grouping, not you. Addition and max are associative; subtraction is not, which is why you can par_iter().sum() but must never expect a parallel "running subtraction" to match the sequential one.

Here is the by-hand version of a parallel max, which makes the associativity requirement concrete: each thread computes the max of its chunk, and we take the max of the maxes:

use std::thread;

fn parallel_max(data: &[i64], threads: usize) -> Option<i64> {
    if data.is_empty() {
        return None;
    }
    let chunk = (data.len() + threads - 1) / threads;
    let partials: Vec<i64> = thread::scope(|s| {
        let handles: Vec<_> = data
            .chunks(chunk)
            .map(|c| s.spawn(move || c.iter().copied().max().unwrap()))
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });
    partials.into_iter().max() // max of the per-chunk maxima
}

fn main() {
    let data: Vec<i64> = vec![3, 9, 1, 7, 4, 8, 2];
    println!("{:?}", parallel_max(&data, 3)); // Some(9)
}

In rayon this whole function collapses to data.par_iter().copied().max(). The reason it is safe to split is precisely that max is associative and has an identity-free combine -- the max of two chunk-maxima is the max of the whole. Whenever you reach for rayon's reduce, ask yourself first: is my combine operation associative? If not, parallelising it will quietly give wrong answers, and no compiler will warn you.

Fork-join with join

For divide-and-conquer, rayon offers join, which runs two closures potentially in parallel and returns both results. It is the primitive that par_iter is itself built on top of:

// requires the `rayon` crate: shown for illustration, not compiled locally
use rayon::join;

fn sum_slice(s: &[i64]) -> i64 { s.iter().sum() }

fn main() {
    let data: Vec<i64> = (0..1000).collect();
    let (left, right) = data.split_at(500);
    let (a, b) = join(|| sum_slice(left), || sum_slice(right)); // both halves, in parallel
    println!("{}", a + b); // 499500
}

The word "potentially" is doing real work in that sentence. join does not promise two threads; it promises that the second closure may be stolen and run on another worker if one is idle. If every core is already busy, join just runs both closures sequentially on the current thread. This is why rayon scales gracefully: it never spawns more work than there are cores to absorb, so recursively join-ing a million times does not create a million threads. It creates a million cheap tasks, and the fixed pool of workers drains them.

You can build the same recursive fork-join in std with a scope that spawns one half and runs the other itself, stopping when a chunk is small enough to do sequentially. That "small enough" threshold matters a lot -- go too fine and the coordination overhead swamps the work:

use std::thread;

fn par_sum(data: &[i64]) -> i64 {
    if data.len() <= 256 {
        return data.iter().sum(); // base case: too small to be worth splitting
    }
    let mid = data.len() / 2;
    let (l, r) = data.split_at(mid);
    thread::scope(|s| {
        let left = s.spawn(|| par_sum(l)); // spawn one half
        let right_sum = par_sum(r);        // run the other half on THIS thread
        left.join().unwrap() + right_sum
    })
}

fn main() {
    let data: Vec<i64> = (0..10_000).collect();
    println!("{}", par_sum(&data)); // 49995000
}

Running the right half on the current thread rather than spawning a second one is the same trick rayon's join uses: never leave the calling thread idle while it waits. The <= 256 base case is our manual version of rayon's automatic decision about when a task is too small to bother splitting. In production you would tune that threshold with the benchmarking from episode 55, not guess it.

Configuring the thread pool

By default rayon uses a global pool sized to your logical core count, which is what you want the vast majority of the time. But sometimes you need control -- to cap threads on a shared server, or to keep a latency-sensitive task off the pool that a batch job is hammering. Rayon lets you build a scoped pool and run work inside it:

// requires the `rayon` crate: shown for illustration, not compiled locally
use rayon::prelude::*;

fn main() {
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(4)
        .build()
        .unwrap();

    let total: i64 = pool.install(|| {
        (0..1_000i64).into_par_iter().map(|n| n * 2).sum()
    });
    println!("{total}"); // 999000
}

Everything inside pool.install(...) uses that four-thread pool instead of the global one. This is the seam where rayon stops hiding its machinery and hands you the knobs -- and it is a good preview of what a custom thread pool with its own work-stealing scheduler looks like from the inside, which is somewhere we are heading soon. For now, the practical advice is: use the global pool unless you have a measured reason not to.

How Python and Go would frame this

If you come from Python, the mental model closest to par_iter is multiprocessing.Pool, because Python's global interpreter lock means threads do not give you CPU parallelism for pure-Python work -- you have to spawn processes:

from multiprocessing import Pool

def square(n):
    return n * n

if __name__ == "__main__":
    with Pool(4) as pool:
        result = pool.map(square, range(100))
    print(len(result))  # 100

That works, but notice the cost that rayon does not pay: multiprocessing copies data across process boundaries by pickling it, and you cannot share a mutable structure without extra machinery. Rayon shares memory directly and safely, because the borrow checker already proved there are no data races.

Go would reach for goroutines and a sync.WaitGroup, splitting the slice by hand much like our std version:

var wg sync.WaitGroup
sums := make([]int, 4)
chunk := len(data) / 4
for i := 0; i < 4; i++ {
    wg.Add(1)
    go func(i int) {
        defer wg.Done()
        for _, v := range data[i*chunk : (i+1)*chunk] {
            sums[i] += v
        }
    }(i)
}
wg.Wait()

Go's version is idiomatic and fast, but you own the correctness: the loop-variable capture, the index arithmetic, the fact that each goroutine writes only its own sums[i] slot (write to a shared slot and Go's race detector, if you remember to run it, will scold you). Rust's answer is to make the whole thing one word and let the type system guarantee the absence of the exact bug Go asks you to avoid by discipline. Different languages, same underlying work-stealing idea; Rust just moves the most error-prone parts from your head into the compiler.

When parallelism actually pays

Here is the crucial judgement, and it is the part beginners skip. Parallelism has overhead -- splitting work, waking threads, coordinating, combining results -- and it only wins when the work per element is large enough to dwarf that cost. A trivial integer sum is memory-bound: the CPU spends its time waiting for data to arrive from memory, not computing, so throwing more cores at it barely helps and can even hurt. Heavy per-element computation is where a thread pool earns its keep:

fn main() {
    let data: Vec<u64> = (0..1_000_000).collect();
    // memory-bound: threads add little here, the bottleneck is memory bandwidth
    let cheap: u64 = data.iter().sum();
    // compute-heavy per element: THIS is where par_iter earns its keep
    let expensive: u64 = data.iter().map(|&n| (0..n % 100).sum::<u64>()).sum();
    println!("{cheap} {expensive}");
}

The cheap line does almost no work per element; parallelising it might even lose to the sequential version once you count the coordination. The expensive line does a nested loop per element, so each item is a meaningful unit of work and the split pays off handsomely. The rule of thumb: parallelise when per-element work is substantial and the collection is large, and do not bother when either is small.

Having said that, rayon's real charm is that trying parallelism costs you almost nothing. Change one word, measure with the Criterion benchmarking from episode 55, and keep the change only if the numbers actually improve. That measured, low-risk approach -- change, measure, keep-or-revert -- is the right way to add any parallelism, and it is a world away from the atomics-and-raw-pointers work we did in episodes 66 through 69. Rayon is where all that low-level understanding pays off: you now know exactly what it is doing, which means you know when to trust it and when to reach lower.

De groeten, en tot de volgende!

This closes our first ten episodes of Phase 5, from the Send/Sync foundations, through channels and locks and atomics, into the hand-built lock-free stack and queue, and finally out the other side to the effortless data parallelism of rayon. If there is one throughline I want you to carry forward, it is this: rayon feels like a holiday after raw pointers precisely because you did the hard episodes first. You are not trusting a black box -- you are trusting a tool whose internals you could now sketch on a napkin. We are not done with the pool machinery, though; there is a lot more to say about schedulers that balance work across cores, and we will get our hands into that next. Bedankt voor het meelezen, en tot de volgende keer ;-)

Exercises

Write a parallel_sum that splits a slice across a chosen number of scoped threads and adds the partial sums, using ceil-division so the last chunk is not dropped when the length is not evenly divisible.
Implement a recursive divide-and-conquer parallel max with thread::scope, a sequential base case for small slices, and the "run one half on the current thread" trick from par_sum.
In a comment, describe one concrete workload where par_iter clearly helps and one where it does not, and explain each answer in terms of per-element work and collection size.

@scipio