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:
@@ -1,3 +1,5 @@
|
|||||||
|
#![forbid(unsafe_code)]
|
||||||
|
|
||||||
use std::{
|
use std::{
|
||||||
cmp::Reverse,
|
cmp::Reverse,
|
||||||
f64::consts::{FRAC_1_SQRT_2, FRAC_2_SQRT_PI, SQRT_2},
|
f64::consts::{FRAC_1_SQRT_2, FRAC_2_SQRT_PI, SQRT_2},
|
||||||
|
|||||||
+67
-38
@@ -124,18 +124,19 @@ impl Event {
|
|||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Direct in-loop update: mutates self and `skills` inline with no
|
/// Run inference for this event and return its per-item likelihoods.
|
||||||
/// intermediate allocation. Used by both the sequential sweep path and,
|
///
|
||||||
/// via unsafe, by the parallel rayon path for events in the same color
|
/// Reads `skills` immutably and does not touch `self`, so every event in
|
||||||
/// group (which have disjoint agent sets — see `sweep_color_groups`).
|
/// a color group can run concurrently without any aliasing question —
|
||||||
fn iteration_direct<T: Time, D: Drift<T>>(
|
/// the mutation is deferred to `apply`.
|
||||||
&mut self,
|
fn compute<T: Time, D: Drift<T>>(
|
||||||
skills: &mut SkillStore,
|
&self,
|
||||||
|
skills: &SkillStore,
|
||||||
agents: &CompetitorStore<T, D>,
|
agents: &CompetitorStore<T, D>,
|
||||||
p_draw: f64,
|
p_draw: f64,
|
||||||
convergence: crate::ConvergenceOptions,
|
convergence: crate::ConvergenceOptions,
|
||||||
arena: &mut ScratchArena,
|
arena: &mut ScratchArena,
|
||||||
) {
|
) -> EventUpdate {
|
||||||
let teams = self.within_priors(false, false, skills, agents);
|
let teams = self.within_priors(false, false, skills, agents);
|
||||||
let result = self.outputs();
|
let result = self.outputs();
|
||||||
let g = match self.kind {
|
let g = match self.kind {
|
||||||
@@ -152,19 +153,49 @@ impl Event {
|
|||||||
),
|
),
|
||||||
};
|
};
|
||||||
|
|
||||||
for (t, team) in self.teams.iter_mut().enumerate() {
|
EventUpdate {
|
||||||
for (i, item) in team.items.iter_mut().enumerate() {
|
log_evidence: g.log_evidence,
|
||||||
let old_likelihood = skills.get(item.agent).unwrap().likelihood;
|
likelihoods: g.likelihoods,
|
||||||
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];
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
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)]
|
#[derive(Debug)]
|
||||||
pub struct TimeSlice<T: Time = i64> {
|
pub struct TimeSlice<T: Time = i64> {
|
||||||
pub(crate) events: Vec<Event>,
|
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
|
/// Full event sweep using the color-group partition. Colors are processed
|
||||||
/// sequentially; within each color the inner loop is parallel under rayon.
|
/// sequentially; within each color the inner loop is parallel under rayon.
|
||||||
///
|
///
|
||||||
/// Events within each color group touch disjoint agent sets (guaranteed by
|
/// Events in one color group touch disjoint agent sets, so none of them
|
||||||
/// the greedy coloring). This lets each rayon thread write directly to its
|
/// can observe another's writes. That makes the sweep separable: inference
|
||||||
/// events' skill likelihoods without a deferred-apply step, matching the
|
/// runs concurrently over shared `&self.skills`, and the resulting updates
|
||||||
/// sequential path's allocation profile. The unsafe block is sound because:
|
/// are folded in afterwards in index order. Splitting it this way needs no
|
||||||
/// 1. `self.events[range]` and `self.skills` are separate fields → disjoint.
|
/// `unsafe` and no aliasing argument, and it keeps results bit-identical
|
||||||
/// 2. Events in the same color group access disjoint `Index` values in
|
/// across thread counts because the apply order does not depend on which
|
||||||
/// `self.skills`, so concurrent writes land on different memory locations.
|
/// worker finished first.
|
||||||
/// 3. Each event only writes to its own items' likelihoods (no sharing).
|
|
||||||
#[cfg(feature = "rayon")]
|
#[cfg(feature = "rayon")]
|
||||||
fn sweep_color_groups<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
|
fn sweep_color_groups<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
@@ -411,29 +441,28 @@ impl<T: Time> TimeSlice<T> {
|
|||||||
if group_len == 0 {
|
if group_len == 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let range = self.color_groups.color_range(color_idx);
|
let range = self.color_groups.color_range(color_idx);
|
||||||
let p_draw = self.p_draw;
|
let p_draw = self.p_draw;
|
||||||
let convergence = self.convergence;
|
let convergence = self.convergence;
|
||||||
|
|
||||||
if group_len >= RAYON_THRESHOLD {
|
if group_len >= RAYON_THRESHOLD {
|
||||||
// Obtain a raw pointer from the unique `&mut self.skills` reference.
|
let skills = &self.skills;
|
||||||
// Casting back to `&mut` inside the closure is sound because:
|
let updates: Vec<EventUpdate> = self.events[range.clone()]
|
||||||
// 1. The pointer originates from a `&mut` — no aliasing with shared refs.
|
.par_iter()
|
||||||
// 2. Events in the same color group touch disjoint `Index` slots in the
|
.map(|ev| {
|
||||||
// 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| {
|
ARENA.with(|cell| {
|
||||||
let mut arena = cell.borrow_mut();
|
let mut arena = cell.borrow_mut();
|
||||||
arena.reset();
|
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 {
|
} else {
|
||||||
for ev in &mut self.events[range] {
|
for ev in &mut self.events[range] {
|
||||||
ev.iteration_direct(
|
ev.iteration_direct(
|
||||||
|
|||||||
Reference in New Issue
Block a user