ogeom_core/
parallel.rs

1//! Deterministic parallelism for the kernel's embarrassingly parallel
2//! stages.
3//!
4//! The rule that makes parallelism admissible here at all: **the answer must
5//! be bit-identical at any thread count.** [`map_ordered`] guarantees it
6//! structurally: each item is computed independently from shared read-only
7//! input, results are collected in item order, and nothing about scheduling
8//! can reach the output. A stage that cannot meet that bar stays sequential.
9//!
10//! The thread count comes from [`threads`]: the machine's parallelism by
11//! default, overridable process-wide with [`set_threads`], including down
12//! to one, which is also what tiny workloads collapse to on their own.
13//! Worker threads re-install the caller's progress watch, so cancellation
14//! reaches into the workers.
15
16use std::sync::atomic::{AtomicUsize, Ordering};
17
18use crate::progress;
19
20/// 0 means "ask the machine".
21static THREADS: AtomicUsize = AtomicUsize::new(0);
22
23/// The thread count parallel stages will use.
24#[must_use]
25pub fn threads() -> usize {
26    let configured = THREADS.load(Ordering::Relaxed);
27    if configured != 0 {
28        return configured;
29    }
30    std::thread::available_parallelism().map_or(1, std::num::NonZero::get)
31}
32
33/// Set the process-wide thread count for parallel stages. `0` restores the
34/// machine default. The answer never depends on this; only the wall clock
35/// does.
36pub fn set_threads(count: usize) {
37    THREADS.store(count, Ordering::Relaxed);
38}
39
40/// Map `f` over `items` on up to [`threads`] scoped threads, returning
41/// results in item order. `f` receives the item index and the item.
42///
43/// Determinism holds by construction: items are computed independently and
44/// results placed by index, so the output is identical at any thread count.
45/// The caller's progress watch is re-installed in every worker; `f` may
46/// checkpoint through it.
47pub fn map_ordered<T, R>(items: &[T], f: impl Fn(usize, &T) -> R + Sync) -> Vec<R>
48where
49    T: Sync,
50    R: Send,
51{
52    let workers = threads().clamp(1, items.len().max(1));
53    if workers <= 1 || items.len() <= 1 {
54        return items.iter().enumerate().map(|(i, t)| f(i, t)).collect();
55    }
56
57    let snapshot = progress::snapshot();
58    // Work is *taken*, not dealt: expensive items cluster (one spline-heavy
59    // face's edges sit adjacent in a reader's job list), and a worker dealt
60    // that region as a contiguous chunk finishes last while the rest idle.
61    // Each worker pulls the next undone index instead, so the wall clock
62    // tracks the total work rather than the heaviest deal. The answer cannot
63    // tell the difference: every index is computed by the same call exactly
64    // once, and the merge reassembles by index, so the output is the item
65    // order however the indices were claimed.
66    let next = std::sync::atomic::AtomicUsize::new(0);
67    let mut parts: Vec<Vec<(usize, R)>> = Vec::with_capacity(workers);
68    std::thread::scope(|scope| {
69        let mut handles = Vec::with_capacity(workers);
70        for _ in 0..workers {
71            let f = &f;
72            let next = &next;
73            let snapshot = snapshot.clone();
74            handles.push(scope.spawn(move || {
75                progress::with_snapshot(snapshot.as_ref(), || {
76                    let mut mine = Vec::new();
77                    loop {
78                        let i = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
79                        let Some(item) = items.get(i) else { break };
80                        mine.push((i, f(i, item)));
81                    }
82                    mine
83                })
84            }));
85        }
86        for handle in handles {
87            match handle.join() {
88                Ok(part) => parts.push(part),
89                Err(panic) => std::panic::resume_unwind(panic),
90            }
91        }
92    });
93    let mut indexed: Vec<(usize, R)> = parts.into_iter().flatten().collect();
94    indexed.sort_unstable_by_key(|(i, _)| *i);
95    indexed.into_iter().map(|(_, r)| r).collect()
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn order_is_item_order_at_any_thread_count() {
104        let items: Vec<usize> = (0..137).collect();
105        let serial: Vec<usize> = items.iter().map(|x| x * 3).collect();
106        for count in [1, 2, 7] {
107            set_threads(count);
108            let parallel = map_ordered(&items, |i, x| {
109                assert_eq!(i, *x);
110                x * 3
111            });
112            assert_eq!(parallel, serial);
113        }
114        set_threads(0);
115    }
116
117    #[test]
118    fn cancellation_reaches_the_workers() {
119        let watch = progress::Watch::new();
120        watch.canceller().cancel();
121        set_threads(4);
122        let items: Vec<usize> = (0..64).collect();
123        let outcomes = progress::watched(&watch, || {
124            map_ordered(&items, |_, _| progress::checkpoint())
125        });
126        set_threads(0);
127        assert!(outcomes.iter().all(std::result::Result::is_err));
128    }
129}
ogeom_core/parallel.rs

ogeom_core/
parallel.rs
