kime_cpu/
par.rs

1//! Splitting work across threads.
2//!
3//! This is the plain version the reference kernels use: scoped threads per call and a shared
4//! counter that hands out tasks in order. The pinned worker groups from spec/10-cpu.md replace it
5//! when the plan executor lands. Nothing a task computes may depend on which thread runs it, which
6//! is what keeps the kernels deterministic under any thread count.
7
8use std::marker::PhantomData;
9use std::num::NonZero;
10use std::sync::Mutex;
11use std::sync::atomic::{AtomicUsize, Ordering};
12
13/// The number of threads the machine offers.
14#[must_use]
15pub fn available() -> usize {
16    std::thread::available_parallelism().map_or(1, NonZero::get)
17}
18
19/// Runs `f(i)` for every `i` in `0..n` on up to `threads` threads.
20pub fn for_each(n: usize, threads: usize, f: impl Fn(usize) + Sync) {
21    let threads = threads.clamp(1, n.max(1));
22    if threads == 1 {
23        (0..n).for_each(f);
24        return;
25    }
26    let next = AtomicUsize::new(0);
27    let work = || {
28        loop {
29            let i = next.fetch_add(1, Ordering::Relaxed);
30            if i >= n {
31                break;
32            }
33            f(i);
34        }
35    };
36    std::thread::scope(|s| {
37        for _ in 1..threads {
38            s.spawn(work);
39        }
40        work();
41    });
42}
43
44/// `(0..n).map(f).collect()` on up to `threads` threads, in order.
45///
46/// # Panics
47///
48/// If `f` panics.
49pub fn map<T: Send>(n: usize, threads: usize, f: impl Fn(usize) -> T + Sync) -> Vec<T> {
50    let slots: Vec<Mutex<Option<T>>> = (0..n).map(|_| Mutex::new(None)).collect();
51    for_each(n, threads, |i| *slots[i].lock().unwrap() = Some(f(i)));
52    slots.into_iter().map(|s| s.into_inner().unwrap().unwrap()).collect()
53}
54
55/// A buffer many threads write into at once, each to elements no other thread touches.
56#[derive(Debug)]
57pub struct Shared<'a> {
58    ptr: *mut f32,
59    len: usize,
60    _borrow: PhantomData<&'a mut [f32]>,
61}
62
63// SAFETY: Shared is a &mut [f32] split across threads. The only access is `set`, whose contract
64// makes every element written by at most one thread, so sending and sharing it is sound.
65unsafe impl Send for Shared<'_> {}
66// SAFETY: as above.
67unsafe impl Sync for Shared<'_> {}
68
69impl<'a> Shared<'a> {
70    /// Borrows `buf` for the lifetime of the Shared.
71    pub fn new(buf: &'a mut [f32]) -> Self {
72        Self { ptr: buf.as_mut_ptr(), len: buf.len(), _borrow: PhantomData }
73    }
74
75    /// Writes `v` at `i`.
76    ///
77    /// # Safety
78    ///
79    /// No other thread may read or write element `i` while this Shared is alive.
80    ///
81    /// # Panics
82    ///
83    /// If `i` is out of bounds.
84    #[inline(always)]
85    pub unsafe fn set(&self, i: usize, v: f32) {
86        assert!(i < self.len);
87        // SAFETY: in bounds by the assert, and the caller guarantees no other thread touches i.
88        unsafe { self.ptr.add(i).write(v) }
89    }
90}
91
92impl Shared<'_> {
93    /// Elements `start..start + len` as a slice.
94    ///
95    /// # Safety
96    ///
97    /// No other thread may read or write those elements while the slice lives.
98    ///
99    /// # Panics
100    ///
101    /// If the range is out of bounds.
102    #[inline(always)]
103    #[allow(clippy::mut_from_ref)]
104    pub unsafe fn slice_mut(&self, start: usize, len: usize) -> &mut [f32] {
105        assert!(start + len <= self.len);
106        // SAFETY: in bounds by the assert, and the caller guarantees no other thread touches it.
107        unsafe { std::slice::from_raw_parts_mut(self.ptr.add(start), len) }
108    }
109
110    /// Reads element `i`.
111    ///
112    /// # Safety
113    ///
114    /// No other thread may write element `i` while this Shared is alive.
115    ///
116    /// # Panics
117    ///
118    /// If `i` is out of bounds.
119    #[inline(always)]
120    pub unsafe fn get(&self, i: usize) -> f32 {
121        assert!(i < self.len);
122        // SAFETY: in bounds by the assert, and the caller guarantees no other thread writes i.
123        unsafe { self.ptr.add(i).read() }
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn every_task_runs_once() {
133        for threads in [1, 2, 7] {
134            for n in [0, 1, 5, 100] {
135                let got = map(n, threads, |i| i * 2);
136                assert_eq!(got, (0..n).map(|i| i * 2).collect::<Vec<_>>());
137            }
138        }
139    }
140}
kime_cpu/par.rs

kime_cpu/
par.rs
