1#![allow(missing_docs)]
2pub mod bound;
4pub mod cost;
6#[cfg(feature = "threadteam")]
8pub mod team;
9#[doc(hidden)]
10pub mod thread {
11 use crate::*;
12 struct Slot<T>(core::cell::UnsafeCell<Option<T>>);
15 unsafe impl<T: Send> Sync for Slot<T> {}
17 impl<T> Slot<T> {
18 fn new(value: T) -> Self {
19 Self(core::cell::UnsafeCell::new(Some(value)))
20 }
21 unsafe fn take(&self) -> T {
24 unsafe { (*self.0.get()).take() }.expect("one child per slot")
25 }
26 }
27 #[inline]
30 pub fn join_raw(
31 op_a: impl Send + FnOnce(Par),
32 op_b: impl Send + FnOnce(Par),
33 parallelism: Par,
34 ) {
35 fn implementation<'a>(
36 op_a: &'a mut (dyn Send + FnMut(Par)),
37 op_b: &'a mut (dyn Send + FnMut(Par)),
38 mut parallelism: Par,
39 ) {
40 match parallelism {
41 Par::Seq => {
42 (op_a(parallelism.reborrow()), op_b(parallelism));
43 },
44 #[cfg(feature = "threadteam")]
45 Par::Team(team) => {
46 let n_threads = team.size().get();
47 if n_threads == 1 {
48 (op_a(Par::Team(team)), op_b(Par::Seq));
49 } else {
50 let half = n_threads / 2;
53 let ops = [Slot::new(op_a), Slot::new(op_b)];
54 super::team::note_split(2);
55 let _ = team.split(&[n_threads - half, half], |c, child| {
56 unsafe { (ops[c].take())(Par::Team(child)) }
58 });
59 }
60 },
61 };
62 }
63 let mut op_a = Some(op_a);
64 let mut op_b = Some(op_b);
65 implementation(
66 &mut |parallelism| op_a.take().unwrap()(parallelism),
67 &mut |parallelism| op_b.take().unwrap()(parallelism),
68 parallelism,
69 )
70 }
71 pub fn with_width<R>(par: Par, width: usize, f: impl FnOnce(Par) -> R) -> R {
77 if width >= par.degree() {
78 return f(par);
79 }
80 if width < 2 {
81 return f(Par::Seq);
82 }
83 match par {
84 Par::Seq => f(Par::Seq),
85 #[cfg(feature = "threadteam")]
86 Par::Team(team) => {
87 super::team::note_width(width);
88 struct Once<F, R>(core::cell::UnsafeCell<(Option<F>, Option<R>)>);
89 unsafe impl<F, R> Sync for Once<F, R> {}
92 impl<F, R> Once<F, R> {
93 fn get(&self) -> *mut (Option<F>, Option<R>) {
94 self.0.get()
95 }
96 }
97 let once = Once(core::cell::UnsafeCell::new((Some(f), None)));
98 team.split(&[width], |_, child| {
99 let (f, out) = unsafe { &mut *once.get() };
101 *out = Some((f.take().expect("one child"))(Par::Team(child)));
102 })
103 .expect("2 <= width < size is a valid split");
104 once.0.into_inner().1.expect("the child ran")
105 },
106 #[cfg(not(feature = "threadteam"))]
107 Par::__Lifetime(x, _) => match x {},
108 }
109 }
110 pub struct Ptr<T>(pub *mut T);
112 unsafe impl<T> Send for Ptr<T> {}
113 unsafe impl<T> Sync for Ptr<T> {}
114 impl<T> Copy for Ptr<T> {}
115 impl<T> Clone for Ptr<T> {
116 #[inline]
117 fn clone(&self) -> Self {
118 *self
119 }
120 }
121 #[inline]
124 pub fn parallelism_degree(parallelism: &Par) -> usize {
125 parallelism.degree()
126 }
127
128 #[cfg(feature = "rayon")]
131 pub fn for_each<
132 T: Send,
133 I: rayon::iter::IntoParallelIterator<Item = T>,
134 >(
135 par: Par,
136 n_jobs: usize,
137 iter: I,
138 f: impl Sync + Fn(T),
139 ) where
140 I::Iter: rayon::iter::IndexedParallelIterator,
141 {
142 match par {
143 Par::Seq => {
144 use rayon::iter::plumbing::{Producer, ProducerCallback};
145 struct Each<F>(F);
146 impl<T, F: Fn(T)> ProducerCallback<T> for Each<F> {
147 type Output = ();
148 fn callback<P: Producer<Item = T>>(self, p: P) {
149 p.into_iter().for_each(self.0);
150 }
151 }
152 rayon::iter::IndexedParallelIterator::with_producer(
153 iter.into_par_iter(),
154 Each(f),
155 );
156 },
157 #[cfg(feature = "threadteam")]
158 Par::Team(team) => super::team::for_each(team, n_jobs, iter, f),
159 }
160 }
161 #[cfg(feature = "rayon")]
164 pub fn for_each_with_par<
165 T: Send,
166 I: rayon::iter::IntoParallelIterator<Item = T>,
167 >(
168 par: Par,
169 n_jobs: usize,
170 iter: I,
171 f: impl Sync + Fn(T, Par),
172 ) where
173 I::Iter: rayon::iter::IndexedParallelIterator,
174 {
175 match par {
176 Par::Seq => for_each(Par::Seq, n_jobs, iter, |item| f(item, Par::Seq)),
177 #[cfg(feature = "threadteam")]
178 Par::Team(team) => {
179 super::team::for_each_child(team, n_jobs, iter, |item, child| {
181 f(item, Par::Team(child))
182 })
183 },
184 }
185 }
186
187 #[cfg(feature = "rayon")]
190 pub fn join_sized(
191 par: Par,
192 size_a: usize,
193 op_a: impl Send + FnOnce(Par),
194 op_b: impl Send + FnOnce(Par),
195 ) {
196 let degree = par.degree();
197 if degree < 2 {
198 op_a(Par::Seq);
199 op_b(Par::Seq);
200 return;
201 }
202 let size_a = size_a.clamp(1, degree - 1);
203 let (a, b) = (Slot::new(op_a), Slot::new(op_b));
205 split(par, &[size_a, degree - size_a], &|c, child| {
206 unsafe {
208 if c == 0 {
209 (a.take())(child)
210 } else {
211 (b.take())(child)
212 }
213 }
214 });
215 }
216
217 #[cfg(feature = "rayon")]
221 pub fn split(par: Par, sizes: &[usize], f: &(dyn Sync + Fn(usize, Par))) {
222 core::assert!(sizes.iter().all(|&s| s > 0));
223 match par {
224 Par::Seq => {
225 for c in 0..sizes.len() {
226 f(c, Par::Seq);
227 }
228 },
229 #[cfg(feature = "threadteam")]
230 Par::Team(team) => {
231 if sizes.iter().sum::<usize>() <= team.size().get() {
232 super::team::note_split(sizes.len());
233 let _ = team.split(sizes, |c, child| f(c, Par::Team(child)));
234 } else {
235 for c in 0..sizes.len() {
237 f(c, Par::Seq);
238 }
239 }
240 },
241 }
242 }
243
244 #[inline]
250 pub fn par_split_indices(
251 n: usize,
252 idx: usize,
253 chunk_count: usize,
254 ) -> (usize, usize) {
255 let chunk_size = n / chunk_count;
256 let rem = n % chunk_count;
257 let idx_to_col_start = move |idx| {
258 if idx < rem {
259 idx * (chunk_size + 1)
260 } else {
261 rem + idx * chunk_size
262 }
263 };
264 let start = idx_to_col_start(idx);
265 let end = idx_to_col_start(idx + 1);
266 (start, end - start)
267 }
268
269 #[cfg(test)]
270 mod tests {
271 use super::{for_each, for_each_with_par, join_raw, join_sized, split};
272 use crate::Par;
273 use core::sync::atomic::{AtomicUsize, Ordering};
274
275 struct Leaves {
276 count: AtomicUsize,
277 active: AtomicUsize,
278 peak: AtomicUsize,
279 }
280
281 fn walk(parallelism: Par, leaves: &Leaves) {
284 if parallelism.degree() > 1 {
285 join_raw(
286 |par| walk(par, leaves),
287 |par| walk(par, leaves),
288 parallelism,
289 );
290 } else {
291 leaves.count.fetch_add(1, Ordering::Relaxed);
292 let now = leaves.active.fetch_add(1, Ordering::AcqRel) + 1;
293 leaves.peak.fetch_max(now, Ordering::AcqRel);
294 std::thread::sleep(std::time::Duration::from_millis(5));
295 leaves.active.fetch_sub(1, Ordering::AcqRel);
296 }
297 }
298
299 #[test]
300 fn join_raw_leaves_do_not_exceed_requested_width() {
301 crate::team::with_team(8, |mut par| {
304 for n in [2usize, 3, 5, 6, 7, 8] {
305 let l = Leaves {
306 count: AtomicUsize::new(0),
307 active: AtomicUsize::new(0),
308 peak: AtomicUsize::new(0),
309 };
310 split(par.reborrow(), &[n], &|_c, child| walk(child, &l));
311 assert_eq!(l.count.into_inner(), n, "leaves for n = {n}");
312 assert!(
313 l.peak.into_inner() <= n,
314 "concurrent leaves exceed n = {n}"
315 );
316 }
317 });
318 }
319
320 #[test]
326 fn for_each_region_reaches_the_team() {
327 use std::time::{Duration, Instant};
328 let active = AtomicUsize::new(0);
329 let peak = AtomicUsize::new(0);
330 let deadline = Instant::now() + Duration::from_secs(2);
331 crate::team::with_team(4, |par| {
332 use rayon::prelude::*;
333 for_each(par, 4, (0..4u32).into_par_iter(), |_| {
334 let now = active.fetch_add(1, Ordering::AcqRel) + 1;
335 peak.fetch_max(now, Ordering::AcqRel);
336 while peak.load(Ordering::Acquire) < 2 && Instant::now() < deadline {
337 std::hint::spin_loop();
338 }
339 std::thread::sleep(Duration::from_millis(2));
340 active.fetch_sub(1, Ordering::AcqRel);
341 });
342 });
343 assert!(
344 peak.into_inner() > 1,
345 "no member of the team of 4 entered the region: it ran on the caller alone"
346 );
347 }
348
349 #[test]
352 fn join_sized_children_stay_within_their_share() {
353 for n in [2usize, 3, 5, 8] {
354 crate::team::with_team(n, |mut par| {
355 use rayon::prelude::*;
356 for left in 0..=n + 1 {
357 let active = AtomicUsize::new(0);
358 let peak = AtomicUsize::new(0);
359 let widths = [AtomicUsize::new(0), AtomicUsize::new(0)];
360 let side = |which: usize, par: Par| {
361 let jobs = par.degree();
362 widths[which].store(jobs, Ordering::Relaxed);
363 for_each_with_par(par, jobs, (0..jobs).into_par_iter(), |_, _| {
364 let now = active.fetch_add(1, Ordering::AcqRel) + 1;
365 peak.fetch_max(now, Ordering::AcqRel);
366 std::thread::sleep(std::time::Duration::from_millis(3));
367 active.fetch_sub(1, Ordering::AcqRel);
368 });
369 };
370 join_sized(
371 par.reborrow(),
372 left,
373 |par| side(0, par),
374 |par| side(1, par),
375 );
376 let total = widths[0].load(Ordering::Relaxed)
377 + widths[1].load(Ordering::Relaxed);
378 assert_eq!(total, n, "n = {n}, left = {left}");
379 assert!(
380 peak.load(Ordering::Relaxed) <= n,
381 "concurrent jobs exceed n = {n} with left = {left}"
382 );
383 }
384 });
385 }
386 }
387 }
388}
389pub mod approx;
391pub mod simd;