fix(rayon): remove the aliasing unsafe from the parallel sweep

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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DnsaJg74eNSva3PJjK2eej
This commit is contained in:
2026-08-04 21:54:34 +02:00
co-authored by Claude Opus 5
parent c088214fed
commit 06b6a68499
2 changed files with 68 additions and 37 deletions
+2
View File
@@ -1,3 +1,5 @@
#![forbid(unsafe_code)]
use std::{
cmp::Reverse,
f64::consts::{FRAC_1_SQRT_2, FRAC_2_SQRT_PI, SQRT_2},
+67 -38
View File
@@ -124,18 +124,19 @@ impl Event {
.collect::<Vec<_>>()
}
/// 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<T: Time, D: Drift<T>>(
&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<T: Time, D: Drift<T>>(
&self,
skills: &SkillStore,
agents: &CompetitorStore<T, D>,
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,19 +153,49 @@ impl Event {
),
};
for (t, team) in self.teams.iter_mut().enumerate() {
for (i, item) in team.items.iter_mut().enumerate() {
let old_likelihood = skills.get(item.agent).unwrap().likelihood;
let new_likelihood = (old_likelihood / item.likelihood) * g.likelihoods[t][i];
skills.get_mut(item.agent).unwrap().likelihood = new_likelihood;
item.likelihood = g.likelihoods[t][i];
EventUpdate {
log_evidence: g.log_evidence,
likelihoods: g.likelihoods,
}
}
self.log_evidence = g.log_evidence;
/// 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) * fresh;
skills.get_mut(item.agent).unwrap().likelihood = new_likelihood;
item.likelihood = fresh;
}
}
self.log_evidence = update.log_evidence;
}
/// Compute and apply in one step — the sequential sweep.
fn iteration_direct<T: Time, D: Drift<T>>(
&mut self,
skills: &mut SkillStore,
agents: &CompetitorStore<T, D>,
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<Vec<Gaussian>>,
}
#[derive(Debug)]
pub struct TimeSlice<T: Time = i64> {
pub(crate) events: Vec<Event>,
@@ -384,14 +415,13 @@ impl<T: Time> TimeSlice<T> {
/// 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<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
use rayon::prelude::*;
@@ -411,29 +441,28 @@ impl<T: Time> TimeSlice<T> {
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) };
let skills = &self.skills;
let updates: Vec<EventUpdate> = self.events[range.clone()]
.par_iter()
.map(|ev| {
ARENA.with(|cell| {
let mut arena = cell.borrow_mut();
arena.reset();
ev.iteration_direct(skills, agents, p_draw, convergence, &mut arena);
});
});
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(