From 06b6a6849914e1975021d2f264580f88e6397d94 Mon Sep 17 00:00:00 2001 From: Anders Olsson Date: Tue, 4 Aug 2026 21:54:34 +0200 Subject: [PATCH] fix(rayon): remove the aliasing unsafe from the parallel sweep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parallel color-group sweep passed a `*mut SkillStore` through a `usize` and cast it back inside the rayon closure, so every worker materialised its own `&mut SkillStore` to the same store. Two live `&mut` to one object is an aliasing violation whatever the workers subsequently touch — `&mut` carries `noalias` down to LLVM — and laundering the pointer through `usize` also discarded provenance. The existing SAFETY comment argued element disjointness, which is true and is why nothing miscompiled in practice, but it is not the property the aliasing rules ask about. Events in a color group touch disjoint agents, so none can observe another's writes. That makes the sweep separable rather than merely safe-in-practice: `Event::compute` runs inference over shared `&self.skills` with no mutation, and `Event::apply` folds the results in afterwards in index order. No `unsafe`, no aliasing argument, and the apply order does not depend on which worker finished first, so results stay bit-identical across thread counts. The crate now contains no `unsafe` at all, locked in with `#![forbid(unsafe_code)]`. Splitting compute from apply also removes the duplicated sweep body: the `from > 0` branch of `TimeSlice::iteration` was a verbatim copy of `iteration_direct`, and both now share one implementation. Cost, measured on the three `history_converge` workloads (sequential vs parallel, this machine): 500x100@10perslice 4.02ms -> 4.21ms 2000x200@20perslice 19.70ms -> 19.76ms 1v1-5000x50000 11.75ms -> 10.46ms The deferred apply gives back part of the parallel win on the only workload where rayon ever helped (1.12x here, against the 1.3x T3 reported), and the sequential path is unchanged. Trading a fraction of a 1.3x speedup on one pathological shape for the removal of undefined behaviour is the right side of that bargain. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01DnsaJg74eNSva3PJjK2eej --- src/lib.rs | 2 + src/time_slice.rs | 103 +++++++++++++++++++++++++++++----------------- 2 files changed, 68 insertions(+), 37 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 227c4f1..ae109f5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![forbid(unsafe_code)] + use std::{ cmp::Reverse, f64::consts::{FRAC_1_SQRT_2, FRAC_2_SQRT_PI, SQRT_2}, diff --git a/src/time_slice.rs b/src/time_slice.rs index 0f001b1..b762a81 100644 --- a/src/time_slice.rs +++ b/src/time_slice.rs @@ -124,18 +124,19 @@ impl Event { .collect::>() } - /// Direct in-loop update: mutates self and `skills` inline with no - /// intermediate allocation. Used by both the sequential sweep path and, - /// via unsafe, by the parallel rayon path for events in the same color - /// group (which have disjoint agent sets — see `sweep_color_groups`). - fn iteration_direct>( - &mut self, - skills: &mut SkillStore, + /// Run inference for this event and return its per-item likelihoods. + /// + /// Reads `skills` immutably and does not touch `self`, so every event in + /// a color group can run concurrently without any aliasing question — + /// the mutation is deferred to `apply`. + fn compute>( + &self, + skills: &SkillStore, agents: &CompetitorStore, p_draw: f64, convergence: crate::ConvergenceOptions, arena: &mut ScratchArena, - ) { + ) -> EventUpdate { let teams = self.within_priors(false, false, skills, agents); let result = self.outputs(); let g = match self.kind { @@ -152,17 +153,47 @@ impl Event { ), }; + EventUpdate { + log_evidence: g.log_evidence, + likelihoods: g.likelihoods, + } + } + + /// Fold a computed update into the skill store and cache it on the items. + fn apply(&mut self, skills: &mut SkillStore, update: EventUpdate) { for (t, team) in self.teams.iter_mut().enumerate() { for (i, item) in team.items.iter_mut().enumerate() { + let fresh = update.likelihoods[t][i]; let old_likelihood = skills.get(item.agent).unwrap().likelihood; - let new_likelihood = (old_likelihood / item.likelihood) * g.likelihoods[t][i]; + let new_likelihood = (old_likelihood / item.likelihood) * fresh; skills.get_mut(item.agent).unwrap().likelihood = new_likelihood; - item.likelihood = g.likelihoods[t][i]; + item.likelihood = fresh; } } - self.log_evidence = g.log_evidence; + self.log_evidence = update.log_evidence; } + + /// Compute and apply in one step — the sequential sweep. + fn iteration_direct>( + &mut self, + skills: &mut SkillStore, + agents: &CompetitorStore, + p_draw: f64, + convergence: crate::ConvergenceOptions, + arena: &mut ScratchArena, + ) { + let update = self.compute(skills, agents, p_draw, convergence, arena); + self.apply(skills, update); + } +} + +/// The result of running inference for one event, before it is folded back +/// into the shared skill store. +#[derive(Debug)] +struct EventUpdate { + log_evidence: f64, + likelihoods: Vec>, } #[derive(Debug)] @@ -384,14 +415,13 @@ impl TimeSlice { /// Full event sweep using the color-group partition. Colors are processed /// sequentially; within each color the inner loop is parallel under rayon. /// - /// Events within each color group touch disjoint agent sets (guaranteed by - /// the greedy coloring). This lets each rayon thread write directly to its - /// events' skill likelihoods without a deferred-apply step, matching the - /// sequential path's allocation profile. The unsafe block is sound because: - /// 1. `self.events[range]` and `self.skills` are separate fields → disjoint. - /// 2. Events in the same color group access disjoint `Index` values in - /// `self.skills`, so concurrent writes land on different memory locations. - /// 3. Each event only writes to its own items' likelihoods (no sharing). + /// Events in one color group touch disjoint agent sets, so none of them + /// can observe another's writes. That makes the sweep separable: inference + /// runs concurrently over shared `&self.skills`, and the resulting updates + /// are folded in afterwards in index order. Splitting it this way needs no + /// `unsafe` and no aliasing argument, and it keeps results bit-identical + /// across thread counts because the apply order does not depend on which + /// worker finished first. #[cfg(feature = "rayon")] fn sweep_color_groups>(&mut self, agents: &CompetitorStore) { use rayon::prelude::*; @@ -411,29 +441,28 @@ impl TimeSlice { if group_len == 0 { continue; } + let range = self.color_groups.color_range(color_idx); let p_draw = self.p_draw; let convergence = self.convergence; if group_len >= RAYON_THRESHOLD { - // Obtain a raw pointer from the unique `&mut self.skills` reference. - // Casting back to `&mut` inside the closure is sound because: - // 1. The pointer originates from a `&mut` — no aliasing with shared refs. - // 2. Events in the same color group touch disjoint `Index` slots in the - // underlying Vec, so concurrent writes from different threads land on - // different memory locations — no data race. - // 3. `self.events[range]` and `self.skills` are separate struct fields, - // so the borrow splits cleanly. - let skills_addr: usize = (&mut self.skills as *mut SkillStore) as usize; - self.events[range].par_iter_mut().for_each(move |ev| { - // SAFETY: see above. - let skills: &mut SkillStore = unsafe { &mut *(skills_addr as *mut SkillStore) }; - ARENA.with(|cell| { - let mut arena = cell.borrow_mut(); - arena.reset(); - ev.iteration_direct(skills, agents, p_draw, convergence, &mut arena); - }); - }); + let skills = &self.skills; + let updates: Vec = self.events[range.clone()] + .par_iter() + .map(|ev| { + ARENA.with(|cell| { + let mut arena = cell.borrow_mut(); + arena.reset(); + + ev.compute(skills, agents, p_draw, convergence, &mut arena) + }) + }) + .collect(); + + for (ev, update) in self.events[range].iter_mut().zip(updates) { + ev.apply(&mut self.skills, update); + } } else { for ev in &mut self.events[range] { ev.iteration_direct(