moirai_parallel/ops/unit_tasks/
layout.rs

1//! Unit-task planning: the bytes-per-task arithmetic shared by the
2//! [`for_each_unit_task_*`](self) operators (ADR 0059).
3
4use crate::policy::ExecutionPolicy;
5
6/// Bytes of work one scheduled task carries.
7///
8/// One length-64 lane per task made apollo's 64Â³ transform pass 3.3x slower
9/// than serial: moirai's per-task dispatch measures about 180 ns against a
10/// 55â€“130 ns lane. Tasks of 64 KiB made the same pass faster than serial and
11/// than quarter-megabyte tasks (apollo `dimension_3d::pass_attribution`,
12/// 2026-09-09), and leto-ops' batched transpose settled on the same width. It
13/// stays inside one core's L2.
14pub const UNIT_TASK_BYTES: usize = 64 * 1024;
15
16/// Units one task carries when each unit moves `unit_bytes`, never fewer than
17/// one: a unit at or above [`UNIT_TASK_BYTES`] is a task on its own.
18#[must_use]
19pub const fn units_per_task(unit_bytes: usize) -> usize {
20    let per_task = UNIT_TASK_BYTES / if unit_bytes == 0 { 1 } else { unit_bytes };
21    if per_task == 0 { 1 } else { per_task }
22}
23
24/// Task layout of a unit-task pass over `len` elements.
25#[derive(Clone, Copy)]
26pub(super) struct UnitTaskPlan {
27    /// Units one task carries.
28    pub(super) per_task: usize,
29    /// Elements one task carries: `per_task` whole units.
30    pub(super) task_len: usize,
31    /// Tasks the pass splits into; the last may be shorter.
32    pub(super) tasks: usize,
33    /// Whether the policy spreads the tasks over workers.
34    pub(super) parallel: bool,
35}
36
37/// Rejects data that does not divide into whole `unit_len`-element units.
38#[track_caller]
39pub(super) fn assert_whole_units(len: usize, unit_len: usize) {
40    assert!(
41        unit_len > 0 && len.is_multiple_of(unit_len),
42        "unit tasks need whole units: data length {len} is not a multiple of unit length {unit_len}",
43    );
44}
45
46/// The task layout and policy decision for `len` elements of whole
47/// `unit_len`-element units that each move `unit_bytes`, or `None` when there
48/// is no unit to run. `P::parallelize_work` is consulted only when the pass
49/// spans more than one task.
50pub(super) fn plan_unit_tasks<P: ExecutionPolicy>(
51    len: usize,
52    unit_len: usize,
53    unit_bytes: usize,
54) -> Option<UnitTaskPlan> {
55    let units = len / unit_len;
56    if units == 0 {
57        return None;
58    }
59    let per_task = units_per_task(unit_bytes);
60    let tasks = units.div_ceil(per_task);
61    Some(UnitTaskPlan {
62        per_task,
63        task_len: per_task * unit_len,
64        tasks,
65        parallel: tasks > 1 && P::parallelize_work(len, tasks, units.saturating_mul(unit_bytes)),
66    })
67}
moirai_parallel/ops/unit_tasks/layout.rs

moirai_parallel/ops/unit_tasks/
layout.rs
