From d8b4c5cc1f6f6119bde2ec044425bdaf2c31fb0f Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Thu, 29 Oct 2026 13:14:09 +0900 Subject: [PATCH] placement: opt-in participant pinning, PAUSE-only reserved spin, oversubscribed park (#9) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the decided design for tensor4all/threadteam-rs#9 on top of the broccoli A/B/C measurement in tensor4all/faer#107 (pinning removes the 3x process mode: 0 of 20 slow with the participants pinned, 11 of 20 unpinned, 7 of 20 with one spare CPU). A. Opt-in placement. `PoolBuilder::pin_participants(bool)` (one private `PIN_PARTICIPANTS_DEFAULT: bool = false` switch), `Pool::pin_participants()` getter, `src/pin.rs`. At the start of a root team scope (`Exec::team_with` on the acquisition path and on the carried path of a root token minted by a `Team`) the leader reads its own mask S and runs one pin section through the section machinery (it replaces `acknowledge`): member k (participant k >= 1) pins itself to the k-th set CPU of S, wrapping modulo |S|. **The leader is never pinned** and its mask is never changed, not even for the scope: a vendor library builds its thread team from the thread that calls it, and a one-CPU caller silently confines the whole provider team (tenferro-rs PERFORMANCE_TIPS.md: 3-20x slower dgemm/LU), so S[0] is left to the leader without enforcing it; a pinned member must not build such a team directly inside a share (use `Exec::external`, which restores its mask first, or keep pinning off). Only the thread being moved changes its affinity, with `sched_setaffinity(0, ...)`. A member saves its previous mask in its own slot (UnsafeCell + a thread-owned flag, SAFETY comment) the first time it is pinned and restores it itself in `member_loop` before it leaves `RESERVED`, before it serves Rayon while lent, and (`sched::unpin_worker`) when it is itself the thread that runs foreign work, e.g. the member that leads a `fork` child and calls `external`. Best effort everywhere; Linux only (accepted and inert elsewhere); `libc` is now also a Linux dependency (no libc type is public); stack `cpu_set_t` only, so the path allocates nothing. The pin section is tagged by a bool on `Section`; its closure captures the mask. B. Reserved spin without yield: `spin_wait` (one PAUSE per poll) replaces the shared `backoff` in the reserved `gen` poll and in `Section::join`/ `Section::barrier` when the scope is not oversubscribed. `backoff` stays for acquisition, the cold claim, `IDLE` and the leaving states. C. Oversubscription: a root scope whose participants exceed |S| (Linux affinity count, else `available_parallelism`) sets the members' park delay to zero for its reservation and resets `polls`/`spins` per idle episode (64 polls before the park check); `join`/`barrier` spin ~100 times then yield per poll. Root pressure and the park delay are decided on the thread that enters the scope — at `Exec::team_with` on both paths and, for a `Team`, also at `Team::exec`, so a token minted on one CPU set and entered on another is paced and placed by the entering thread's set. The flag is carried through `TeamAuth`, `TeamExec`, reborrow/limit/split/fork children and every `Section`. D. `src/park.rs`, the handshake and publication are untouched; the loom models pass unchanged. Tests: `tests/pinning.rs` (member placement on distinct CPUs, wrap-around on a 2-CPU caller mask, the leader's mask identical before, inside and after the scope for an acquisition and for a persistent-team root token and after a panic in a share, persistent members across two scopes, pinning off, restoration while lent by `external`, and a member-led `fork` child calling `external`), `tests/oversubscribed.rs` (a 2-CPU mask and a 1-CPU scope entry, both with a policy park delay far above the hang guard, so parking at all proves the scope's zero delay; nested split included), `src/tests/placement.rs` (flag propagation, entry-time pressure), and pinning-on allocation windows in `tests/alloc_free.rs` and `tests/team_alloc_free.rs`. `tests/idle_clock.rs` now checks the order of each member's own trace events instead of a wall-clock premise ("sections every 5 ms" breaks on a slow runner once members no longer yield): never a park less than the delay after the member's own last share, wake-up or start; every member parks once no section arrives; a following section wakes every parked member. Docs: `README.md`, `docs/design/placement.md` (new "Participant placement (opt-in)" and "How a reserved member and a joining leader wait" sections). Co-Authored-By: Claude Opus 5.0 --- Cargo.toml | 5 + README.md | 3 +- docs/design/placement.md | 147 +++++++++++- src/awake.rs | 169 +++++++++++++- src/exec.rs | 111 ++++++++- src/lib.rs | 1 + src/pin.rs | 259 +++++++++++++++++++++ src/pool.rs | 118 +++++++++- src/reserve.rs | 15 +- src/sched.rs | 18 ++ src/team.rs | 120 +++++++--- src/tests/mod.rs | 2 + src/tests/placement.rs | 180 +++++++++++++++ tests/alloc_free.rs | 24 ++ tests/idle_clock.rs | 186 +++++++++++---- tests/oversubscribed.rs | 212 +++++++++++++++++ tests/pinning.rs | 484 +++++++++++++++++++++++++++++++++++++++ tests/team_alloc_free.rs | 37 +++ 18 files changed, 1989 insertions(+), 102 deletions(-) create mode 100644 src/pin.rs create mode 100644 src/tests/placement.rs create mode 100644 tests/oversubscribed.rs create mode 100644 tests/pinning.rs diff --git a/Cargo.toml b/Cargo.toml index 051ed78..aeb95fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,9 +35,14 @@ bench = false [target.'cfg(loom)'.dependencies] loom = "0.7" +# Worker QoS (`pthread_set_qos_class_self_np`). [target.'cfg(target_vendor = "apple")'.dependencies] libc = "0.2" +# Participant placement (`sched_getaffinity` / `sched_setaffinity`). +[target.'cfg(target_os = "linux")'.dependencies] +libc = "0.2" + [dev-dependencies] libc = "0.2" diff --git a/README.md b/README.md index d0e7dfe..efd0948 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,8 @@ the token to kernels by value. Kernels never store a token and never mint one. | `Exec::team(f)` / `team_with(over, f)` | reserve `budget - 1` awake pool members for a scope and run barrier-bearing sections through `TeamExec` (`run`, `run_with`, `split`, `fork`, `reborrow`), or `Err(Refused::Unavailable)`; see the design record | | `Pool::into_team(width)` -> `Team` (owns the pool; `IntoTeamError` gives it back), `Team::exec(budget)` / `try_exec`, `Team::pool`, `Team::into_pool` | persistent reservation for an engine that owns its pool (`Pool` is not `Clone`); minted tokens are ordinary `Exec`s whose `fork`, `for_ranges`, `broadcast`, `team` dispatch on the reserved members; one leader at a time (`exec` blocks, `try_exec` -> `Refused::Busy`); idle members park while reserved | | `PoolBuilder::qos(Qos)`, `raise_current_thread_qos(Qos)` | worker QoS on macOS (default `UserInitiated`); raise the leading thread without ever lowering it | -| `PoolBuilder::idle_spin`, `idle_runs_rayon`, `acquire_timeout`; `Pool::release_idle` | waiting policy of idle team members | +| `PoolBuilder::idle_spin`, `idle_runs_rayon`, `acquire_timeout`; `Pool::release_idle` | waiting policy of idle team members; a reserved member spins with one PAUSE per poll (never `sched_yield`), parks after the delay, and parks after a short spin in a root scope that has more participants than the caller's CPUs | +| `PoolBuilder::pin_participants(bool)`, `Pool::pin_participants()` | opt-in placement of the **members** of a root team scope (Linux; off by default, which is consistent with the rule it protects): member `k` (participant `k >= 1`) pins itself onto the `k`-th CPU of the calling thread's affinity set, `S[0]` is left to the leader, and a member restores its own mask before it returns to Rayon or serves foreign work. **The leader is never pinned**, because a vendor library creates its thread team from the calling thread and a one-CPU caller would confine the whole provider team; a member that calls such a library directly inside a share does confine it, so use `Exec::external` (which restores the member's mask first) or leave pinning off. For tensor4all/faer#107; give each team its own CPU set (`taskset -c`), there is no CPU allocator | | `DispatchCosts`, `WidthCosts`, `PoolBuilder::dispatch_costs`, `Exec::dispatch_cost(width)`, `TeamExec::dispatch_cost(width)` | cost of one parallel section per kind of token (team, in pool, outside), steady state; one row per `HostClass` (`HostClass::detect()`, keyed by kernel dispatch: `x86_64-avx512`, `x86_64-avx2`, `x86_64-scalar`, `aarch64-apple`, `aarch64`, `generic`), `DispatchCosts::for_class(class)`; measured rows for `x86_64-avx512` (team anchors at widths 2, 4, 6 and 8) and `aarch64-apple`, a measured Zen2 team row for `x86_64-avx2` (anchors at ten widths from 2 to 64; its plain kinds stay generic), every other class takes the conservative `generic` row (3 / 15 / 40 us at the anchors 2, 4, 6); a row holds measured anchors (`WidthCosts::new`, `WidthCosts::flat`, or `WidthCosts::from_anchors` for measured anchors up to 64 and beyond), interpolates between them and grows linearly above the largest one, so a width above the anchors is never charged the cost of the largest anchor; `DispatchCosts::host_default()` is the detected class's row; `Duration::MAX` when serial; override per pool | | `pays_at(est, sections, cost)`, `Exec::pays(est, sections, width)`, `Exec::width_for(est, sections)`, `width_for_by`, `PAYS` | the one rule: parallel iff `est >= PAYS * sections * D(width)`; `width_for` minimises `sections * D(k) + est / k` | | `Exec::external(f)` | enter the pool once and hand the budget as a plain width to an outside library (faer, a vendor BLAS); how it divides the width is that library's concern | diff --git a/docs/design/placement.md b/docs/design/placement.md index 8194054..7d489ee 100644 --- a/docs/design/placement.md +++ b/docs/design/placement.md @@ -70,6 +70,7 @@ impl PoolBuilder { pub fn thread_name(self, f: impl Fn(usize) -> String + 'static) -> Self; pub fn stack_size(self, bytes: usize) -> Self; // default sized for LU recursion pub fn on_start(self, f: impl Fn(usize) + Send + Sync + 'static) -> Self; // e.g. pinning + pub fn pin_participants(self, on: bool) -> Self; // opt-in, Linux; off by default pub fn build(self) -> Result; } @@ -324,11 +325,12 @@ A section is a stack object (`Section`): a type-erased pointer to the closure, the participant count, a padded completion counter, a sense-reversing barrier, a cancellation flag and an atomic cell for the first panic payload. The leader stores the section address and the participant index into each -member slot and bumps its `gen` counter with release ordering; members spin on -`gen` and `state` (pure spin, then `yield_now` between polls) and run their -share and decrement the completion counter. After the reservation's park delay -(the effective `idle_spin`) without work, a member parks (see below). The leader runs -share 0, then waits for the counter. Nothing is allocated and nobody is woken +member slot and bumps its `gen` counter with release ordering; members poll +`gen` and `state` with one PAUSE per poll (`spin_loop`, no `sched_yield`) and +run their share and decrement the completion counter. After the reservation's +park delay (the effective `idle_spin`, or zero in an oversubscribed root scope, +see below) without work, a member parks (see below). The leader runs +share 0, then waits for the counter, polling the same way. Nothing is allocated and nobody is woken (`tests/team_alloc_free.rs`; acquisition allocates a small `Vec` and, for members not already awake, the claim jobs). @@ -450,6 +452,54 @@ shim (`std`, or `loom` under `--cfg loom`), and checked three ways: No allocation happens on the dispatch path. Dropping the team stores `IDLE` with a zero idle spin and broadcasts, so the members return to Rayon at once. +### How a reserved member and a joining leader wait + +Two waits exist, both allocation-free and both unchanged in the handshake: + +* **Reserved members polling `gen`, and a leader in `Section::join` or + `Section::barrier`, in a scope that is not oversubscribed**: one + `std::hint::spin_loop()` per poll (PAUSE on x86, YIELD/ISB on aarch64), never + a `yield_now`. A thread that only issues PAUSE stays runnable but does not + re-enter the scheduler on every poll, so a pair the scheduler stacked on one + CPU can be separated instead of preempting each other in a yield loop, and a + share does not pay a syscall per poll. This is libgomp's shape. A reserved + member still parks after its park delay, which is what ends the spin; nothing + else wakes it. +* **Acquisition, the cold claim, and the `IDLE` / leaving states**: `backoff` + as before (a bounded pure spin, then one `yield_now` per poll). Those threads + are not reserved yet or are on their way out, so they must cooperate to make + progress on a loaded host. + +**Oversubscription.** A root team scope is oversubscribed when its participants +number more than the CPUs in the **entering thread's** affinity mask (on +non-Linux targets: `std::thread::available_parallelism()`). It is decided at +root-scope entry on that thread, whether or not participant pinning is on: at +`Exec::team_with` (acquisition path and the carried path of a root token minted +by a `Team`), and for a `Team` also when its root token is minted +(`Team::exec`/`try_exec`), which is the pressure a token used without a team +scope reports. Because the scope entry repeats it, a token minted on one CPU set +and used on another is paced and placed by the set of the thread that enters the +scope, never by the one it was minted on. The result is carried into +`TeamExec`, team-carrying `Exec` tokens, `reborrow`, `limit`, `split` and `fork` +children and every `Section`, so a nested section of an oversubscribed root +knows it too; it is never recomputed from a child's size. The members' park +delay is written to their slots at the same entry. In such a scope: + +* the park delay of the reserved members is **zero**, so a member that finds no + section spins one short episode (64 polls) and parks instead of holding a CPU + for the whole `idle_spin`. `polls` and `spins` are reset at the start of each + idle episode, so the episode is counted afresh; +* `Section::join` and `Section::barrier` spin about 100 times and then + `yield_now` per poll, the shape libomp takes when it detects more threads than + CPUs, so the waiting participants do not hold every CPU while their peers are + queued behind them. + +The flag is per scope and per root; overlapping teams or pools are the caller's +business (see the placement warning below). This is the behaviour libgomp +throttles to about 100 spins when oversubscribed; measured on broccoli, parking +alone did not remove the tensor4all/faer#107 slow mode, so the zero park delay is +about the CPU a member burns, not about the mode. + **Quality of service.** `PoolBuilder::qos(Qos)` sets the class of each worker when it starts (`pthread_set_qos_class_self_np` on Apple targets, nothing elsewhere; no platform type is public). The default is `UserInitiated` on @@ -462,10 +512,89 @@ measurable difference to the wake latency on this host. Dropping the team stores `IDLE` with a zero idle spin and wakes them, so the members return to Rayon at once. +### Participant placement (opt-in) + +`PoolBuilder::pin_participants(bool)` (default off; `Pool::pin_participants()` +reads it back) makes a root team scope place its members. It exists because +of the measurement in tensor4all/faer#107: `qr_factor` f64 512 at 16 threads on +16 allowed CPUs, 20 processes on broccoli (EPYC 7702P, boost 0, governor +`performance`), 20 repetitions each. + +| arm | slow processes (>= 6 ms) | samples with two runnable team threads on one CPU | publish-to-share | +| --- | --- | --- | --- | +| A, width 16, unpinned | 11 / 20 | 92 % in the slow half | 225 us fast / 835 us slow | +| B, width 15 on the same 16 CPUs (one spare CPU) | 7 / 20 | 91 % in the slow half | 196 / 757 us | +| C, width 16, every participant pinned to CPUs 48-63 | **0 / 20** | **0 %** | 218 us | + +Arm C pinned the leader too; `threadteam` pins only the members and leaves the +leader on the caller's whole set (the rule below), which relaxes the placement of +participant 0 only. + +12.4 ms is the one-thread time of the call, so the slow mode is the members' +work not overlapping: they are alive, unparked and receiving the same shares, +but queued on one CPU. A spare CPU does not stop the scheduler from stacking the +pair; deterministic placement does. That is why placement is the load-bearing +part of the change and parking is not, and why placement stays **opt-in**: all +four compared runtimes leave threads unpinned by default, and pinning takes CPUs +away from the rest of the process and fights the caller's `taskset`. + +How it works: + +* At the start of a **root team scope** — `Exec::team_with` on the acquisition + path, and on the carried path when the token is the root token minted by a + persistent `Team` (`Team::exec`/`try_exec`; not a child, a reborrow or a + sub-team) — the leader reads **its own** mask `S` with + `sched_getaffinity(0, ...)` and runs one pin section through the ordinary + section machinery (it also serves as the rendezvous, replacing + `acknowledge`). Member `k` (participant `k >= 1`) pins itself with + `sched_setaffinity(0, ...)` onto the `k`-th set CPU of `S`, wrapping modulo + `|S|` when the team is wider than the set. A team of one is left alone. +* **The leader is never pinned**, and its mask is never changed, not even for + the scope. `S[0]` is left to it without enforcing it. A vendor library (an + OpenMP or OpenBLAS runtime, a LAPACK provider) builds its thread team from the + thread that calls it, and a team built from a one-CPU thread silently confines + the whole provider team to that CPU (`tenferro-rs` `PERFORMANCE_TIPS.md`: 3-20x + slower `dgemm` and LU at 4 threads); engines call such libraries from the + coordinating thread, so that thread keeps its CPU set. Consequently a **pinned + member must not build a foreign thread team directly inside a share**: it + would confine that library to its one CPU. Route the call through + `Exec::external`, which restores the member's own mask first, or keep pinning + off — it is off by default. +* **Every affinity change is made by the thread itself.** A member + saves its previous mask in its own slot (an `UnsafeCell` plus a + flag that only the member thread of that slot touches) the first time it is + pinned, and restores it itself, in the member loop, before it leaves + `RESERVED` (so before it returns its thread to Rayon) and before it serves + Rayon work while lent by `Exec::external` — and it does the same when it is + the thread that runs work outside its team itself, for example the member + that leads a `fork` child and calls `external` (there the call runs inline on + that worker, so it restores its own mask on its own thread before entering); + it re-pins only at the next pin section. A persistent team therefore keeps the + mask of its first scope and re-pins at each scope relative to that scope's + leader. +* **Best effort.** A refused syscall (a kernel mask larger than `cpu_set_t`, a + cpuset that changed under us) leaves that thread as it was; an empty mask pins + nobody; a failed pin or restore never fails the scope. +* **Linux only.** Other targets accept the option and do nothing. No libc type + appears in the public API; the masks are stack `cpu_set_t` values and the + `k`-th set CPU is found by scanning one, so the path allocates nothing + (`tests/team_alloc_free.rs`). +* **Assumption: the caller's CPU set belongs to this team.** Placement maps the + participants onto the caller's mask, so two teams whose leaders see the same + mask stack on the same CPUs. Give each team (in practice: each process) its + own CPU set with `taskset -c`; there is no CPU allocator in this crate and none + is planned for this fix. Overlapping teams and pools of one process stay + caller-managed. +* Only the members of a team are placed. A token of a `Team` that never opens a + team scope (`exec.for_ranges(...)` directly) is not re-pinned; its members + keep the placement of the last root scope, or of `Pool::into_team`. + ### Waiting policy `PoolBuilder::idle_spin` (default 200 ms, may be zero) and -`PoolBuilder::idle_runs_rayon` (default `false`) set `WaitPolicy` defaults; +`PoolBuilder::idle_runs_rayon` (default `false`) set `WaitPolicy` defaults; an +oversubscribed root scope overrides the park delay with zero for its own +reservation only; `acquire_timeout` is separate. `WaitOverride` (partial policy) applies to one operation: `Exec::team_with` for a root reservation, `TeamExec::with_wait` for a scope. Children inherit the effective policy and may override it for their own @@ -549,8 +678,10 @@ on a pool worker grants no reservation. Reproducible reductions are left to numerical kernels (fixed logical chunks, chunk-owned partials, fixed merge order): the crate gives no stealing within a -team, so a section's shares are exactly the participants. Cache-domain layout and -pinning are not handled. +team, so a section's shares are exactly the participants. Cache-domain layout is +not handled, and placement is opt-in and per team (see below): there is no CPU +allocator, so two teams whose leaders see the same CPU set would place onto the +same CPUs. ## Dispatch cost and the parallel rule diff --git a/src/awake.rs b/src/awake.rs index 1f56498..b6f2faf 100644 --- a/src/awake.rs +++ b/src/awake.rs @@ -56,7 +56,13 @@ const RESERVED: usize = 3; pub(crate) struct Padded(pub(crate) T); /// Pure spin for a while, then yield the CPU between polls. +/// +/// Used where progress needs cooperation with other threads on the same CPUs: +/// `acquire`, the cold claim, and the loops that are not reserved dispatch +/// polls (`IDLE`, and the states a loop is leaving). Reserved members and a +/// leader waiting on a section poll with [`spin_wait`] instead. #[inline] +#[cfg_attr(not(feature = "parallel"), allow(dead_code))] pub(crate) fn backoff(spins: &mut u32) { if *spins < (1 << 14) { *spins += 1; @@ -66,6 +72,43 @@ pub(crate) fn backoff(spins: &mut u32) { } } +/// One PAUSE per poll (YIELD on aarch64), never a yield. +/// +/// The wait of a reserved member polling `gen`, and of a leader in +/// `Section::join` / `Section::barrier`, in a scope that is not oversubscribed: +/// a thread that only spins with PAUSE stays runnable but does not repeatedly +/// re-enter the scheduler, so a stacked pair can be separated and a share does +/// not pay for a `sched_yield` on every poll. A parked member is woken by the +/// leader through [`Parker`]. +#[inline] +pub(crate) fn spin_wait() { + std::hint::spin_loop(); +} + +/// How many polls a joined section spins before it yields, in an oversubscribed +/// scope (more participants than CPUs in the caller's set). +const OVERSUBSCRIBED_SPINS: u32 = 100; + +/// One poll of a leader waiting on a section. +/// +/// Not oversubscribed: [`spin_wait`]. Oversubscribed: about +/// [`OVERSUBSCRIBED_SPINS`] PAUSE spins, then one `yield_now` per poll, so the +/// waiting participants do not hold every CPU while their peers are queued +/// behind them (the shape libomp takes when it detects oversubscription). +#[inline] +fn wait_poll(over: bool, polls: &mut u32) { + if !over { + spin_wait(); + return; + } + if *polls >= OVERSUBSCRIBED_SPINS { + std::thread::yield_now(); + } else { + *polls += 1; + std::hint::spin_loop(); + } +} + /// One member's dispatch slot. #[repr(align(128))] struct Slot { @@ -82,6 +125,10 @@ struct Slot { /// Set by the leader while the member's thread is lent to Rayon /// (`Exec::external` on a team token). lent: AtomicBool, + /// The mask this member saved before pinning itself; only the member thread + /// of this slot touches it (see [`crate::pin::SavedMask`]). + #[cfg_attr(not(feature = "parallel"), allow(dead_code))] + saved: crate::pin::SavedMask, /// Where this member sleeps once its park delay has passed; woken by the /// leader after it changed `gen`, `state` or `lent` of this slot. park: Parker, @@ -98,6 +145,7 @@ impl Slot { idle_yield: AtomicBool::new(false), park_ns: AtomicU64::new(u64::MAX), lent: AtomicBool::new(false), + saved: crate::pin::SavedMask::new(), park: Parker::new(), } } @@ -111,6 +159,8 @@ pub(crate) struct Awake { pub(crate) release_epoch: AtomicUsize, pub(crate) defaults: WaitPolicy, pub(crate) costs: crate::cost::DispatchCosts, + /// Opt-in participant pinning (`PoolBuilder::pin_participants`). + pub(crate) pin: bool, #[cfg_attr(not(feature = "parallel"), allow(dead_code))] pub(crate) acquire_timeout: Duration, #[cfg(feature = "trace")] @@ -123,6 +173,7 @@ impl Awake { defaults: WaitPolicy, acquire_timeout: Duration, costs: crate::cost::DispatchCosts, + pin: bool, ) -> Self { Awake { slots: (0..width).map(|_| Slot::new()).collect(), @@ -130,6 +181,7 @@ impl Awake { release_epoch: AtomicUsize::new(0), defaults, costs, + pin, acquire_timeout, #[cfg(feature = "trace")] trace: crate::trace::Trace::new(width), @@ -157,6 +209,34 @@ impl Awake { } } + /// Sets the park delay of the given reserved members, in nanoseconds. + /// + /// A root team scope writes it once at entry: zero (park after the short + /// spin) when the scope is oversubscribed, the effective `idle_spin` + /// otherwise. Nothing else changes it, so a member keeps the value of the + /// scope it was reserved for. + pub(crate) fn set_park_ns(&self, ids: &[usize], ns: u64) { + for &id in ids { + self.slots[id].park_ns.store(ns, Ordering::Relaxed); + } + } + + /// Restores the saved mask of member slot `id`, on the calling thread. + /// + /// Only the worker whose index is `id` may call this (`sched::unpin_worker` + /// does), and that worker owns the slot's saved mask (see + /// [`crate::pin::SavedMask`]). It is used when a pinned participant's own + /// thread is about to run work that is not its team's (`Exec::external`, + /// for example in the member that leads a `fork` child), so the foreign + /// work does not keep the team's placement. The member re-pins only at the + /// next pin section. + #[cfg_attr(not(feature = "parallel"), allow(dead_code))] + pub(crate) fn unpin_member(&self, id: usize) { + if let Some(slot) = self.slots.get(id) { + slot.saved.restore(); + } + } + /// Wakes every parked member (see [`Parker`]); a test hook, the entries wake /// only the slots they changed. #[cfg_attr(not(feature = "test-hooks"), allow(dead_code))] @@ -412,7 +492,9 @@ fn claim_job(aw: &std::sync::Arc, acq: &std::sync::Arc) { /// another member does not move it, and it is read only while polling. When it /// is older than the member's park delay the member parks and stays reserved; /// a member whose team runs no section therefore costs CPU for the park delay -/// at most, then nothing. +/// at most, then nothing. The delay is zero in an oversubscribed root scope, so +/// `polls` and `spins` are reset at the start of each idle episode there: the +/// member then spins one short episode of `polls` and parks. #[cfg(feature = "parallel")] fn member_loop(aw: &Awake, id: usize, mut seen: usize, epoch: usize) { use std::time::Instant; @@ -420,14 +502,27 @@ fn member_loop(aw: &Awake, id: usize, mut seen: usize, epoch: usize) { let mut idle_since: Option = None; let mut last_work: Option = None; let mut was_lent = false; + let mut reserved = true; let mut spins = 0u32; let mut polls = 0u32; trace!(aw, Some(id), Enter); loop { - match slot.state.load(Ordering::Acquire) { + let state = slot.state.load(Ordering::Acquire); + // Pinning ends with the reservation, and it ends on this thread: a + // member restores its own mask before it leaves `RESERVED` and before + // it serves Rayon work (a lend restores it below). + if state == RESERVED { + reserved = true; + } else if reserved { + reserved = false; + slot.saved.restore(); + } + match state { RESERVED if slot.lent.load(Ordering::Acquire) => { if !was_lent { was_lent = true; + // The thread serves the pool's jobs now: undo our pin first. + slot.saved.restore(); trace!(aw, Some(id), LentBegin); } idle_since = None; @@ -447,8 +542,14 @@ fn member_loop(aw: &Awake, id: usize, mut seen: usize, epoch: usize) { idle_since = None; let g = slot.gen.load(Ordering::Acquire); if g == seen { + let fresh = last_work.is_none(); let since = *last_work.get_or_insert_with(Instant::now); - backoff(&mut spins); + if fresh { + // A new idle episode: the park check counts from here. + spins = 0; + polls = 0; + } + spin_wait(); polls = polls.wrapping_add(1); if polls & 0x3F == 0 { let park_ns = slot.park_ns.load(Ordering::Relaxed); @@ -458,18 +559,26 @@ fn member_loop(aw: &Awake, id: usize, mut seen: usize, epoch: usize) { trace!(aw, Some(id), ParkExit); last_work = None; spins = 0; + polls = 0; } } } else { seen = g; spins = 0; + polls = 0; last_work = None; let ptr = slot.ptr.load(Ordering::Relaxed); let arg = slot.arg.load(Ordering::Relaxed); - trace!(aw, Some(id), SectionBegin); // SAFETY: the leader keeps the section alive until its // completion counter reaches zero, which happens only - // after this call decremented it. + // after the call below decremented it. + if unsafe { (*ptr).pin } { + // A pin section: this member saves its previous mask + // here, before it pins itself in its share. + slot.saved.save_if_unpinned(); + } + trace!(aw, Some(id), SectionBegin); + // SAFETY: as above. unsafe { Section::execute_member(ptr, arg) }; trace!(aw, Some(id), SectionEnd); } @@ -534,6 +643,16 @@ pub(crate) struct Section { dispatched: usize, slot_of: unsafe fn(*const (), usize) -> usize, slot_data: *const (), + /// The root scope this section belongs to has more participants than the + /// caller's CPU set has CPUs. Nested sections of the same root inherit it; + /// it is never recomputed from a child's size. + pub(crate) over: bool, + /// This is the *pin* section of a root scope (the extra section a root + /// scope runs when participant pinning is on): a member that sees it saves + /// its own mask before its share pins it. The CPU set itself lives in the + /// section's closure, which captures it. + #[cfg_attr(not(feature = "parallel"), allow(dead_code))] + pin: bool, /// Dispatched participants that have not finished. remaining: Padded, arrived: Padded, @@ -545,18 +664,39 @@ pub(crate) struct Section { } // SAFETY: `data` points to a `Sync` closure that outlives the section; every -// other field is an atomic or a mutex. +// other field is an atomic, a plain flag written before the section is +// published (and read only after a member observed the publishing `gen` with +// acquire ordering), or a function pointer. unsafe impl Sync for Section {} +/// The geometry of a section: its shares, its participants, and the flags of +/// the root scope it belongs to. +#[derive(Clone, Copy)] +pub(crate) struct Geometry { + /// Number of dispatched targets (participants minus the leader). + pub(crate) dispatched: usize, + /// Number of participants, leader included. + pub(crate) size: usize, + /// The root scope is oversubscribed. + pub(crate) over: bool, + /// This is the pin section of a root scope (see `Section::pin`). + pub(crate) pin: bool, +} + impl Section { pub(crate) fn new( call: unsafe fn(*const (), &Section, usize), data: *const (), - size: usize, - dispatched: usize, slot_of: unsafe fn(*const (), usize) -> usize, slot_data: *const (), + geometry: Geometry, ) -> Self { + let Geometry { + dispatched, + size, + over, + pin, + } = geometry; Section { call, data, @@ -564,6 +704,8 @@ impl Section { dispatched, slot_of, slot_data, + over, + pin, remaining: Padded(AtomicUsize::new(dispatched)), arrived: Padded(AtomicUsize::new(0)), epoch: Padded(AtomicUsize::new(0)), @@ -639,10 +781,13 @@ impl Section { } /// Waits until every dispatched participant has finished. + /// + /// Polls with [`spin_wait`] (one PAUSE, no yield) unless the root scope is + /// oversubscribed, where it spins about 100 times and then yields per poll. pub(crate) fn join(&self) { - let mut spins = 0; + let mut polls = 0; while self.remaining.0.load(Ordering::Acquire) != 0 { - backoff(&mut spins); + wait_poll(self.over, &mut polls); } } @@ -661,7 +806,7 @@ impl Section { self.epoch.0.store(epoch.wrapping_add(1), Ordering::Release); return true; } - let mut spins = 0; + let mut polls = 0; loop { if self.epoch.0.load(Ordering::Acquire) != epoch { return true; @@ -669,7 +814,7 @@ impl Section { if self.is_cancelled() { return false; } - backoff(&mut spins); + wait_poll(self.over, &mut polls); } } } diff --git a/src/exec.rs b/src/exec.rs index 1b0143c..e568933 100644 --- a/src/exec.rs +++ b/src/exec.rs @@ -94,6 +94,10 @@ pub(crate) struct TeamAuth<'p> { pub(crate) awake: &'p Awake, pub(crate) ids: &'p [usize], pub(crate) policy: WaitPolicy, + /// The root scope of the token is oversubscribed (see [`TeamExec`]). + /// Fixed when the root token is minted and carried into `reborrow`, + /// `limit`, `fork` children and every section. + pub(crate) over: bool, } impl Exec<'static> { @@ -434,7 +438,9 @@ impl<'p> Exec<'p> { size: budget, children: tasks, }; - split_raw(Some(t.awake), t.ids, &plan, |k, ids| body(k, Some(ids))); + split_raw(Some(t.awake), t.ids, &plan, t.over, |k, ids| { + body(k, Some(ids)); + }); } (None, Some(pool)) if tasks > 1 => { sched::run_tasks(pool, tasks, &self.track, |k| body(k, None)); @@ -484,7 +490,9 @@ impl<'p> Exec<'p> { f(split::chunk_range(len, chunk, k)); }; match (self.team, self.pool) { - (Some(t), _) if tasks > 1 => spread_raw(Some(t.awake), t.ids, tasks, body), + (Some(t), _) if tasks > 1 => { + spread_raw(Some(t.awake), t.ids, tasks, t.over, body); + } (None, Some(pool)) if tasks > 1 => sched::run_tasks(pool, tasks, track, body), _ => { for k in 0..tasks { @@ -550,7 +558,7 @@ impl<'p> Exec<'p> { // runs at once, also for a child produced by `fork`. let _released = self.track.release(); let track = &self.track; - spread_raw(Some(t.awake), t.ids, self.budget.get(), |i| { + spread_raw(Some(t.awake), t.ids, self.budget.get(), t.over, |i| { let _live = track.task(); f(i); }); @@ -635,6 +643,15 @@ impl<'p> Exec<'p> { /// its own responsibility. Without the `parallel` feature, `f` runs inline /// and receives a width of one. /// + /// # Placement + /// + /// The token's reserved members are **lent** for the call: each restores its + /// own affinity mask and serves the pool's jobs until the call returns. The + /// calling thread does the same when it is itself a pinned participant of + /// this pool (for example a member that leads a [`fork`](Exec::fork) child), + /// so the foreign work never keeps a team's single-CPU placement. Both + /// re-pin at the next pin section, that is at the next root team scope. + /// /// # Examples /// /// ``` @@ -652,6 +669,11 @@ impl<'p> Exec<'p> { // A team token lends its reserved members to the pool for the // call: the library's Rayon work runs on them as well. let _lent = self.team.map(|t| LendGuard::new(t.awake, t.ids)); + // The calling thread may itself be a pinned participant of this + // pool (the member that leads a `fork` child, or a token used + // from inside a share): it restores its own mask before it + // serves foreign work, as a lent member does. + sched::unpin_worker(pool); return sched::enter(pool, &self.track, move || f(width)); } } @@ -680,6 +702,16 @@ impl<'p> Exec<'p> { /// [`fork`](Exec::fork). It is conserved as usual: while `f` runs the token /// is borrowed. /// + /// This call opens a **root team scope**: it reads the calling thread's + /// affinity mask, which decides two things. With + /// [`PoolBuilder::pin_participants`](crate::PoolBuilder::pin_participants) on + /// it places every member on its own CPU of that mask for the scope (the + /// leader is never pinned). + /// Whether or not pinning is on, it records whether the scope is + /// *oversubscribed* (more participants than CPUs in the mask), which decides + /// how the members park and how a joining leader waits. See the design + /// record. + /// /// # Errors /// /// [`Refused::Unavailable`] when the members could not be reserved. @@ -722,12 +754,37 @@ impl<'p> Exec<'p> { ) -> Result { if let Some(t) = self.team { // Carried members: the (sub)team is available without acquisition. - return Ok(f(TeamExec { + let mut team = TeamExec { awake: Some(t.awake), ids: t.ids, policy: t.policy.with(over), track: self.track.alias(), - })); + over: t.over, + }; + // Only the root token minted by a `Team` opens a root scope; a + // child, a reborrow or a sub-team re-pins nothing of its own. + if self._leader.is_some() { + // Root-scope entry, on the executing thread: the root pressure + // and the members' park delay are decided by *this* thread's CPU + // set. A token minted on another thread (or another mask) is + // paced by the thread that enters the scope, exactly as its pin + // section is placed by that thread. Descendants inherit it. + let oversubscribed = crate::pin::oversubscribed(team.size().get()); + team.over = oversubscribed; + if let Some(awake) = team.awake { + awake.set_park_ns( + team.ids, + crate::pin::park_delay(team.policy, oversubscribed), + ); + } + if let Some(auth) = self.team.as_mut() { + auth.over = oversubscribed; + } + if team.size().get() > 1 { + pin_scope(&mut team); + } + } + return Ok(f(team)); } let pool = match self.pool { Some(pool) if self.budget.get() > 1 => pool, @@ -738,20 +795,37 @@ impl<'p> Exec<'p> { ids: &[], policy: policy.with(over), track: self.track.alias(), + over: false, })); } }; let policy = pool.wait_policy().with(over); - let reservation = - pool.awake - .acquire(pool, self.budget.get() - 1, policy, policy.idle_spin())?; + // Root-scope entry: the CPU set of this calling thread decides both the + // placement of the participants and whether the scope is + // oversubscribed, whatever the pool's pinning setting. + let oversubscribed = crate::pin::oversubscribed(self.budget.get()); + let park_spin = if oversubscribed { + Duration::ZERO + } else { + policy.idle_spin() + }; + let reservation = pool + .awake + .acquire(pool, self.budget.get() - 1, policy, park_spin)?; let mut team = TeamExec { awake: Some(&pool.awake), ids: &reservation.ids, policy, track: self.track.alias(), + over: oversubscribed, }; - team.acknowledge(); + // The pin section is also the acknowledgement; without pinning the + // participants acknowledge as before. The leader is never pinned, so + // there is nothing to restore when the scope ends or unwinds. + let pinned = team.size().get() > 1 && pin_scope(&mut team); + if !pinned { + team.acknowledge(); + } Ok(f(team)) } @@ -778,6 +852,25 @@ impl<'p> Exec<'p> { } } +/// Runs the pin section of a root team scope: every member places itself on its +/// own CPU of the calling thread's set. The leader is never pinned. +/// +/// `false` when member pinning is off for this pool, when the platform has no +/// CPU set (off Linux), or when the calling thread's mask could not be read: +/// then nothing ran and the caller acknowledges as usual. +fn pin_scope(team: &mut TeamExec<'_>) -> bool { + if !team.awake.is_some_and(|awake| awake.pin) { + return false; + } + // The leader reads its own mask, and never changes it: that set is where + // the members are placed. + let Some(set) = crate::pin::read_here() else { + return false; + }; + team.pin_section(&set); + true +} + impl fmt::Debug for Exec<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Exec") diff --git a/src/lib.rs b/src/lib.rs index c775aec..286262f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -47,6 +47,7 @@ mod exec; #[cfg(feature = "ledger")] pub mod ledger; mod park; +mod pin; mod policy; mod pool; mod reserve; diff --git a/src/pin.rs b/src/pin.rs new file mode 100644 index 0000000..53a2ddb --- /dev/null +++ b/src/pin.rs @@ -0,0 +1,259 @@ +//! Opt-in participant placement on the calling thread's CPU set. +//! +//! With [`PoolBuilder::pin_participants`](crate::PoolBuilder::pin_participants) +//! on, the leader of a root team scope reads its own affinity mask `S` and every +//! **member** pins itself onto one CPU of `S`: participant `k >= 1` onto the +//! `k`-th set CPU of `S`, wrapping modulo `|S|` when `k >= |S|`. The leader +//! (participant 0, the calling thread) is **never pinned**: its mask is never +//! changed, not even for the scope, and `S[0]` is left to it without enforcing +//! it. A vendor library creates its thread team from the thread that calls it, +//! and a team created from a one-CPU thread silently confines the whole library +//! team to that CPU (tenferro-rs `PERFORMANCE_TIPS.md`), so the calling thread +//! keeps its CPU set. A member saves its previous mask in its own slot and +//! restores it itself before it stops being reserved. +//! +//! Every affinity change is made by the thread that is being moved, with +//! `sched_setaffinity(0, ...)`, so no thread ever changes another thread's +//! placement and there is no thread-id bookkeeping. Linux only: on other +//! targets the option is accepted and does nothing. Nothing here allocates: the +//! masks are `cpu_set_t` values on the stack and the `k`-th set CPU is found by +//! scanning one. +//! +//! The same CPU set decides whether a root scope is *oversubscribed* (see +//! [`available`]): a team of more participants than the caller's mask has CPUs +//! parks its members early instead of spinning (see `awake.rs`), and +//! [`oversubscribed`] and [`park_delay`] are what the root-scope entries use. + +use core::sync::atomic::{AtomicBool, Ordering}; + +#[cfg(target_os = "linux")] +pub(crate) use sys::count; +#[cfg(not(target_os = "linux"))] +#[allow(unused_imports)] +pub(crate) use sys::count; +pub(crate) use sys::{nth, read_here, set_here, set_here_to, CpuMask}; + +/// Number of CPUs a root team scope may spread its participants over. +/// +/// On Linux the number of CPUs in the calling thread's affinity mask, i.e. the +/// same set the pin section places participants on. Elsewhere +/// [`std::thread::available_parallelism`]. Zero is never returned. +pub(crate) fn available() -> usize { + #[cfg(target_os = "linux")] + if let Some(set) = read_here() { + let n = count(&set); + if n > 0 { + return n; + } + } + std::thread::available_parallelism().map_or(1, |n| n.get()) +} + +/// Whether a root scope with `participants` participants is oversubscribed: it +/// has more participants than the calling thread's CPU set has CPUs. +/// +/// Evaluated on the thread that enters the root scope, whether or not pinning +/// is on. For a [`Team`](crate::Team) that is the thread that mints the root +/// token *and* again the thread that opens each team scope through it, so a +/// token that travels to a thread with a different CPU set is paced by that +/// thread's set. +pub(crate) fn oversubscribed(participants: usize) -> bool { + participants > available() +} + +/// The park delay of the reserved members of a root scope, in nanoseconds: +/// zero when the scope is oversubscribed (park after the short spin), the +/// scope's policy delay otherwise. +pub(crate) fn park_delay(policy: crate::policy::WaitPolicy, oversubscribed: bool) -> u64 { + if oversubscribed { + 0 + } else { + u64::try_from(policy.idle_spin().as_nanos()).unwrap_or(u64::MAX) + } +} + +/// Pins the calling thread onto the `k`-th CPU of `set` (wrapping). +/// +/// Called by a member for itself (`k >= 1`); the leader never calls it. Best +/// effort: an empty set (never produced by [`available`]), a refused syscall, or +/// the calling thread already being outside `set` leaves the thread where it +/// was. +pub(crate) fn pin_to(set: &CpuMask, k: usize) { + if let Some(cpu) = nth(set, k) { + set_here_to(cpu); + } +} + +/// The mask a member saved before it pinned itself, and whether it is pinned. +/// +/// Written and read only by the thread that owns the slot: a slot id `i` is +/// only ever entered by the worker with index `i` (`claim_job` uses +/// `rayon::current_thread_index()`), and only that thread touches this value — +/// from its `member_loop`, and from `Awake::unpin_member` when the same thread +/// is about to run work that is not its team's (`Exec::external`). The leader +/// never reads it, which is why one saved mask per slot and no +/// cross-reservation cache is enough. +#[cfg_attr(not(feature = "parallel"), allow(dead_code))] +pub(crate) struct SavedMask { + mask: core::cell::UnsafeCell>, + pinned: AtomicBool, +} + +// SAFETY: the `UnsafeCell` is only reached through `save_if_unpinned` and +// `restore`, and both are called exclusively from threads that own the slot: a +// pool worker with index `i` is one fixed thread, it runs the member loop of +// slot `i` only, and `Awake::unpin_member(i)` (`sched::unpin_worker`) is called +// on that same worker. The state those calls leave behind (a saved mask, +// `pinned`) is restored before the loop returns the thread to Rayon, so no +// other thread observes it. `pinned` itself is an atomic, so even a spurious +// read from elsewhere would be sound. +unsafe impl Sync for SavedMask {} + +#[cfg_attr(not(feature = "parallel"), allow(dead_code))] +impl SavedMask { + pub(crate) const fn new() -> Self { + SavedMask { + mask: core::cell::UnsafeCell::new(core::mem::MaybeUninit::uninit()), + pinned: AtomicBool::new(false), + } + } + + /// Saves the calling thread's mask unless this member is already pinned. + /// + /// A persistent team re-pins its members at every root scope and keeps the + /// first saved mask, so a member that was never returned to Rayon restores + /// what it had before the first scope. + pub(crate) fn save_if_unpinned(&self) { + if self.pinned.load(Ordering::Relaxed) { + return; + } + // `read_here` is `None` off Linux and when the kernel refuses a mask + // larger than `cpu_set_t`: then nothing was changed and nothing has to + // be restored. + if let Some(mask) = read_here() { + // SAFETY: this thread owns the cell (see the `Sync` note). + unsafe { (*self.mask.get()).write(mask) }; + self.pinned.store(true, Ordering::Relaxed); + } + } + + /// Restores the saved mask on the calling thread, if this member is pinned. + /// + /// Best effort: when the kernel refuses the saved mask (a cpuset changed + /// under us) the saved value is kept for a later attempt. + pub(crate) fn restore(&self) { + if !self.pinned.load(Ordering::Relaxed) { + return; + } + // SAFETY: written by `save_if_unpinned` on this thread (see `Sync`). + let mask = unsafe { (*self.mask.get()).assume_init_ref() }; + if set_here(mask) { + self.pinned.store(false, Ordering::Relaxed); + } + } +} + +#[cfg(target_os = "linux")] +mod sys { + use core::mem; + + /// The most CPUs a `cpu_set_t` can hold: its representable bits (1024 on + /// Linux, 128 bytes). + const SETSIZE: usize = 8 * mem::size_of::(); + + /// A thread's CPU set, on the stack (`cpu_set_t`, 128 bytes). + #[derive(Clone, Copy)] + pub(crate) struct CpuMask { + set: libc::cpu_set_t, + } + + /// The calling thread's affinity mask, or `None` when the syscall refuses + /// it (a kernel mask larger than `cpu_set_t`, no current thread). + pub(crate) fn read_here() -> Option { + // SAFETY: a `cpu_set_t` is plain data; the call writes at most its size. + let mut set = unsafe { mem::zeroed::() }; + // SAFETY: `set` is a valid, correctly sized `cpu_set_t`; pid 0 is the + // calling thread, so this thread is the one that is read. + let r = unsafe { libc::sched_getaffinity(0, mem::size_of::(), &mut set) }; + (r == 0).then_some(CpuMask { set }) + } + + /// Sets the calling thread's mask. `false` when the kernel refuses it; the + /// thread keeps the mask it had. + pub(crate) fn set_here(mask: &CpuMask) -> bool { + // SAFETY: reads a valid `cpu_set_t` of the right size for pid 0 (the + // calling thread): the calling thread is the one being moved. + let r = unsafe { libc::sched_setaffinity(0, mem::size_of::(), &mask.set) }; + r == 0 + } + + /// Pins the calling thread onto exactly `cpu`. + pub(crate) fn set_here_to(cpu: usize) -> bool { + if cpu >= SETSIZE { + return false; + } + // SAFETY: a zeroed `cpu_set_t`; `CPU_ZERO`/`CPU_SET` only clear and set + // bits, and `cpu < SETSIZE` keeps the index inside the set. + unsafe { + let mut set = mem::zeroed::(); + libc::CPU_ZERO(&mut set); + libc::CPU_SET(cpu, &mut set); + set_here(&CpuMask { set }) + } + } + + /// Number of CPUs in the set. + pub(crate) fn count(mask: &CpuMask) -> usize { + // SAFETY: only counts the bits of a valid set. + let n = unsafe { libc::CPU_COUNT(&mask.set) }; + n.max(0) as usize + } + + /// Index of the `k`-th set CPU, wrapping modulo the set size; `None` for an + /// empty set (never pin to an empty mask). + pub(crate) fn nth(mask: &CpuMask, k: usize) -> Option { + let n = count(mask); + if n == 0 { + return None; + } + let mut want = k % n; + for cpu in 0..SETSIZE { + // SAFETY: reads one bit of a valid set. + if unsafe { libc::CPU_ISSET(cpu, &mask.set) } { + if want == 0 { + return Some(cpu); + } + want -= 1; + } + } + None + } +} + +#[cfg(not(target_os = "linux"))] +#[allow(dead_code)] // the option is accepted and does nothing off Linux +mod sys { + /// No CPU set off Linux: every operation is a no-op. + #[derive(Clone, Copy)] + pub(crate) struct CpuMask(()); + + pub(crate) fn read_here() -> Option { + None + } + + pub(crate) fn set_here(_mask: &CpuMask) -> bool { + false + } + + pub(crate) fn set_here_to(_cpu: usize) -> bool { + false + } + + pub(crate) fn count(_mask: &CpuMask) -> usize { + 0 + } + + pub(crate) fn nth(_mask: &CpuMask, _k: usize) -> Option { + None + } +} diff --git a/src/pool.rs b/src/pool.rs index c24d56d..4b177a0 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -19,6 +19,18 @@ const DEFAULT_STACK_BYTES: usize = 16 << 20; /// Default time an acquisition waits for idle workers to wake: 2 ms. const DEFAULT_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(2); +/// Whether a pool pins the members of a root team scope by default. +/// +/// The one switch for the default of [`PoolBuilder::pin_participants`]: off, as +/// in every compared runtime (libgomp, libomp, OpenBLAS, faer) and as measured +/// on broccoli, where the slow mode of tensor4all/faer#107 disappears only with +/// member placement (arm C: 0 of 20 processes slow, 0 % stacked samples) while it +/// persists with a free CPU (arm B: 7 of 20). Turning it on by default would take +/// CPUs away from other users of the process, fight the caller's `taskset`, and +/// confine any vendor thread team a member creates, so the decision stays per +/// pool. +const PIN_PARTICIPANTS_DEFAULT: bool = false; + /// A fixed set of worker threads that [`Exec`](crate::Exec) tokens place work on. /// /// `Pool` is the only way to make a pool. It wraps the scheduler privately: no @@ -146,6 +158,20 @@ impl Pool { self.awake.costs } + /// Whether the pool pins the **members** of a root team scope (see + /// [`PoolBuilder::pin_participants`]); the leader is never pinned. + /// + /// # Examples + /// + /// ``` + /// let pool = threadteam::Pool::builder().threads(1).build()?; + /// assert!(!pool.pin_participants()); // off by default + /// # Ok::<(), threadteam::Error>(()) + /// ``` + pub fn pin_participants(&self) -> bool { + self.awake.pin + } + /// Number of workers currently kept awake as team members, reserved or /// idle. Zero once idle members returned to Rayon. A diagnostic: the value /// can change at any moment. @@ -225,6 +251,7 @@ pub struct PoolBuilder { acquire_timeout: Duration, qos: Option, costs: DispatchCosts, + pin_participants: bool, } impl PoolBuilder { @@ -238,6 +265,7 @@ impl PoolBuilder { acquire_timeout: DEFAULT_ACQUIRE_TIMEOUT, qos: None, costs: DispatchCosts::host_default(), + pin_participants: PIN_PARTICIPANTS_DEFAULT, } } @@ -423,6 +451,86 @@ impl PoolBuilder { self } + /// Pins the **members** of a root team scope onto the calling thread's CPU + /// set. Off by default; Linux only, where it does something. + /// + /// With it on, the leader of a root team scope + /// ([`Exec::team`](crate::Exec::team), or the root token of a + /// [`Team`](crate::Team)) reads **its own** affinity mask `S` with + /// `sched_getaffinity`, and every member pins itself with + /// `sched_setaffinity` onto one CPU of `S`: member `k` (participant `k`, + /// counting the leader as 0) onto the `k`-th set CPU of `S`, wrapping modulo + /// `|S|` when the team is wider than the set. `S[0]` is thereby left to the + /// leader, without enforcing it. + /// + /// **The leader is never pinned.** Its mask is never changed, not even for + /// the scope: it keeps all of `S` while the members work. A vendor library + /// (OpenMP, OpenBLAS, a LAPACK provider) creates its thread team from the + /// thread that calls it, and a team created from a one-CPU thread silently + /// confines the whole provider team to that CPU — measured at 3-20x slower + /// dgemm and LU at 4 threads. Engines call such libraries from the + /// coordinating thread, so that thread must keep its CPU set. + /// + /// A member saves its previous mask in its own slot and restores it + /// *itself* before it stops being reserved: before it returns to Rayon, and + /// before its thread serves work that is not its team's — while it is lent + /// by [`Exec::external`](crate::Exec::external), and when it runs such work + /// itself (a member that leads a `fork` child and calls `external`). No + /// thread ever changes another thread's affinity. + /// + /// # What it is for + /// + /// It keeps the members of a team apart. Measured on broccoli (16 threads on + /// 16 allowed CPUs, 20 processes, tensor4all/faer#107): with the scheduler + /// free to place the threads, 11 of 20 processes ran a `qr_factor` 512 at + /// the one-thread time, with 92 % of the samples showing two runnable team + /// threads on one CPU; with the participants pinned, 0 of 20 did, and no + /// sample was stacked. Giving the team one spare CPU did not help (7 of 20 + /// slow). That arm pinned every participant including the leader; this crate + /// leaves the leader on the caller's whole set, which is what the rule above + /// requires and only relaxes the placement of participant 0. + /// + /// # Assumptions and limits + /// + /// * **Linux only.** The CPU set is a stack `cpu_set_t` and the placement is + /// `sched_setaffinity(0, ...)`. On every other target the option is + /// accepted and does nothing. + /// * **The caller's CPU set is the source of placement**, so the caller's + /// mask should belong to this team alone. Two teams whose leaders see the + /// same mask place their members on the same CPUs; give each team (in + /// practice: each process) its own CPU set (`taskset -c`). + /// * **Best effort.** A refused syscall (a mask larger than `cpu_set_t`, a + /// cpuset that changed under us, a CPU offlined) leaves that thread where + /// it was; an empty mask pins nobody; a failed pin or restore never fails + /// the scope. + /// * **A pinned member must not create a foreign thread team.** A member + /// that calls OpenMP, OpenBLAS or any library that builds its team from + /// the calling thread *directly inside a share* confines that library's + /// team to one CPU. Call it through + /// [`Exec::external`](crate::Exec::external), which restores the member's + /// mask first, or keep this option off. + /// * Only the members of a team are placed. A token of a + /// [`Team`](crate::Team) that never opens a team scope + /// (`exec.for_ranges(...)` directly) is not re-pinned; its members keep + /// the placement of the last root scope (or of `Pool::into_team`). + /// + /// Off by default, consistently with the rule above and with libgomp, + /// libomp, OpenBLAS and faer, which leave threads unpinned unless asked. + /// + /// # Examples + /// + /// ``` + /// let pool = threadteam::Pool::builder().threads(2).pin_participants(true).build()?; + /// assert!(pool.pin_participants()); + /// # let _ = pool; + /// # Ok::<(), threadteam::Error>(()) + /// ``` + #[must_use] + pub fn pin_participants(mut self, on: bool) -> Self { + self.pin_participants = on; + self + } + /// Starts the workers. /// /// Without the `parallel` feature no thread is started and the pool has @@ -504,6 +612,7 @@ impl PoolBuilder { self.wait, self.acquire_timeout, self.costs, + self.pin_participants, )), }) } @@ -512,7 +621,13 @@ impl PoolBuilder { fn build_checked(self) -> Result { Ok(Pool { width: NonZeroUsize::MIN, - awake: Arc::new(Awake::new(1, self.wait, self.acquire_timeout, self.costs)), + awake: Arc::new(Awake::new( + 1, + self.wait, + self.acquire_timeout, + self.costs, + self.pin_participants, + )), }) } } @@ -524,6 +639,7 @@ impl fmt::Debug for PoolBuilder { .field("stack_size", &self.stack_size) .field("wait", &self.wait) .field("acquire_timeout", &self.acquire_timeout) + .field("pin_participants", &self.pin_participants) .finish_non_exhaustive() } } diff --git a/src/reserve.rs b/src/reserve.rs index 7d488ae..4ec6b96 100644 --- a/src/reserve.rs +++ b/src/reserve.rs @@ -303,10 +303,23 @@ impl Team { fn token(&self, budget: usize) -> Exec<'_> { let budget = budget.clamp(1, self.size().get()); + let ids = &self.reservation.ids[..budget - 1]; + // Root-scope entry of the token: `exec` is serialized (one leader at a + // time), so the calling thread's CPU set decides here whether the scope + // is oversubscribed. When it is, the members park after their short + // spin instead of spinning for the whole park delay; otherwise they get + // the policy's delay. `Exec::team_with` on the minted token repeats + // this on the thread that enters each scope, so a token that travels to + // a thread with a different CPU set is paced by that set. + let oversubscribed = crate::pin::oversubscribed(budget); + self.pool + .awake + .set_park_ns(ids, crate::pin::park_delay(self.policy, oversubscribed)); let auth = TeamAuth { awake: &self.pool.awake, - ids: &self.reservation.ids[..budget - 1], + ids, policy: self.policy, + over: oversubscribed, }; Exec::from_team(&self.pool, auth, LeaderGuard(&self.gate)) } diff --git a/src/sched.rs b/src/sched.rs index 9d13dc4..b1f8a21 100644 --- a/src/sched.rs +++ b/src/sched.rs @@ -18,6 +18,24 @@ pub(crate) fn in_pool(_pool: &Pool) -> bool { false } +/// Restores the calling thread's own affinity mask when it is a reserved, +/// pinned member of `pool`. +/// +/// A participant can be the thread that runs work which is not its team's: the +/// member that leads a `fork` child and calls +/// [`Exec::external`](crate::Exec::external), or a token used from inside a +/// share. Such a thread must not keep the team's single-CPU placement while it +/// serves that work, and only it may change its own affinity, so it restores +/// the mask it saved before pinning (see `pin::SavedMask`). A no-op when the +/// calling thread is not a worker of `pool` or is not pinned; the member +/// re-pins at the next pin section. +#[cfg(feature = "parallel")] +pub(crate) fn unpin_worker(pool: &Pool) { + if let Some(worker) = pool.inner.current_thread_index() { + pool.awake.unpin_member(worker); + } +} + /// Runs `task(0..tasks)` concurrently on `pool` and waits for all of them. /// /// The tasks form a recursive binary `rayon::join` tree, whose jobs live on the diff --git a/src/team.rs b/src/team.rs index a8a9a67..9bac2b4 100644 --- a/src/team.rs +++ b/src/team.rs @@ -1,6 +1,6 @@ //! Reserved awake members: the scoped authority [`TeamExec`]. -use crate::awake::{Awake, JoinGuard, Section}; +use crate::awake::{Awake, Geometry, JoinGuard, Section}; use crate::cost::{pays_at, width_for_by, DispatchCosts, DispatchKind}; use crate::error::Error; use crate::policy::{WaitOverride, WaitPolicy}; @@ -79,10 +79,18 @@ impl Participant<'_> { /// `budget - 1` pool workers. Before any section runs, every participant /// acknowledged on its own thread (a rendezvous), so barriers inside /// sections cannot deadlock for lack of threads. -/// * Reserved members spin, then park once the park delay (`idle_spin`) has -/// passed since their own last share; they run no unrelated Rayon job and do -/// not leave on a timeout, also between sections and during a long serial -/// phase of the leader. They are returned when the scope ends. +/// * Reserved members spin with one PAUSE per poll, then park once the park +/// delay (`idle_spin`) has passed since their own last share; they run no +/// unrelated Rayon job and do not leave on a timeout, also between sections +/// and during a long serial phase of the leader. They are returned when the +/// scope ends. In a root scope with more participants than CPUs in the +/// caller's affinity mask the delay is zero, so they park after a short spin +/// instead of holding a CPU. +/// * If the pool pins members +/// ([`PoolBuilder::pin_participants`](crate::PoolBuilder::pin_participants)), +/// member `k` of a root scope (participant `k >= 1`) runs on the `k`-th CPU of +/// the calling thread's set, and restores its own mask afterwards. The leader +/// is never pinned. /// * A panic in any participant propagates to the leader after every /// dispatched share has finished; the reservation is released. /// @@ -180,6 +188,12 @@ pub struct TeamExec<'team> { pub(crate) ids: &'team [usize], pub(crate) policy: WaitPolicy, pub(crate) track: Track, + /// The root scope of this authority is oversubscribed: it has more + /// participants than the caller's CPU set has CPUs. Fixed at root-scope + /// entry, inherited by reborrows, splits and `fork` children, and put on + /// every [`Section`] so that nested sections of an oversubscribed root + /// know it too. + pub(crate) over: bool, } /// How the members of a split are divided. @@ -318,6 +332,7 @@ impl<'team> TeamExec<'team> { ids: self.ids, policy: self.policy, track: self.track.alias(), + over: self.over, } } @@ -361,6 +376,17 @@ impl<'team> TeamExec<'team> { /// A participant that panics cancels the section's barriers, so peers /// blocked in [`Participant::barrier`] return `Err(Cancelled)`. pub fn run_with(self, f: F) + where + F: Fn(&Participant<'_>) + Sync, + { + self.run_with_pin(f, false); + } + + /// [`run_with`](Self::run_with), or the pin section of a root scope when + /// `pin` is set: the members save their own mask before their share runs + /// (see `member_loop`) and every member pins itself in the share. The + /// leader's share never pins. + pub(crate) fn run_with_pin(self, f: F, pin: bool) where F: Fn(&Participant<'_>) + Sync, { @@ -371,7 +397,37 @@ impl<'team> TeamExec<'team> { f(p); }; let ids = self.ids; - section(self.awake, &|pos| ids[pos], ids.len(), ids.len() + 1, &body); + let geometry = Geometry { + dispatched: ids.len(), + size: ids.len() + 1, + over: self.over, + pin, + }; + section(self.awake, &|pos| ids[pos], &body, geometry); + } + + /// Places every member on its own CPU of `set`: member `k` (participant + /// `k >= 1`) pins itself onto the `k`-th set CPU, wrapping modulo `|set|`. + /// The leader, participant 0, is **never pinned**: it keeps its own mask, + /// `set`, so a vendor team it creates inherits the caller's CPU set instead + /// of one CPU (see [`PoolBuilder::pin_participants`]). `set[0]` is therefore + /// left to the leader, without enforcing it. + /// + /// One section of the existing machinery, so all participants are on their + /// own threads; it also serves as the rendezvous that precedes any user + /// section. Each member changes its own affinity, and saves its previous + /// mask first (once per reservation). Best effort: an empty or refused set + /// leaves a member where it was. + pub(crate) fn pin_section(&mut self, set: &crate::pin::CpuMask) { + self.reborrow().run_with_pin( + |p| { + let _ = p.barrier(); + if p.index() > 0 { + crate::pin::pin_to(set, p.index()); + } + }, + true, + ); } /// Divides the team into child reservations of the given sizes. @@ -461,14 +517,15 @@ impl<'team> TeamExec<'team> { { let track = &self.track; let _released = track.release(); - let (awake, ids, policy) = (self.awake, self.ids, self.policy); - split_raw(awake, ids, plan, |c, child_ids| { + let (awake, ids, policy, over) = (self.awake, self.ids, self.policy, self.over); + split_raw(awake, ids, plan, over, |c, child_ids| { let _live = track.task(); let child = TeamExec { awake, ids: child_ids, policy, track: track.child(), + over, }; f(c, child); }); @@ -491,13 +548,8 @@ impl<'team> TeamExec<'team> { /// /// The one implementation of a section: `TeamExec` and team-carrying `Exec` /// tokens both dispatch through it. -pub(crate) fn section( - awake: Option<&Awake>, - slot_of: &S, - dispatched: usize, - size: usize, - body: &B, -) where +pub(crate) fn section(awake: Option<&Awake>, slot_of: &S, body: &B, geometry: Geometry) +where B: Fn(&Participant<'_>) + Sync, S: Fn(usize) -> usize + Sync, { @@ -516,11 +568,11 @@ pub(crate) fn section( let sec = Section::new( call::, (body as *const B).cast(), - size, - dispatched, slot_call::, (slot_of as *const S).cast(), + geometry, ); + let dispatched = geometry.dispatched; let guard = JoinGuard(&sec); if dispatched > 0 { if let Some(awake) = awake { @@ -540,7 +592,7 @@ pub(crate) fn section( /// Runs `f(c, child_ids)` for every child of `plan` at once: child 0 on the /// calling thread, child `c` on the first member of its index range. -pub(crate) fn split_raw(awake: Option<&Awake>, ids: &[usize], plan: &Plan<'_>, f: F) +pub(crate) fn split_raw(awake: Option<&Awake>, ids: &[usize], plan: &Plan<'_>, over: bool, f: F) where F: Fn(usize, &[usize]) + Sync, { @@ -550,28 +602,38 @@ where f(c, &ids[start..start + size - 1]); }; let n = plan.children(); + let geometry = Geometry { + dispatched: n - 1, + size: ids.len() + 1, + over, + pin: false, + }; section( awake, &|pos| ids[plan.child(pos + 1).0 - 1], - n - 1, - ids.len() + 1, &body, + geometry, ); } /// Runs `f(k)` for `k in 0..participants` at once (leader is `k = 0`). -pub(crate) fn spread_raw(awake: Option<&Awake>, ids: &[usize], participants: usize, f: F) -where +pub(crate) fn spread_raw( + awake: Option<&Awake>, + ids: &[usize], + participants: usize, + over: bool, + f: F, +) where F: Fn(usize) + Sync, { let body = |p: &Participant<'_>| f(p.index()); - section( - awake, - &|pos| ids[pos], - participants - 1, - participants, - &body, - ); + let geometry = Geometry { + dispatched: participants - 1, + size: participants, + over, + pin: false, + }; + section(awake, &|pos| ids[pos], &body, geometry); } impl fmt::Debug for TeamExec<'_> { diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 4cf30e1..ea7a00b 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1 +1,3 @@ +#[cfg(all(feature = "parallel", target_os = "linux"))] +mod placement; mod split; diff --git a/src/tests/placement.rs b/src/tests/placement.rs new file mode 100644 index 0000000..5678605 --- /dev/null +++ b/src/tests/placement.rs @@ -0,0 +1,180 @@ +//! Root-scope flags that the placement work carries: whether a root team scope +//! is oversubscribed reaches reborrows, `split`/`fork` children, team tokens and +//! every section of the root. + +use crate::awake::{Geometry, Section}; +use crate::{Exec, Pool}; +use std::time::Duration; + +fn own_cpus() -> Vec { + // SAFETY: reads the calling thread's mask into a zeroed `cpu_set_t`. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + let r = libc::sched_getaffinity(0, std::mem::size_of::(), &mut set); + assert_eq!(r, 0, "sched_getaffinity failed"); + (0..libc::CPU_SETSIZE as usize) + .filter(|&cpu| libc::CPU_ISSET(cpu, &set)) + .collect() + } +} + +fn set_own_cpus(cpus: &[usize]) -> bool { + // SAFETY: builds a `cpu_set_t` and applies it to pid 0, i.e. this thread. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + libc::CPU_ZERO(&mut set); + for &cpu in cpus { + libc::CPU_SET(cpu, &mut set); + } + libc::sched_setaffinity(0, std::mem::size_of::(), &set) == 0 + } +} + +fn pool(width: usize) -> Pool { + Pool::builder() + .threads(width) + .idle_spin(Duration::ZERO) + .acquire_timeout(Duration::from_secs(10)) + .build() + .unwrap() +} + +/// Runs `f` with the calling thread restricted to two CPUs. +fn on_two_cpus(f: impl FnOnce()) { + let all = own_cpus(); + if all.len() < 3 { + eprintln!("skipped: need at least 3 allowed CPUs to restrict to 2"); + return; + } + if !set_own_cpus(&all[..2]) { + eprintln!("skipped: cannot restrict the calling thread to two CPUs"); + return; + } + f(); + assert!(set_own_cpus(&all), "restoring the calling thread's mask"); +} + +#[test] +fn a_section_carries_the_oversubscription_flag_it_was_built_with() { + unsafe fn call(_: *const (), _: &Section, _: usize) {} + unsafe fn slot_of(_: *const (), pos: usize) -> usize { + pos + } + let geometry = |over| Geometry { + dispatched: 0, + size: 1, + over, + pin: false, + }; + let quiet = Section::new( + call, + std::ptr::null(), + slot_of, + std::ptr::null(), + geometry(false), + ); + assert!(!quiet.over); + let crowded = Section::new( + call, + std::ptr::null(), + slot_of, + std::ptr::null(), + geometry(true), + ); + assert!(crowded.over); +} + +#[test] +fn an_oversubscribed_root_keeps_the_flag_through_reborrows_splits_and_fork_children() { + on_two_cpus(|| { + let width = 4; + let pool = pool(width); + + // A scoped acquisition whose participants outnumber the caller's CPUs. + let mut exec = Exec::new(&pool, width).unwrap(); + exec.team(|mut team| { + assert!(team.over, "4 participants on 2 CPUs is oversubscribed"); + let over = team.reborrow().over; // reborrows inherit it + assert!(over); + team.reborrow() + .split(&[2, 2], |_, child| { + assert!(child.over, "a split child of an oversubscribed root"); + child.run(|_, _| {}); + }) + .unwrap(); + team.fork(2, |_, child| { + assert!(child.over, "a fork child of an oversubscribed root"); + }); + }) + .expect("an idle pool"); + drop(exec); + + // A persistent team fixes the flag at its root token and carries it + // into the carried team and into `fork` children of the token. + let team = pool.into_team(width).expect("an idle pool"); + let mut exec = team.exec(width); + assert!(exec.team(|t| t.over).expect("a team token opens its scope")); + exec.reborrow().fork(2, |_, mut child| { + assert!( + child + .team(|t| t.over) + .expect("a fork child carries the sub-team"), + "a fork child of an oversubscribed team token" + ); + }); + }); +} + +#[test] +fn a_root_scope_that_fits_the_callers_cpus_is_not_oversubscribed() { + let all = own_cpus(); + let width = all.len().min(3); + let pool = pool(width); + let mut exec = Exec::new(&pool, width).unwrap(); + exec.team(|team| { + assert!( + !team.over, + "{width} participants on {} CPUs is not oversubscribed", + all.len() + ); + }) + .expect("an idle pool"); +} + +#[test] +fn a_team_scope_takes_the_oversubscription_from_the_thread_that_enters_it() { + let all = own_cpus(); + if all.len() < 4 { + eprintln!("skipped: need at least 4 allowed CPUs for a one-CPU scope mask"); + return; + } + let width = 4; + let pool = pool(width); + let team = pool.into_team(width).expect("an idle pool"); + // Mint the root token on the full mask: four participants on four CPUs is + // not oversubscribed at mint. + let mut exec = team.exec(width); + // Enter the scope from a thread whose mask holds one CPU (the token could + // just as well have travelled to another thread): the scope decides from + // that set, exactly as its pin section does. + if !set_own_cpus(&all[..1]) { + eprintln!("skipped: cannot restrict the calling thread to one CPU"); + return; + } + let scoped = exec + .team(|mut team| { + let root = team.over; + let child = std::sync::atomic::AtomicBool::new(false); + team.reborrow().fork(2, |_, child_team| { + child.store(child_team.over, std::sync::atomic::Ordering::Relaxed); + }); + (root, child.load(std::sync::atomic::Ordering::Relaxed)) + }) + .expect("a team token opens its scope"); + assert!(set_own_cpus(&all), "restoring the calling thread's mask"); + assert!( + scoped.0, + "four participants on a one-CPU scope mask is oversubscribed" + ); + assert!(scoped.1, "a fork child inherits the root's pressure"); +} diff --git a/tests/alloc_free.rs b/tests/alloc_free.rs index ed18236..28e73ff 100644 --- a/tests/alloc_free.rs +++ b/tests/alloc_free.rs @@ -80,4 +80,28 @@ fn fork_for_ranges_external_do_not_allocate() { 0, "in-pool entries allocated" ); + + // Participant pinning on: the extra pin section a root team scope runs must + // not allocate either. The window contains the acquisition's id `Vec`, the + // one allocation of an acquisition documented in the design record, so it + // is compared against the same window with pinning off rather than + // against zero. + let team_scope_window = |pin: bool| { + let pool = Pool::builder() + .threads(3) + .pin_participants(pin) + .idle_spin(std::time::Duration::from_millis(1)) + .acquire_timeout(std::time::Duration::from_secs(5)) + .build() + .unwrap(); + let scope = |pool: &Pool| { + let mut exec = Exec::new(pool, 3).unwrap(); + let _live = exec.team(|mut team| team.reborrow().run(|_, _| {})); + }; + measure(10, || scope(&pool)) + }; + let off = team_scope_window(false); + let on = team_scope_window(true); + assert!(off <= 1, "a team scope allocates at most the id Vec: {off}"); + assert_eq!(on, off, "participant pinning added an allocation"); } diff --git a/tests/idle_clock.rs b/tests/idle_clock.rs index 4def7eb..b14ce83 100644 --- a/tests/idle_clock.rs +++ b/tests/idle_clock.rs @@ -1,16 +1,32 @@ //! The idle clock of a reserved member (feature `trace`): the time since the -//! member's own last share (or since it started), per member. A member parks -//! once that exceeds the park delay and never earlier; sections that arrive -//! sooner keep it spinning. Timings are generous; the test takes well under a -//! second. +//! member's own last share (or since it started, or since it woke from a park). +//! +//! Everything here is judged by the **order of a member's own trace events**, +//! never by a wall-clock bound on how fast the leader gets around to the next +//! section: on a small or loaded host one `sleep` can wake tens of milliseconds +//! late, so "sections every 5 ms" is not something a test may assume. Member +//! slot ids are read from the trace as well, because the acquisition may claim +//! any free slots. +//! +//! * Never before the delay: for every park of member `m` at time `t`, +//! `t - max(m's last share end, m's last wake-up, m's start) >= SPIN`. This +//! holds for every park in the whole test. +//! * After the delay: once no section arrives, every member parks; polled until +//! it happens, with no upper bound on how late. +//! * A following section wakes every member that was parked: the number of +//! `ParkExit`s rises by the number of members parked at that moment. #![cfg(all(feature = "trace", feature = "parallel"))] -use std::time::Duration; +use std::time::{Duration, Instant}; use threadteam::trace::{self, Kind}; use threadteam::Pool; const SPIN: Duration = Duration::from_millis(40); +/// Only a hang guard: no assertion below is a bound on how long the members or +/// the leader take. +const DEADLINE: Duration = Duration::from_secs(10); + fn team(width: usize) -> threadteam::Team { Pool::builder() .threads(width) @@ -22,59 +38,147 @@ fn team(width: usize) -> threadteam::Team { .unwrap() } -fn kinds(team: &threadteam::Team, kind: Kind) -> Vec<(usize, u64)> { +/// The member events of the pool, in time order. +fn member_events(team: &threadteam::Team) -> Vec { trace::events(team.pool()) .into_iter() - .filter(|e| e.kind == kind && e.member != usize::MAX) - .map(|e| (e.member, e.ns)) + .filter(|e| e.member != usize::MAX) .collect() } +fn count(team: &threadteam::Team, kind: Kind) -> usize { + member_events(team) + .iter() + .filter(|e| e.kind == kind) + .count() +} + +/// The members that ran a share, in the order of their first share. +fn members_that_ran(events: &[trace::Event]) -> Vec { + let mut out: Vec = Vec::new(); + for e in events { + if e.kind == Kind::SectionEnd && !out.contains(&e.member) { + out.push(e.member); + } + } + out +} + +/// Members asleep now: their latest park event is a `ParkEnter`. +fn parked_now(team: &threadteam::Team) -> usize { + let mut parked = std::collections::HashSet::new(); + for e in member_events(team) { + match e.kind { + Kind::ParkEnter => { + parked.insert(e.member); + } + Kind::ParkExit => { + parked.remove(&e.member); + } + _ => {} + } + } + parked.len() +} + +/// Waits (hang guard only) until `cond` holds, polling the trace. +fn wait_until(what: &str, mut cond: impl FnMut() -> bool) { + let start = Instant::now(); + while !cond() { + assert!(start.elapsed() < DEADLINE, "timeout waiting for {what}"); + std::thread::sleep(Duration::from_millis(2)); + } +} + +/// Did every member that ran a share park after its own last share end? +fn every_member_parked_after_its_last_share(events: &[trace::Event]) -> bool { + let members = members_that_ran(events); + !members.is_empty() + && members.iter().all(|&member| { + let last_end = events + .iter() + .filter(|e| e.member == member && e.kind == Kind::SectionEnd) + .map(|e| e.ns) + .max(); + let park = events + .iter() + .filter(|e| e.member == member && e.kind == Kind::ParkEnter) + .map(|e| e.ns) + .max(); + matches!((last_end, park), (Some(end), Some(park)) if park > end) + }) +} + +/// Asserts the "never before the delay" property over every park in the trace. +fn assert_no_park_before_the_delay(team: &threadteam::Team) { + let events = member_events(team); + let parks: Vec<_> = events + .iter() + .filter(|e| e.kind == Kind::ParkEnter) + .collect(); + assert!(!parks.is_empty(), "no member parked at all"); + for park in parks { + // The park delay runs from the later of the member's own last share + // end, its own last wake-up and its start. + let anchor = events + .iter() + .filter(|e| e.member == park.member && e.ns <= park.ns) + .filter(|e| matches!(e.kind, Kind::SectionEnd | Kind::ParkExit | Kind::Enter)) + .map(|e| e.ns) + .max(); + let Some(anchor) = anchor else { + panic!("member {} parked without an anchor event", park.member); + }; + assert!( + park.ns - anchor >= SPIN.as_nanos() as u64, + "member {} parked {} us after its own last share, wake-up or start", + park.member, + (park.ns - anchor) / 1000 + ); + } +} + #[test] -fn members_park_after_the_park_delay_since_their_last_share_and_not_before() { +fn members_park_after_the_park_delay_since_their_own_last_share_or_wake() { let team = team(3); let mut exec = team.exec(3); if team.size().get() < 3 { return; } - // Sections 5 ms apart for 150 ms: nobody parks (park delay 40 ms). + // Sections a few milliseconds apart for about 150 ms. This does *not* + // assert that nobody parks: on a loaded host the leader's sleeps can wake + // late enough for a member to idle through its delay, and that is allowed. for _ in 0..30 { exec.reborrow().broadcast(|_| {}).unwrap(); std::thread::sleep(Duration::from_millis(5)); } - assert!( - kinds(&team, Kind::ParkEnter).is_empty(), - "a member parked although sections arrived every 5 ms" + // Then no section at all: every member parks, however long it takes. + wait_until("every member parked after its own last share", || { + every_member_parked_after_its_last_share(&member_events(&team)) + }); + let members = members_that_ran(&member_events(&team)); + assert_eq!( + members.len(), + team.size().get() - 1, + "every member ran a share: {members:?}" ); - // Then no section: every member parks once, not before the delay. - let last_end = kinds(&team, Kind::SectionEnd) - .iter() - .map(|&(_, ns)| ns) - .max() - .unwrap(); - std::thread::sleep(Duration::from_millis(300)); - let parks = kinds(&team, Kind::ParkEnter); + // The "not before" property holds for every park of the whole test, the + // ones inside the phase of regular sections included. + assert_no_park_before_the_delay(&team); + // A following section wakes every member that is parked now: the wake-ups + // are exactly the parked members' `ParkExit`s. + let parked = parked_now(&team); assert_eq!( - parks.len(), - 2, - "each of the two members parks once: {parks:?}" + parked, + team.size().get() - 1, + "every member is parked after the last section" ); - for (m, ns) in parks { - let after = ns.saturating_sub(last_end); - let own_end = kinds(&team, Kind::SectionEnd) - .iter() - .filter(|&&(mm, _)| mm == m) - .map(|&(_, n)| n) - .max() - .unwrap(); - assert!( - ns - own_end >= SPIN.as_nanos() as u64, - "member {m} parked {} us after its own last share", - (ns - own_end) / 1000 - ); - assert!(after < 250_000_000, "member {m} parked too late"); - } - // A section wakes them; the clock restarts from that share. + let exits = count(&team, Kind::ParkExit); exec.reborrow().broadcast(|_| {}).unwrap(); - assert_eq!(kinds(&team, Kind::ParkExit).len(), 2); + wait_until("every parked member woke", || { + count(&team, Kind::ParkExit) >= exits + parked + }); + // The clock restarts from the share they just ran: a later park is again at + // least SPIN after it. + assert_no_park_before_the_delay(&team); } diff --git a/tests/oversubscribed.rs b/tests/oversubscribed.rs new file mode 100644 index 0000000..d9bd759 --- /dev/null +++ b/tests/oversubscribed.rs @@ -0,0 +1,212 @@ +//! A root team scope that has more participants than CPUs in the caller's set +//! parks its members after a short spin instead of spinning for the whole park +//! delay, and its nested sections keep that knowledge. +//! +//! Nothing here is a wall-clock bound on the members: the park delay of the +//! policy is set far above the deadline of the hang guard, so *reaching* the +//! park proves the scope used its own (zero) delay, whatever the host's speed, +//! and the wait itself is a poll of the trace with no upper bound. +#![cfg(all(feature = "trace", feature = "parallel", target_os = "linux"))] + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; +use threadteam::trace::{self, Kind}; +use threadteam::Pool; + +/// The policy's park delay. Far above `DEADLINE`: a member that waited for it +/// could not park within the deadline. +const PARK_DELAY: Duration = Duration::from_secs(60); + +/// Only a hang guard. +const DEADLINE: Duration = Duration::from_secs(10); + +/// A hang is a failure: abort the whole process if the test body stalls. +fn watchdog() -> impl Drop { + struct Stop(std::sync::mpsc::Sender<()>); + impl Drop for Stop { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + let (tx, rx) = std::sync::mpsc::channel::<()>(); + std::thread::spawn(move || { + if rx.recv_timeout(Duration::from_secs(120)).is_err() { + eprintln!("watchdog: test stalled"); + std::process::exit(101); + } + }); + Stop(tx) +} + +fn own_cpus() -> Vec { + // SAFETY: reads the calling thread's mask into a zeroed `cpu_set_t`. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + let r = libc::sched_getaffinity(0, std::mem::size_of::(), &mut set); + assert_eq!(r, 0, "sched_getaffinity failed"); + (0..libc::CPU_SETSIZE as usize) + .filter(|&cpu| libc::CPU_ISSET(cpu, &set)) + .collect() + } +} + +fn set_own_cpus(cpus: &[usize]) -> bool { + // SAFETY: builds a `cpu_set_t` and applies it to pid 0, i.e. this thread. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + libc::CPU_ZERO(&mut set); + for &cpu in cpus { + libc::CPU_SET(cpu, &mut set); + } + libc::sched_setaffinity(0, std::mem::size_of::(), &set) == 0 + } +} + +/// Restores the calling thread's mask on drop. +struct Restrict(Vec); + +impl Drop for Restrict { + fn drop(&mut self) { + set_own_cpus(&self.0); + } +} + +/// The members that ran a share, in the order of their first share. Their slot +/// ids are whatever the acquisition claimed, so they are read from the trace. +fn members_that_ran(events: &[trace::Event]) -> Vec { + let mut out: Vec = Vec::new(); + for e in events { + if e.kind == Kind::SectionEnd && !out.contains(&e.member) { + out.push(e.member); + } + } + out +} + +/// Polls (hang guard only) until every member that ran a share has a +/// `ParkEnter` later than its own last `SectionEnd`. +/// +/// No upper bound on how late: what the deadline rules out is the opposite — +/// with `PARK_DELAY` = 60 s, a member that kept the policy's delay cannot park +/// before it, so parking within `DEADLINE` can only come from the scope's own +/// (zero) delay. +fn wait_until_members_parked(team: &threadteam::Team) { + let start = Instant::now(); + loop { + let events = trace::events(team.pool()); + let members = members_that_ran(&events); + let parked = !members.is_empty() + && members.iter().all(|&member| { + let last_end = events + .iter() + .filter(|e| e.member == member && e.kind == Kind::SectionEnd) + .map(|e| e.ns) + .max(); + let park = events + .iter() + .filter(|e| e.member == member && e.kind == Kind::ParkEnter) + .map(|e| e.ns) + .max(); + matches!((last_end, park), (Some(end), Some(park)) if park > end) + }); + if parked { + assert_eq!( + members.len(), + team.size().get() - 1, + "every member ran a share" + ); + return; + } + assert!( + start.elapsed() < DEADLINE, + "no member parked within {DEADLINE:?}: the scope kept the policy's {PARK_DELAY:?} delay" + ); + std::thread::sleep(Duration::from_millis(2)); + } +} + +#[test] +fn an_oversubscribed_root_scope_parks_its_members_quickly() { + let _w = watchdog(); + let all = own_cpus(); + if all.len() < 3 { + eprintln!("skipped: need at least 3 allowed CPUs to restrict to 2"); + return; + } + let cpus = vec![all[0], all[1]]; + if !set_own_cpus(&cpus) { + eprintln!("skipped: cannot restrict the calling thread to {cpus:?}"); + return; + } + let _restore = Restrict(all); + + let pool = Pool::builder() + .threads(4) + .idle_spin(PARK_DELAY) + .acquire_timeout(Duration::from_secs(10)) + .build() + .unwrap(); + let team = pool.into_team(4).expect("an idle pool"); + let width = team.size().get(); + assert!(width > cpus.len(), "4 participants on 2 CPUs"); + + let sum = AtomicUsize::new(0); + let exec = team.exec(width); + exec.broadcast(|i| { + sum.fetch_add(i + 1, Ordering::Relaxed); + }) + .expect("a team token broadcasts"); + assert_eq!(sum.load(Ordering::Relaxed), (1..=width).sum::()); + // The root scope is oversubscribed, so its members park without waiting for + // the 60 s park delay. + wait_until_members_parked(&team); + + // Nested sections of the same root keep that knowledge: a split, then a + // section in each child, and the members park again. + let mut exec = team.exec(width); + exec.team(|team| { + team.split(&[2, 2], |_, child| child.run(|_, _| {})) + .unwrap(); + }) + .expect("a team token opens its scope"); + wait_until_members_parked(&team); +} + +#[test] +fn the_scope_that_enters_decides_the_park_delay_not_the_mint() { + let _w = watchdog(); + let all = own_cpus(); + if all.len() < 2 { + eprintln!("skipped: need at least 2 allowed CPUs to restrict to 1"); + return; + } + // Mint the root token of a persistent team on the full mask: three + // participants on all CPUs are not oversubscribed, so the members start + // with the policy's 60 s park delay. + let pool = Pool::builder() + .threads(3) + .idle_spin(PARK_DELAY) + .acquire_timeout(Duration::from_secs(10)) + .build() + .unwrap(); + let team = pool.into_team(3).expect("an idle pool"); + let width = team.size().get(); + let mut exec = team.exec(width); + // Enter the scope from a one-CPU mask: it is oversubscribed there, so its + // members must park after the short spin, not after 60 s. + let cpus = vec![all[0]]; + if !set_own_cpus(&cpus) { + eprintln!("skipped: cannot restrict the calling thread to {cpus:?}"); + return; + } + let _restore = Restrict(all); + let sum = AtomicUsize::new(0); + exec.team(|mut team| { + team.reborrow().run(|i, _| { + sum.fetch_add(i + 1, Ordering::Relaxed); + }); + }) + .expect("a team token opens its scope"); + assert_eq!(sum.load(Ordering::Relaxed), (1..=width).sum::()); + wait_until_members_parked(&team); +} diff --git a/tests/pinning.rs b/tests/pinning.rs new file mode 100644 index 0000000..60918d4 --- /dev/null +++ b/tests/pinning.rs @@ -0,0 +1,484 @@ +//! Opt-in member placement (Linux only): every member pins itself onto one CPU +//! of the calling thread's set for a root team scope and restores its own mask +//! afterwards, while the leader (the calling thread) is never moved. +#![cfg(all(feature = "parallel", target_os = "linux"))] + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use threadteam::{Exec, Pool}; + +/// A hang is a failure: abort the whole process if the test body stalls. +fn watchdog() -> impl Drop { + struct Stop(std::sync::mpsc::Sender<()>); + impl Drop for Stop { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + let (tx, rx) = std::sync::mpsc::channel::<()>(); + std::thread::spawn(move || { + if rx.recv_timeout(Duration::from_secs(45)).is_err() { + eprintln!("watchdog: test stalled"); + std::process::exit(101); + } + }); + Stop(tx) +} + +/// The set CPUs of the calling thread's affinity mask, ascending. +fn own_cpus() -> Vec { + // SAFETY: reads the calling thread's mask into a zeroed `cpu_set_t` of the + // right size; pid 0 is this thread. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + let r = libc::sched_getaffinity(0, std::mem::size_of::(), &mut set); + assert_eq!(r, 0, "sched_getaffinity failed"); + (0..libc::CPU_SETSIZE as usize) + .filter(|&cpu| libc::CPU_ISSET(cpu, &set)) + .collect() + } +} + +fn set_own_cpus(cpus: &[usize]) -> bool { + // SAFETY: builds a `cpu_set_t` and applies it to pid 0, i.e. this thread. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + libc::CPU_ZERO(&mut set); + for &cpu in cpus { + libc::CPU_SET(cpu, &mut set); + } + libc::sched_setaffinity(0, std::mem::size_of::(), &set) == 0 + } +} + +/// Restricts the calling thread to `cpus`; restores the previous mask on drop. +struct Restrict(Vec); + +impl Restrict { + /// `None` when the restriction is refused (the test then skips). + fn to(cpus: &[usize]) -> Option { + let before = own_cpus(); + set_own_cpus(cpus).then_some(Restrict(before)) + } +} + +impl Drop for Restrict { + fn drop(&mut self) { + set_own_cpus(&self.0); + } +} + +fn builder(width: usize, pin: bool) -> threadteam::PoolBuilder { + Pool::builder() + .threads(width) + .pin_participants(pin) + .idle_spin(Duration::from_millis(1)) + .acquire_timeout(Duration::from_secs(10)) +} + +/// Worker index -> the CPUs it started with. +type Masks = Arc)>>>; + +/// The masks the workers had when they started. +#[derive(Clone, Default)] +struct Origin(Masks); + +impl Origin { + fn record(&self, worker: usize, cpus: Vec) { + self.0.lock().unwrap().push((worker, cpus)); + } + + fn of(&self, worker: usize) -> Vec { + let table = self.0.lock().unwrap(); + let row = table.iter().find(|(i, _)| *i == worker); + assert!(row.is_some(), "no mask recorded for worker {worker}"); + row.unwrap().1.clone() + } + + fn all(&self) -> Vec> { + self.0 + .lock() + .unwrap() + .iter() + .map(|(_, m)| m.clone()) + .collect() + } +} + +/// A pool whose workers report their initial masks, and the table. +fn origin_pool(width: usize, pin: bool) -> (Pool, Origin) { + let origin = Origin::default(); + let store = origin.clone(); + let pool = builder(width, pin) + .on_start(move |worker| store.record(worker, own_cpus())) + .build() + .unwrap(); + (pool, origin) +} + +/// Collects the mask each participant sees while it runs a share. +#[derive(Default)] +struct Seen { + masks: Mutex)>>, +} + +impl Seen { + fn record(&self, index: usize) { + self.masks.lock().unwrap().push((index, own_cpus())); + } + + fn take(&self) -> Vec<(usize, Vec)> { + let mut v = self.masks.lock().unwrap().clone(); + v.sort_by_key(|&(i, _)| i); + v + } +} + +#[test] +fn a_root_team_scope_places_each_member_on_its_own_cpu_and_never_moves_the_leader() { + let _w = watchdog(); + let set = own_cpus(); + if set.len() < 2 { + eprintln!("skipped: the process has fewer than 2 allowed CPUs"); + return; + } + let width = set.len().min(4); + let (pool, origin) = origin_pool(width, true); + let before = own_cpus(); + let seen = Seen::default(); + let mut exec = Exec::new(&pool, width).unwrap(); + exec.team(|mut team| { + assert_eq!(team.size().get(), width); + team.reborrow().run_with(|p| seen.record(p.index())); + }) + .expect("an idle pool"); + + let masks = seen.take(); + assert_eq!(masks.len(), width, "every participant ran a share"); + // Participant 0 is the leader: it runs its share on the calling thread and + // its mask is untouched, inside the scope included. + assert_eq!( + masks[0].1, before, + "the leader was moved: {:?} instead of {before:?}", + masks[0].1 + ); + for (index, mask) in &masks[1..] { + // Member k sits alone on the k-th CPU of the caller's set. + assert_eq!( + mask, + &vec![set[*index]], + "participant {index} of {width} on the caller's set {set:?}" + ); + } + let distinct: std::collections::HashSet<_> = masks[1..].iter().map(|(_, m)| m[0]).collect(); + assert_eq!(distinct.len(), width - 1, "members are on distinct CPUs"); + // And the calling thread's mask is the same after the scope. + assert_eq!(own_cpus(), before, "the leader's mask after the scope"); + + // A member restores its own mask before it returns its thread to Rayon: + // every worker of the pool shows its original mask again. + pool.release_idle(); + Exec::new(&pool, width) + .unwrap() + .broadcast(|worker| { + assert_eq!( + own_cpus(), + origin.of(worker), + "worker {worker} kept a pinned mask after the scope" + ); + }) + .expect("a root token from a plain thread"); +} + +#[test] +fn a_persistent_team_pins_every_root_scope_and_never_moves_the_leader() { + let _w = watchdog(); + let set = own_cpus(); + if set.len() < 2 { + eprintln!("skipped: the process has fewer than 2 allowed CPUs"); + return; + } + let (pool, origin) = origin_pool(set.len().min(4), true); + let team = pool.into_team(set.len().min(4)).expect("an idle pool"); + let width = team.size().get(); + let seen = Seen::default(); + let before = own_cpus(); + for round in 0..2 { + let mut exec = team.exec(width); + // The faer `with_team` path: a root token of the team, used through + // `team()`, twice on the same persistent reservation. + exec.team(|mut team| { + team.reborrow().run_with(|p| seen.record(p.index())); + }) + .expect("a team token opens its scope"); + drop(exec); + assert_eq!( + own_cpus(), + before, + "round {round}: the leader's mask after the scope" + ); + } + let masks = seen.take(); + assert_eq!(masks.len(), 2 * width); + let mut rounds = vec![0usize; width]; + for (index, mask) in &masks { + let round = rounds[*index]; + rounds[*index] += 1; + if *index == 0 { + // Participant 0 is the leader, which is never pinned. + assert_eq!(mask, &before, "round {round}: the leader was moved"); + } else { + assert_eq!( + mask, + &vec![set[*index]], + "round {round}, participant {index}" + ); + } + } + assert!( + rounds.iter().all(|&n| n == 2), + "both rounds ran: {rounds:?}" + ); + // The members restore what they had before the first scope when the + // reservation ends. + let pool = team.into_pool(); + Exec::new(&pool, width) + .unwrap() + .broadcast(|worker| { + assert_eq!( + own_cpus(), + origin.of(worker), + "worker {worker} after the team" + ); + }) + .expect("a root token from a plain thread"); +} + +#[test] +fn pinning_off_is_the_default_and_leaves_every_mask_alone() { + let _w = watchdog(); + let set = own_cpus(); + let width = set.len().min(3); + let (pool, origin) = origin_pool(width, false); + assert!(!pool.pin_participants(), "off by default"); + let before = own_cpus(); + let seen = Seen::default(); + let mut exec = Exec::new(&pool, width).unwrap(); + exec.team(|mut team| { + team.reborrow().run_with(|p| seen.record(p.index())); + }) + .expect("an idle pool"); + let masks = seen.take(); + assert_eq!(masks.len(), width); + for (index, mask) in &masks { + // Nothing moves: every thread keeps the mask it started with. + assert_eq!(mask, &before, "participant {index} was moved"); + } + assert_eq!(own_cpus(), before); + assert!( + origin.all().iter().all(|m| m == &before), + "a worker started with a different mask than the caller" + ); +} + +#[test] +fn a_lent_member_restores_its_mask_before_it_serves_rayon() { + let _w = watchdog(); + let set = own_cpus(); + if set.len() < 2 { + eprintln!("skipped: the process has fewer than 2 allowed CPUs"); + return; + } + let width = set.len().min(4); + let (pool, origin) = origin_pool(width, true); + let team = pool.into_team(width).expect("an idle pool"); + let width = team.size().get(); + let mut exec = team.exec(width); + // Pin the members for the root scope; a persistent team keeps them pinned + // until they are released or lent. + exec.team(|_| {}).expect("a team token opens its scope"); + // `external` lends the members to Rayon: the closure runs on one of them + // and the jobs it forks run on the others. + let inside = AtomicUsize::new(0); + let job_masks = Mutex::new(Vec::new()); + let lent_caller = exec.external(|width| { + let here = own_cpus(); + Exec::new(team.pool(), width.get()) + .unwrap() + .fork(width.get(), |_, _| { + inside.fetch_add(1, Ordering::Relaxed); + job_masks.lock().unwrap().push(own_cpus()); + }); + here + }); + assert_eq!( + lent_caller, + own_cpus(), + "the thread that ran the lent call kept a pin" + ); + assert_eq!(inside.load(Ordering::Relaxed), width); + for mask in job_masks.into_inner().unwrap() { + // A member that had not restored would show one CPU, not the set. + assert!( + mask.len() > 1, + "a lent job ran on a pinned thread: {mask:?}" + ); + assert!( + origin.all().contains(&mask), + "a lent job ran with an unexpected mask {mask:?}" + ); + } +} + +#[test] +fn a_team_wider_than_the_callers_set_wraps_around_it() { + let _w = watchdog(); + let all = own_cpus(); + if all.len() < 3 { + eprintln!("skipped: need at least 3 allowed CPUs to restrict to 2"); + return; + } + let cpus = vec![all[0], all[1]]; + let Some(_restore) = Restrict::to(&cpus) else { + eprintln!("skipped: cannot restrict the calling thread to {cpus:?}"); + return; + }; + let width = 4; + let pool = builder(width, true).build().unwrap(); + let before = own_cpus(); + assert_eq!(before, cpus); + let seen = Seen::default(); + let mut exec = Exec::new(&pool, width).unwrap(); + exec.team(|mut team| { + assert_eq!(team.size().get(), width); + team.reborrow().run_with(|p| seen.record(p.index())); + }) + .expect("an idle pool"); + let masks = seen.take(); + assert_eq!(masks.len(), width); + // The leader keeps the whole set; member k wraps onto S[k % 2]. + assert_eq!(masks[0].1, before, "the leader was moved"); + for (index, mask) in &masks[1..] { + assert_eq!( + mask, + &vec![cpus[index % cpus.len()]], + "participant {index} of {width} on the caller's set {cpus:?}" + ); + } + assert_eq!(own_cpus(), before, "the leader's mask after the scope"); +} + +#[test] +fn a_member_leading_a_fork_child_restores_its_mask_before_external() { + let _w = watchdog(); + let set = own_cpus(); + // The member's child needs a budget of two to enter the pool, so the team + // must fork into shares of two. + if set.len() < 4 { + eprintln!("skipped: a member-led child of width 2 needs 4 allowed CPUs"); + return; + } + let width = 4; + let (pool, origin) = origin_pool(width, true); + let team = pool.into_team(width).expect("an idle pool"); + let mut exec = team.exec(width); + // Pin the members for the root scope; a persistent team keeps them pinned. + exec.team(|_| {}).expect("a team token opens its scope"); + let pre_pin = own_cpus(); + let seen = Mutex::new(Vec::new()); + exec.fork(2, |child, child_token| { + if child == 1 { + // Child 1 is led by a member, which the scope pinned to one CPU. + // `external` runs inline on that thread (it is a worker of the + // pool), so it must have restored its own mask first. + let mask = child_token.external(|_width| own_cpus()); + seen.lock().unwrap().push(mask); + } + }); + let seen = seen.into_inner().unwrap(); + assert_eq!(seen.len(), 1, "the member-led child ran `external`"); + assert_eq!( + seen[0], pre_pin, + "a pinned member kept its mask while serving foreign work" + ); + assert!( + origin.all().contains(&pre_pin), + "the member restored what it had before the first pin" + ); +} + +#[test] +fn a_panic_in_a_share_leaves_the_callers_affinity_alone() { + use std::panic::{catch_unwind, AssertUnwindSafe}; + let _w = watchdog(); + let set = own_cpus(); + if set.len() < 2 { + eprintln!("skipped: the process has fewer than 2 allowed CPUs"); + return; + } + let width = set.len().min(3); + + // A scoped acquisition: the leader is never pinned, so unwinding through + // the scope cannot leave it on another mask either. + let (pool, origin) = origin_pool(width, true); + let before = own_cpus(); + let mut exec = Exec::new(&pool, width).unwrap(); + let panicked = catch_unwind(AssertUnwindSafe(|| { + let _ = exec.team(|mut team| { + team.reborrow().run(|index, _| { + assert!(index < width); + if index == 1 { + panic!("a member panics in its share"); + } + }); + }); + })); + assert!(panicked.is_err(), "the panic reached the leader"); + assert_eq!(own_cpus(), before, "the caller after an unwinding scope"); + pool.release_idle(); + Exec::new(&pool, width) + .unwrap() + .broadcast(|worker| { + assert_eq!( + own_cpus(), + origin.of(worker), + "worker {worker} after unwind" + ); + }) + .expect("a root token from a plain thread"); + + // A persistent team's root token: same, and the members restore what they + // had before the first pin when the reservation ends. + let (pool, origin) = origin_pool(width, true); + let team = pool.into_team(width).expect("an idle pool"); + let before = own_cpus(); + let mut exec = team.exec(width); + let panicked = catch_unwind(AssertUnwindSafe(|| { + let _ = exec.team(|mut team| { + team.reborrow().run(|index, _| { + if index == 1 { + panic!("a member panics in its share"); + } + }); + }); + })); + assert!(panicked.is_err(), "the panic reached the leader"); + drop(exec); + assert_eq!( + own_cpus(), + before, + "the caller after an unwinding team scope" + ); + let pool = team.into_pool(); + Exec::new(&pool, width) + .unwrap() + .broadcast(|worker| { + assert_eq!( + own_cpus(), + origin.of(worker), + "worker {worker} after the team" + ); + }) + .expect("a root token from a plain thread"); +} diff --git a/tests/team_alloc_free.rs b/tests/team_alloc_free.rs index d519d92..32e0ee5 100644 --- a/tests/team_alloc_free.rs +++ b/tests/team_alloc_free.rs @@ -137,4 +137,41 @@ fn team_dispatch_does_not_allocate_after_warm_up() { "after park, round {round}" ); } + + // Participant pinning on: minting a root token of a persistent team, running + // a root team scope through it (which is where the pin section places the + // participants) and dropping the token must not allocate. The leader saves + // and restores its mask on its stack, a member saves its mask in its own + // slot, and both use `cpu_set_t` values, never a `Vec`. + let pool = Pool::builder() + .threads(4) + .pin_participants(true) + .idle_spin(Duration::from_millis(1)) + .acquire_timeout(Duration::from_secs(5)) + .build() + .unwrap(); + let team = pool.into_team(4).unwrap(); + let token_phases = |team: &threadteam::Team| { + let mut exec = team.exec(4); // mint + exec.team(|mut team| { + team.reborrow().run(|_, _| {}); + team.reborrow().run_with(|p| { + let _ = p.barrier(); + }); + }) + .expect("the team token opens its scope"); + // pinned members dispatch as usual + exec.reborrow().fork(2, |_, mut child| { + child + .team(|team| team.run(|_, _| {})) + .expect("carried sub-team"); + }); + drop(exec); // drop + }; + for _ in 0..50 { + token_phases(&team); // warm-up + } + for round in 0..20 { + assert_eq!(allocs_in(|| token_phases(&team)), 0, "pinned round {round}"); + } }