p3_maybe_rayon/task_size.rs
1//! Minimum task sizes for parallel loops.
2//!
3//! Splitting a loop across workers buys throughput and costs latency.
4//! Both sides of that trade collapse into one number, the fewest items a task may hold:
5//!
6//! ```text
7//! floor >= total items -> one task
8//! floor < total items -> total / floor tasks
9//! ```
10//!
11//! A floor lets a call site write its loop body once.
12//! Branching on a length writes the body twice, and can only ask for one task or for all.
13//!
14//! # Entering a loop
15//!
16//! One task is not yet a serial loop.
17//! rayon reaches a body through its own bridge and consumer, even with nothing to split.
18//!
19//! That costs about 1 ns per item, which is a tenth of a fold of a few field operations.
20//!
21//! Two terminals take the leaf iterator directly once the floor forbids a split:
22//!
23//! - [`TaskSizeExt::for_each_min_task_bytes`]
24//! - [`TaskSizeExt::map_collect_min_task_bytes`]
25//!
26//! A loop written with either runs as a plain serial loop, with no bridge to pay for.
27//!
28//! [`TaskSizeExt::with_min_task_bytes`] floors the split and does nothing else.
29//! Use it where one item is a whole chunk, whose own work dwarfs the per-item cost.
30//!
31//! # Cost model
32//!
33//! One item is charged by the bytes it moves, at a fixed rate per byte.
34//!
35//! Moved bytes are counted as every element the body reads plus every element it writes.
36//! A read-modify-write of one element is therefore two, not one.
37//!
38//! - Bytes track the element type, so the floor follows it without being restated.
39//! - A packed extension element costs 32 times a 4-byte prime field element.
40//! - A task therefore needs 32 times fewer of the wide items than of the narrow ones.
41//!
42//! The rate assumes a body with instruction-level parallelism, as a streaming loop has.
43//! A body that waits on its own previous step charges a multiple of its width.
44//!
45//! Charging too little is safe for how a split loop is cut up.
46//! It is not safe for whether the loop splits at all, which is why a slow body says so.
47//!
48//! # Pool size
49//!
50//! The time a loop must be worth scales with the number of workers.
51//! Waking a wider pool costs more, so a wider pool demands a longer loop.
52//!
53//! A pool of one worker never splits at all, which is also what a build without rayon gets.
54//!
55//! # Overrides
56//!
57//! - `P3_MIN_PARALLEL_NS` fixes the serial time a loop must be worth, ignoring the pool size.
58//! - `P3_MAX_TASK_NS` sets the time budget one task holds once a loop does split.
59//! - Setting both to `0` drops every floor to a single item.
60//! - The default gate is calibrated separately for Linux and macOS dispatch costs.
61//!
62//! A floor of one is rayon's own unbounded splitting, which is what an A/B run of a loop wants.
63//!
64//! A site that uses the floor as its own chunk length reads that same `0` as a chunk of one item.
65//!
66//! That shape is slower than either arm, and it is not one the model ever picks.
67//!
68//! The packed fill behind a table of ascending powers is the site to watch.
69//!
70//! It pays one exponentiation per chunk, so a chunk of one item pays one per item.
71//!
72//! On 32 Linux workers that setting costs the fill 60x at 2^12 powers and about 4x at 2^22.
73//!
74//! Twiddle tables and query power tables both come from it, so a whole-prover A/B mostly times it.
75//!
76//! A lower bound on the chunk does not rescue the setting.
77//!
78//! One fill round as the bound wins back 2 to 3x at 2^22 and nothing at 2^12.
79//!
80//! Such a site is measured against its own body rather than through these knobs.
81//!
82//! A caller that would rather not depend on the ambient environment sets both from code.
83//!
84//! Either way the budgets are fixed once and then hold for the rest of the process.
85//! The first parallel loop fixes them, so a change after that point is refused.
86//!
87//! # Examples
88//!
89//! ```
90//! use p3_maybe_rayon::prelude::*;
91//!
92//! let data = vec![1u64; 1 << 10];
93//!
94//! // An adapter changes only how the loop is divided, never what it covers.
95//! let sum: u64 = data.par_iter().with_min_task::<u64>().sum();
96//! assert_eq!(sum, 1 << 10);
97//!
98//! // A call site that picks its own chunk length asks for the floor directly.
99//! // The answer moves with the pool and with any override, so only its range is fixed.
100//! let chunk = min_task_len(data.len(), size_of::<u64>());
101//! assert!((1..=data.len()).contains(&chunk));
102//! ```
103
104use alloc::vec::Vec;
105
106#[cfg(feature = "parallel")]
107use rayon::iter::plumbing::{Producer, ProducerCallback};
108
109use super::prelude::*;
110
111/// Picoseconds one core spends per byte a memory-bound loop moves.
112///
113/// Measured on the fold kernel `a[i] += (b[i] - a[i]) * r`, over buffers past L2.
114/// One item there reads two elements and writes one, so it moves three:
115///
116/// ```text
117/// 4-byte prime field : 0.75 ns / 12 B -> 62 ps per byte
118/// 16-byte extension : 6.4 ns / 48 B -> 133 ps per byte
119/// ```
120///
121/// The midpoint keeps both within a factor of 1.6, which a floor can absorb.
122const PICOS_PER_BYTE: u64 = 100;
123
124/// Serial time a loop must be worth per worker in the pool, before splitting pays.
125///
126/// Dispatching a fan-out costs roughly 0.4 us per worker:
127///
128/// ```text
129/// 16 workers : 7 us per dispatch
130/// 32 workers : 13 us per dispatch
131/// ```
132///
133/// - Asking 0.625 us per worker demands a loop the model prices at 1.5 dispatches.
134/// - Scaling with the pool is what keeps a 96-core host from splitting at a loss.
135///
136/// ```text
137/// 16 workers : 10 us before a split is worth it
138/// 32 workers : 20 us
139/// 96 workers : 60 us
140/// ```
141///
142/// How much real work that is depends on how closely a body matches the rate.
143/// A narrow-field fold is priced about 1.6x high, so its gate sits nearer one dispatch.
144///
145/// Linear in the worker count down to two.
146///
147/// Each row is the length at which the same fold first beats itself run whole:
148///
149/// ```text
150/// workers break-even serial work per worker
151/// 2 0.76 us 0.38 us
152/// 4 1.51 us 0.38 us
153/// 8 3.00 us 0.37 us
154/// 16 6.02 us 0.38 us
155/// 32 12.06 us 0.38 us
156/// ```
157///
158/// Non-macOS default, calibrated on Linux through 96 workers.
159#[cfg(not(target_os = "macos"))]
160const MIN_PARALLEL_PICOS_PER_WORKER: u64 = 625_000;
161
162/// Serial time a loop must be worth per worker in the pool, on Apple silicon.
163///
164/// Dispatching a fan-out costs 1.7 to 2.0 us per worker, four to five times the Linux price.
165///
166/// Measured on an M4 Pro, forcing a split that the gate would otherwise refuse:
167///
168/// ```text
169/// workers forced split per worker
170/// 2 3.7 us 1.85 us
171/// 4 6.7 us 1.68 us
172/// 8 16 us 2.00 us
173/// ```
174///
175/// - Asking 2.5 us per worker demands a loop the model prices at 1.25 to 1.5 dispatches.
176/// - The Linux gate asks for 1.56, so this reaches that margin but is thinner at eight workers.
177///
178/// The value is a gate, not the dispatch cost itself.
179///
180/// Break-even at 14 workers fell between 21 and 52 us of serial work.
181///
182/// The gate lands inside that bracket:
183///
184/// ```text
185/// 14 workers : 35 us before a split is worth it
186/// ```
187#[cfg(target_os = "macos")]
188const MIN_PARALLEL_PICOS_PER_WORKER: u64 = 2_500_000;
189
190/// Time one task holds once a loop does split, in picoseconds.
191///
192/// - The cap keeps a task's working set inside L1 or L2.
193/// - It also bounds how finely work stealing may cut, which nothing else does.
194///
195/// ```text
196/// 2 us at 100 ps per byte = 20 KiB moved per task
197/// ```
198///
199/// A floor only ever forbids a split, so it never produces a task rayon would not have made.
200///
201/// An uncontended loop stops at roughly one leaf per worker on its own.
202///
203/// Stealing a task resets that budget, so a contended loop keeps halving until the floor stops it.
204///
205/// Measured on the same fold kernel, tasks of that size beat both alternatives:
206///
207/// ```text
208/// vs one task per worker : up to 1.3x faster
209/// vs unbounded splitting : up to 2.8x faster
210/// ```
211const MAX_TASK_PICOS: u64 = 2_000_000;
212
213/// The two time budgets that shape every floor.
214#[derive(Clone, Copy)]
215struct Budget {
216 /// Serial time below which a loop stays on one core.
217 ///
218 /// Absent means scale it with the number of workers.
219 min_parallel_picos: Option<u64>,
220 /// Time one task holds once a loop splits.
221 max_task_picos: u64,
222}
223
224impl Budget {
225 /// The compiled-in budgets, used when the environment says nothing.
226 const DEFAULT: Self = Self {
227 min_parallel_picos: None,
228 max_task_picos: MAX_TASK_PICOS,
229 };
230}
231
232/// The budgets a caller may fix in place of the environment, in nanoseconds.
233///
234/// A build without rayon never splits, so it has nothing to configure and offers no setter.
235#[cfg(feature = "parallel")]
236#[derive(Clone, Copy, Debug, PartialEq, Eq)]
237pub struct TaskBudget {
238 /// Serial time a loop must be worth before splitting pays.
239 ///
240 /// Absent means scale it with the number of workers.
241 pub min_parallel_ns: Option<u64>,
242 /// Time one task holds once a loop splits.
243 pub max_task_ns: u64,
244}
245
246#[cfg(feature = "parallel")]
247impl Default for TaskBudget {
248 fn default() -> Self {
249 Self {
250 min_parallel_ns: None,
251 max_task_ns: MAX_TASK_PICOS / 1_000,
252 }
253 }
254}
255
256/// Returned when the budgets are already fixed and can no longer be changed.
257#[cfg(feature = "parallel")]
258#[derive(Clone, Copy, Debug, PartialEq, Eq)]
259pub struct TaskBudgetLocked;
260
261/// The process-wide budgets, fixed by whichever of a caller or the first loop gets there first.
262#[cfg(feature = "parallel")]
263static BUDGET: std::sync::OnceLock<Budget> = std::sync::OnceLock::new();
264
265/// Fixes the budgets for the rest of the process, in place of the environment.
266///
267/// Call this before the first parallel loop runs, since that is what otherwise fixes them.
268///
269/// # Errors
270///
271/// Fails, changing nothing, once the budgets are fixed.
272#[cfg(feature = "parallel")]
273pub fn set_task_budget(budget: TaskBudget) -> Result<(), TaskBudgetLocked> {
274 // Nanoseconds at the boundary, picoseconds inside, so a caller never sees the finer unit.
275 BUDGET
276 .set(Budget {
277 min_parallel_picos: budget.min_parallel_ns.map(|ns| ns.saturating_mul(1_000)),
278 max_task_picos: budget.max_task_ns.saturating_mul(1_000),
279 })
280 .map_err(|_| TaskBudgetLocked)
281}
282
283/// Reads the budgets from the environment, falling back to the compiled-in values.
284///
285/// The environment is read once, then cached for the process.
286/// A floor is computed on entry to every parallel loop, so the read cannot be repeated.
287#[cfg(feature = "parallel")]
288fn budget() -> Budget {
289 *BUDGET.get_or_init(|| {
290 // Overrides are given in nanoseconds and stored in picoseconds.
291 //
292 // A budget is a tuning knob, so a typo falls back instead of taking the process down.
293 let picos = |name: &str| {
294 std::env::var(name)
295 .ok()
296 .and_then(|v| v.parse::<u64>().ok())
297 .map(|ns| ns.saturating_mul(1_000))
298 };
299 // Start from the compiled-in budgets, then apply whatever the environment states.
300 // Each budget is overridden on its own, so one variable never disturbs the other.
301 let mut budget = Budget::DEFAULT;
302
303 // An override fixes the gate outright, which is what makes an A/B run reproducible.
304 if let Some(picos) = picos("P3_MIN_PARALLEL_NS") {
305 budget.min_parallel_picos = Some(picos);
306 }
307 if let Some(picos) = picos("P3_MAX_TASK_NS") {
308 budget.max_task_picos = picos;
309 }
310 budget
311 })
312}
313
314/// The compiled-in budgets.
315///
316/// A build without rayon has no environment to read, and reports a single worker below.
317#[cfg(not(feature = "parallel"))]
318const fn budget() -> Budget {
319 Budget::DEFAULT
320}
321
322/// Estimated time to move one item, in picoseconds.
323///
324/// A zero-sized item is charged one byte, so that every item carries some cost.
325const fn item_picos(item_bytes: usize) -> u64 {
326 let bytes = if item_bytes == 0 {
327 1
328 } else {
329 item_bytes as u64
330 };
331 bytes.saturating_mul(PICOS_PER_BYTE)
332}
333
334/// Minimum items per task, from budgets and a worker count given explicitly.
335///
336/// Split from the public entry points so tests can pin the arithmetic without a thread pool.
337///
338/// # Returns
339///
340/// A count in `1..=max(len, 1)`, so a chunked call site can divide by it without checking.
341fn min_task_len_with(budget: Budget, threads: usize, len: usize, item_bytes: usize) -> usize {
342 // Phase 1: a pool that cannot hand work anywhere keeps the whole loop as one task.
343 //
344 // This is the answer a build without rayon gets, since it reports a single worker.
345 if threads <= 1 {
346 return len.max(1);
347 }
348
349 // Phase 2: price the whole loop as if it ran on one core.
350 let per_item = item_picos(item_bytes);
351 let total = per_item.saturating_mul(len as u64);
352
353 // Phase 3: a cheap loop is floored at its own length.
354 //
355 // rayon cannot cut below the floor, so the loop stays a single task.
356 // An empty loop still reports one, since a floor of zero would mean nothing.
357 let min_parallel = budget
358 .min_parallel_picos
359 .unwrap_or_else(|| MIN_PARALLEL_PICOS_PER_WORKER.saturating_mul(threads as u64));
360 if total < min_parallel {
361 return len.max(1);
362 }
363
364 // Phase 4: one task per worker is the coarsest split worth allowing.
365 //
366 // The cap then lowers the floor further, so a stolen task can still be cut cache-sized.
367 let per_worker = total / threads as u64;
368 let task_picos = if per_worker < budget.max_task_picos {
369 per_worker
370 } else {
371 budget.max_task_picos
372 };
373
374 // Phase 5: convert a task's time budget back into a count of items.
375 //
376 // An item costlier than a whole task cannot be subdivided, so it yields one per task.
377 // The count cannot exceed the length, because a task holds at most the whole loop.
378 ((task_picos / per_item) as usize).max(1).min(len.max(1))
379}
380
381/// Minimum items one task must hold, for a loop of the given shape.
382///
383/// Use this where a call site picks its own chunk length, and pass the result as that length.
384///
385/// # Arguments
386///
387/// * `len` - total items the loop covers
388/// * `item_bytes` - bytes one item moves: every element read, plus every element written
389///
390/// # Returns
391///
392/// A count in `1..=max(len, 1)`, never zero, so it is always a usable chunk length.
393/// Equal to the total when the loop should not split at all.
394///
395/// Zeroing both budget overrides pins the answer at one, which a chunked caller must expect.
396#[inline]
397pub fn min_task_len(len: usize, item_bytes: usize) -> usize {
398 // The worker count both scales the gate and turns a time budget into a task count,
399 // so it is read here rather than baked into the budgets.
400 min_task_len_with(budget(), current_num_threads(), len, item_bytes)
401}
402
403/// Whether a loop of the given shape is worth splitting at all.
404///
405/// The iterator adapters express this same policy without a branch, so prefer them.
406/// Reach for this only where the two paths differ in more than how the work is divided.
407///
408/// An in-place pass that rewrites its buffer as it walks it is the usual case.
409///
410/// The answer moves with the pool, so both arms run somewhere.
411/// They must agree on every input, bit for bit, and a test has to pin that directly.
412/// Growing a fixture until the host happens to split checks one arm on one machine.
413///
414/// # Arguments
415///
416/// * `len` - total items the loop covers
417/// * `item_bytes` - bytes one item moves: every element read, plus every element written
418#[inline]
419pub fn should_split(len: usize, item_bytes: usize) -> bool {
420 // A floor that reaches the whole length is exactly the answer "do not split".
421 min_task_len(len, item_bytes) < len
422}
423
424/// A body an indexed iterator can run on one core, over its items in order.
425///
426/// A parallel iterator hands its items to a consumer, never to a plain `Iterator`.
427/// Naming the body separately is what lets a floored loop reach the serial path
428/// without the call site writing that body a second time.
429trait SerialOp<T> {
430 /// What running the body over every item produces.
431 type Output;
432
433 /// Runs the body over every item, in order.
434 fn run<I: ExactSizeIterator<Item = T>>(self, items: I) -> Self::Output;
435}
436
437/// Runs a body for its effect on every item.
438struct ForEachSerially<F>(F);
439
440impl<T, F: Fn(T)> SerialOp<T> for ForEachSerially<F> {
441 type Output = ();
442
443 #[inline]
444 fn run<I: ExactSizeIterator<Item = T>>(self, items: I) {
445 items.for_each(self.0);
446 }
447}
448
449/// Maps every item and gathers the results into a vector.
450struct MapCollectSerially<F>(F);
451
452impl<T, B, F: Fn(T) -> B> SerialOp<T> for MapCollectSerially<F> {
453 type Output = Vec<B>;
454
455 #[inline]
456 fn run<I: ExactSizeIterator<Item = T>>(self, items: I) -> Vec<B> {
457 items.map(self.0).collect()
458 }
459}
460
461/// Runs `op` over the whole loop on one core.
462///
463/// An indexed iterator lends out the leaf iterator rayon would walk inside a single task.
464/// Taking it directly skips rayon's bridge, its consumer, and its collect folder.
465#[cfg(feature = "parallel")]
466#[inline]
467fn drive_serially<I, Op>(iter: I, op: Op) -> Op::Output
468where
469 I: IndexedParallelIterator,
470 Op: SerialOp<I::Item>,
471{
472 iter.with_producer(SerialCallback(op))
473}
474
475/// Carries `op` into the one place a producer's concrete type is nameable.
476#[cfg(feature = "parallel")]
477struct SerialCallback<Op>(Op);
478
479#[cfg(feature = "parallel")]
480impl<T: Send, Op: SerialOp<T>> ProducerCallback<T> for SerialCallback<Op> {
481 type Output = Op::Output;
482
483 #[inline]
484 fn callback<P: Producer<Item = T>>(self, producer: P) -> Self::Output {
485 self.0.run(producer.into_iter())
486 }
487}
488
489/// Runs `op` over the whole loop, which is all a serial build ever does.
490#[cfg(not(feature = "parallel"))]
491#[inline]
492fn drive_serially<I, Op>(iter: I, op: Op) -> Op::Output
493where
494 I: IndexedParallelIterator,
495 Op: SerialOp<I::Item>,
496{
497 op.run(iter)
498}
499
500/// Task-size adapters for parallel iterators.
501///
502/// An adapter reads the loop's own length and floors the split from the cost model.
503/// A call site then states what one item costs, instead of hardcoding a task count.
504pub trait TaskSizeExt: IndexedParallelIterator {
505 /// Floors the split for a loop whose item is a single element of the named type.
506 ///
507 /// A body moving several elements per item charges for all of them.
508 ///
509 /// Such a body asks through the byte-counted adapter instead.
510 #[inline]
511 fn with_min_task<T>(self) -> impl IndexedParallelIterator<Item = Self::Item>
512 where
513 Self: Sized,
514 {
515 // One element per item is just the byte-counted form at that element's width.
516 self.with_min_task_bytes(size_of::<T>())
517 }
518
519 /// Floors the split for a loop whose item moves a given number of bytes.
520 ///
521 /// The count sums every element one item reads and every element it writes.
522 ///
523 /// Use this when the width is known only at run time, such as a chunk of a matrix row.
524 #[inline]
525 fn with_min_task_bytes(
526 self,
527 item_bytes: usize,
528 ) -> impl IndexedParallelIterator<Item = Self::Item>
529 where
530 Self: Sized,
531 {
532 // An indexed iterator already knows how long it is, so the call site never repeats it.
533 let min_len = min_task_len(self.len(), item_bytes);
534
535 // A serial build has no split to constrain, and drops the floor on the floor.
536 self.with_min_len(min_len)
537 }
538
539 /// Runs a body over every item, with the split floored by the cost model.
540 ///
541 /// The count sums every element one item reads and every element it writes.
542 ///
543 /// A loop the floor keeps whole is driven as a plain serial iterator, not as one task.
544 /// rayon reaches a body through its own bridge, which costs about 1 ns per item.
545 #[inline]
546 fn for_each_min_task_bytes<F>(self, item_bytes: usize, op: F)
547 where
548 Self: Sized,
549 F: Fn(Self::Item) + Send + Sync,
550 {
551 let len = self.len();
552 let min_len = min_task_len(len, item_bytes);
553
554 // A floor reaching the whole length is the answer "do not split".
555 if min_len >= len {
556 return drive_serially(self, ForEachSerially(op));
557 }
558 self.with_min_len(min_len).for_each(op);
559 }
560
561 /// Maps every item and collects the results, with the split floored by the cost model.
562 ///
563 /// The count sums every element one item reads and every element it writes.
564 ///
565 /// A loop the floor keeps whole is collected serially, not through rayon's collect folder.
566 /// That folder costs about 1 ns per item more than `Extend`, which a cheap body feels.
567 #[inline]
568 fn map_collect_min_task_bytes<B, F>(self, item_bytes: usize, map_op: F) -> Vec<B>
569 where
570 Self: Sized,
571 B: Send,
572 F: Fn(Self::Item) -> B + Send + Sync,
573 {
574 let len = self.len();
575 let min_len = min_task_len(len, item_bytes);
576
577 // A floor reaching the whole length is the answer "do not split".
578 if min_len >= len {
579 return drive_serially(self, MapCollectSerially(map_op));
580 }
581 self.with_min_len(min_len).map(map_op).collect()
582 }
583}
584
585impl<I: IndexedParallelIterator> TaskSizeExt for I {}
586
587#[cfg(test)]
588mod tests {
589 use alloc::vec;
590
591 use super::*;
592
593 /// Fixed budgets keep arithmetic fixtures independent of the host platform.
594 const B: Budget = Budget {
595 min_parallel_picos: Some(20_000_000),
596 max_task_picos: MAX_TASK_PICOS,
597 };
598
599 /// Worker count used by the arithmetic fixtures.
600 const T: usize = 32;
601
602 #[test]
603 fn short_loops_stay_on_one_core() {
604 // A loop worth less than the parallel budget is floored at its own length.
605 // rayon cannot split below the floor, so the loop runs as a single task.
606 //
607 // 1000 items * 4 B * 100 ps = 0.4 us < 20 us
608 assert_eq!(min_task_len_with(B, T, 1000, 4), 1000);
609
610 // The same loop over a wider item crosses the budget and does split.
611 //
612 // 1000 items * 256 B * 100 ps = 25.6 us >= 20 us
613 assert!(min_task_len_with(B, T, 1000, 256) < 1000);
614
615 // An empty loop still reports a usable floor rather than zero.
616 assert_eq!(min_task_len_with(B, T, 0, 4), 1);
617 }
618
619 #[test]
620 fn the_gate_grows_with_the_pool() {
621 // Invariant: waking a wider pool costs more, so a wider pool demands a longer loop.
622 // Price the fixture just above the platform's 32-worker gate.
623 let budget = Budget::DEFAULT;
624 let low_workers = 32;
625 let high_workers = 64;
626 let len = MIN_PARALLEL_PICOS_PER_WORKER
627 .saturating_mul(low_workers as u64)
628 .div_ceil(item_picos(4)) as usize;
629 assert!(min_task_len_with(budget, low_workers, len, 4) < len);
630 assert_eq!(min_task_len_with(budget, high_workers, len, 4), len);
631 }
632
633 #[test]
634 fn an_override_fixes_the_gate_against_the_pool() {
635 // An override states the gate outright, which is what makes an A/B run reproducible.
636 //
637 // Fixture state: gate pinned at 1 us, against a loop worth 20 us.
638 // -> both pool sizes split, where the scaling gate would have disagreed.
639 let fixed = Budget {
640 min_parallel_picos: Some(1_000_000),
641 max_task_picos: MAX_TASK_PICOS,
642 };
643 let len = 50_001;
644 assert!(min_task_len_with(fixed, 32, len, 4) < len);
645 assert!(min_task_len_with(fixed, 64, len, 4) < len);
646 }
647
648 #[test]
649 fn floor_scales_inversely_with_item_width() {
650 // Invariant: a task holds a fixed amount of work, whatever the item width.
651 // So the item count must fall by exactly the factor the item widens.
652 //
653 // 1 << 22 items of 4 B moves 16 MiB, far past the parallel budget
654 let len = 1 << 22;
655 let narrow = min_task_len_with(B, T, len, 4);
656 let wide = min_task_len_with(B, T, len, 16);
657 assert_eq!(narrow, 4 * wide);
658
659 // The cap binds at this length, so a task holds exactly the task budget.
660 //
661 // 2 us / (4 B * 100 ps) = 5000 items
662 // 2 us / (16 B * 100 ps) = 1250 items
663 assert_eq!((narrow, wide), (5000, 1250));
664 }
665
666 #[test]
667 fn the_cap_lowers_the_floor_below_one_task_per_worker() {
668 // Invariant: this is a statement about the floor, not about how many tasks are made.
669 //
670 // A floor only forbids a split, so it is a ceiling on how finely stealing may cut.
671 //
672 // Fixture state: a loop long enough for the cap to bind.
673 // -> the floor drops below one worker's share, leaving room to cut further
674 let len = 1 << 22;
675 let floor = min_task_len_with(B, T, len, 4);
676 assert!(len / floor > T);
677
678 // Just past the parallel budget the cap does not bind.
679 //
680 // One worker's share is then the smallest piece the floor allows.
681 //
682 // 50001 items worth 20.0 us, over 32 workers = 625 ns per task
683 let len = 50_001;
684 let floor = min_task_len_with(B, T, len, 4);
685 assert_eq!(floor, 1562);
686 assert!(len / floor <= T);
687 }
688
689 #[test]
690 fn costly_items_fall_back_to_one_per_task() {
691 // An item worth more than a whole task cannot be subdivided further.
692 //
693 // One item per task is also where an unfloored rayon loop already sits.
694 // So undercharging an arithmetic-bound body can never cut it more finely than that.
695 //
696 // 2 us budget / (64 KiB * 100 ps) = 0.31 -> floored to 1
697 assert_eq!(min_task_len_with(B, T, 1 << 20, 64 << 10), 1);
698 }
699
700 #[test]
701 fn zero_sized_items_are_charged_one_byte() {
702 // Charging nothing would make the total zero and leave every such loop serial.
703 // One byte per item keeps a long loop over zero-sized items splittable.
704 assert_eq!(item_picos(0), PICOS_PER_BYTE);
705 assert!(min_task_len_with(B, T, 1 << 30, 0) < 1 << 30);
706 }
707
708 #[test]
709 fn a_pool_of_one_never_splits() {
710 // Invariant: a single worker has nobody to hand work to, so there is no split.
711 // The whole loop is one task however expensive it is.
712 assert_eq!(min_task_len_with(B, 1, 1 << 20, 4), 1 << 20);
713 assert_eq!(min_task_len_with(B, 1, 1 << 20, 64 << 10), 1 << 20);
714
715 // A pool reported as empty is treated the same way rather than dividing by zero.
716 assert_eq!(min_task_len_with(B, 0, 1 << 20, 4), 1 << 20);
717 }
718
719 #[test]
720 fn floor_is_always_a_usable_chunk_length() {
721 // Invariant: every call site divides a length by this floor, or chunks by it.
722 //
723 // lower bound 1 -> par_chunks never sees zero and panics
724 // upper bound len -> chunk index arithmetic never runs off the end
725 //
726 // Swept exhaustively rather than sampled, since the grid is small and fixed.
727 let fixed = Budget {
728 min_parallel_picos: Some(1_000_000),
729 max_task_picos: MAX_TASK_PICOS,
730 };
731 for budget in [B, fixed] {
732 for threads in [0usize, 1, 2, 8, 32, 96, 192, 1 << 20] {
733 for len in [0usize, 1, 2, 3, 5, 17, 1023, 1 << 10, 1 << 20, usize::MAX] {
734 for item_bytes in [0usize, 1, 3, 4, 8, 16, 32, 4096, 1 << 20] {
735 let floor = min_task_len_with(budget, threads, len, item_bytes);
736 assert!(
737 (1..=len.max(1)).contains(&floor),
738 "floor {floor} out of range for \
739 threads={threads} len={len} item_bytes={item_bytes}"
740 );
741 }
742 }
743 }
744 }
745 }
746
747 #[test]
748 fn should_split_agrees_with_the_floor() {
749 // Invariant: the predicate is exactly the question the floor already answers.
750 //
751 // Read through the public entry points.
752 //
753 // So this holds on any host pool, in either feature configuration.
754 for len in [0usize, 1, 2, 1023, 1 << 10, 1 << 20] {
755 for item_bytes in [1usize, 4, 16, 256, 1 << 20] {
756 assert_eq!(
757 should_split(len, item_bytes),
758 min_task_len(len, item_bytes) < len
759 );
760 }
761 }
762 }
763
764 /// A build without rayon reports one worker, so no loop is ever worth splitting.
765 ///
766 /// A serial build keeps every call site on its in-place, short-circuiting arm.
767 ///
768 /// This pins that behaviour, which those call sites rely on.
769 #[cfg(not(feature = "parallel"))]
770 #[test]
771 fn a_serial_build_never_splits() {
772 // Fixture state: a loop far past any parallel budget, over a 12-byte item.
773 //
774 // 1 << 19 items * 12 B * 100 ps = 629 us, which a 32-worker pool would split
775 assert_eq!(current_num_threads(), 1);
776 assert_eq!(min_task_len(1 << 19, 12), 1 << 19);
777 assert!(!should_split(1 << 19, 12));
778 }
779
780 #[cfg(feature = "parallel")]
781 #[test]
782 fn the_public_default_exposes_the_compiled_in_budgets() {
783 // An absent fixed gate retains the platform-specific per-worker calibration.
784 let budget = TaskBudget::default();
785 assert_eq!(budget.min_parallel_ns, None);
786 // The public nanosecond value exactly represents the internal picosecond cap.
787 //
788 // Multiplying back up pins the round trip, which a truncating division would hide:
789 //
790 // 2_000_500 ps / 1_000 = 2_000 ns -> 2_000 ns * 1_000 = 2_000_000 ps != cap
791 assert_eq!(budget.max_task_ns * 1_000, MAX_TASK_PICOS);
792 }
793
794 #[cfg(feature = "parallel")]
795 #[test]
796 fn the_budgets_can_only_be_fixed_once() {
797 // Invariant: the budgets hold for the whole process, so a later change is refused.
798 //
799 // Every other test in this binary reads them, and the order tests run in is not
800 // fixed, so the first call here may already be too late.
801 // Whichever way that lands, the call after it must fail.
802 let budget = TaskBudget {
803 min_parallel_ns: Some(1_000),
804 max_task_ns: 2_000,
805 };
806 let _ = set_task_budget(budget);
807 assert_eq!(set_task_budget(budget), Err(TaskBudgetLocked));
808 }
809
810 #[test]
811 fn adapters_preserve_iteration() {
812 // The adapters constrain only how work is divided, never what it covers.
813 //
814 // Fixture state: 1000 items of 8 B, worth 0.8 us, well under the default gate.
815 // -> the loop runs as a single task, and every item is still visited once.
816 let data: Vec<u64> = (0..1000).collect();
817 let expected = 1000 * 999 / 2;
818
819 assert_eq!(
820 data.par_iter().with_min_task::<u64>().sum::<u64>(),
821 expected
822 );
823 assert_eq!(
824 data.par_iter().with_min_task_bytes(8).sum::<u64>(),
825 expected
826 );
827
828 // A width that does cross the budget splits the same loop into many tasks,
829 // and the answer is unchanged.
830 assert_eq!(
831 data.par_iter().with_min_task_bytes(1 << 20).sum::<u64>(),
832 expected
833 );
834 }
835
836 #[test]
837 fn the_terminals_match_the_serial_loop_on_both_arms() {
838 // The two terminals take different paths either side of the gate.
839 //
840 // The serial one reads the iterator's own leaf; the split one goes through rayon.
841 // Both must visit every item once, in order, so a reference loop pins them.
842 //
843 // Fixture state: three shapes, chosen to reach every path.
844 //
845 // empty : nothing to visit, and a floor that must not be zero
846 // 64 x 4 B : priced at 25 ns, under the gate on any pool -> serial leaf
847 // 16384 x 4 KiB : priced at 6.7 ms, over the gate on any pool -> split
848 for (len, item_bytes) in [(0usize, 4usize), (64, 4), (1 << 14, 1 << 12)] {
849 let src: Vec<u64> = (0..len as u64).collect();
850 let expected: Vec<u64> = src.iter().map(|&x| x * 3 + 1).collect();
851
852 // A mapped collect keeps the order of the source, whichever arm runs.
853 assert_eq!(
854 src.par_iter()
855 .map_collect_min_task_bytes(item_bytes, |&x| x * 3 + 1),
856 expected
857 );
858
859 // A for-each writes through, so the same order shows up in the output buffer.
860 let mut out = vec![0u64; len];
861 out.par_iter_mut()
862 .zip(src.par_iter())
863 .for_each_min_task_bytes(item_bytes, |(out, &x)| *out = x * 3 + 1);
864 assert_eq!(out, expected);
865 }
866 }
867}