perf(game): eliminate per-event allocations via ScratchArena

Game::likelihoods previously allocated four Vecs (teams, diffs, ties,
margins) on every call. Batch now owns one ScratchArena reused across
all Game::new calls in the iteration loop; likelihoods() clears and
extends the arena buffers instead of allocating fresh.

For log_evidence (called infrequently), a local ScratchArena is created
per invocation so the method signature stays &self.

Also: add #[derive(Debug)] to TeamMessage and DiffMessage (required by
ScratchArena's own Debug derive).

Part of T0 engine redesign.
This commit is contained in:
2026-04-24 07:24:29 +02:00
parent 49d2b317da
commit b1e0fcb817
6 changed files with 234 additions and 61 deletions

44
src/arena.rs Normal file
View File

@@ -0,0 +1,44 @@
use crate::message::{DiffMessage, TeamMessage};
/// Reusable scratch buffers for `Game::likelihoods`.
///
/// The four Vecs previously allocated fresh on every `Game::new` call —
/// `teams`, `diffs`, `ties`, `margins` — are now borrowed from this arena,
/// reset between uses. A `Batch` owns one arena; all events in the slice
/// share it across the convergence iterations.
#[derive(Debug, Default)]
pub struct ScratchArena {
pub(crate) teams: Vec<TeamMessage>,
pub(crate) diffs: Vec<DiffMessage>,
pub(crate) ties: Vec<bool>,
pub(crate) margins: Vec<f64>,
}
impl ScratchArena {
pub fn new() -> Self {
Self::default()
}
#[inline]
pub(crate) fn reset(&mut self) {
self.teams.clear();
self.diffs.clear();
self.ties.clear();
self.margins.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reset_keeps_capacity() {
let mut arena = ScratchArena::new();
arena.teams.push(TeamMessage::default());
let cap = arena.teams.capacity();
arena.reset();
assert_eq!(arena.teams.len(), 0);
assert_eq!(arena.teams.capacity(), cap);
}
}

View File

@@ -3,6 +3,7 @@ use std::collections::HashMap;
use crate::{ use crate::{
Index, N_INF, Index, N_INF,
agent::Agent, agent::Agent,
arena::ScratchArena,
drift::Drift, drift::Drift,
game::Game, game::Game,
gaussian::Gaussian, gaussian::Gaussian,
@@ -111,6 +112,7 @@ pub struct Batch {
pub(crate) skills: SkillStore, pub(crate) skills: SkillStore,
pub(crate) time: i64, pub(crate) time: i64,
p_draw: f64, p_draw: f64,
arena: ScratchArena,
} }
impl Batch { impl Batch {
@@ -120,6 +122,7 @@ impl Batch {
skills: SkillStore::new(), skills: SkillStore::new(),
time, time,
p_draw, p_draw,
arena: ScratchArena::new(),
} }
} }
@@ -219,7 +222,7 @@ impl Batch {
let teams = event.within_priors(false, false, &self.skills, agents); let teams = event.within_priors(false, false, &self.skills, agents);
let result = event.outputs(); let result = event.outputs();
let g = Game::new(teams, &result, &event.weights, self.p_draw); let g = Game::new(teams, &result, &event.weights, self.p_draw, &mut self.arena);
for (t, team) in event.teams.iter_mut().enumerate() { for (t, team) in event.teams.iter_mut().enumerate() {
for (i, item) in team.items.iter_mut().enumerate() { for (i, item) in team.items.iter_mut().enumerate() {
@@ -295,6 +298,9 @@ impl Batch {
forward: bool, forward: bool,
agents: &AgentStore<D>, agents: &AgentStore<D>,
) -> f64 { ) -> f64 {
// log_evidence is infrequent; a local arena avoids needing &mut self.
let mut arena = ScratchArena::new();
if targets.is_empty() { if targets.is_empty() {
if online || forward { if online || forward {
self.events self.events
@@ -306,6 +312,7 @@ impl Batch {
&event.outputs(), &event.outputs(),
&event.weights, &event.weights,
self.p_draw, self.p_draw,
&mut arena,
) )
.evidence .evidence
.ln() .ln()
@@ -331,6 +338,7 @@ impl Batch {
&event.outputs(), &event.outputs(),
&event.weights, &event.weights,
self.p_draw, self.p_draw,
&mut arena,
) )
.evidence .evidence
.ln() .ln()

View File

@@ -1,5 +1,7 @@
use crate::{ use crate::{
N_INF, N00, approx, compute_margin, N_INF, N00, approx,
arena::ScratchArena,
compute_margin,
drift::Drift, drift::Drift,
evidence, evidence,
gaussian::Gaussian, gaussian::Gaussian,
@@ -24,6 +26,7 @@ impl<'a, D: Drift> Game<'a, D> {
result: &'a [f64], result: &'a [f64],
weights: &'a [Vec<f64>], weights: &'a [Vec<f64>],
p_draw: f64, p_draw: f64,
arena: &mut ScratchArena,
) -> Self { ) -> Self {
debug_assert!( debug_assert!(
(result.len() == teams.len()), (result.len() == teams.len()),
@@ -61,56 +64,62 @@ impl<'a, D: Drift> Game<'a, D> {
evidence: 0.0, evidence: 0.0,
}; };
this.likelihoods(); this.likelihoods(arena);
this this
} }
fn likelihoods(&mut self) { fn likelihoods(&mut self, arena: &mut ScratchArena) {
arena.reset();
let o = sort_perm(self.result, true); let o = sort_perm(self.result, true);
let n_teams = o.len();
let mut team = o // Phase 1: team messages into arena (avoids per-call allocation)
.iter() arena.teams.extend(o.iter().map(|&e| {
.map(|&e| {
let performance = self.teams[e] let performance = self.teams[e]
.iter() .iter()
.zip(self.weights[e].iter()) .zip(self.weights[e].iter())
.fold(N00, |p, (player, &weight)| { .fold(N00, |p, (player, &weight)| {
p + (player.performance() * weight) p + (player.performance() * weight)
}); });
TeamMessage { TeamMessage {
prior: performance, prior: performance,
..Default::default() ..Default::default()
} }
}) }));
.collect::<Vec<_>>();
let mut diff = team // Phase 2: diff messages (split-borrow: teams immut, diffs mut)
.windows(2) {
.map(|w| DiffMessage { let (teams, diffs) = (&arena.teams, &mut arena.diffs);
prior: w[0].prior - w[1].prior, for i in 0..n_teams.saturating_sub(1) {
diffs.push(DiffMessage {
prior: teams[i].prior - teams[i + 1].prior,
likelihood: N_INF, likelihood: N_INF,
}) });
.collect::<Vec<_>>(); }
}
let tie = o // Phase 3: tie and margin
.windows(2) arena
.map(|e| self.result[e[0]] == self.result[e[1]]) .ties
.collect::<Vec<_>>(); .extend(o.windows(2).map(|e| self.result[e[0]] == self.result[e[1]]));
let margin = if self.p_draw == 0.0 { if self.p_draw == 0.0 {
vec![0.0; o.len() - 1] arena.margins.resize(n_teams.saturating_sub(1), 0.0);
} else { } else {
o.windows(2) arena.margins.extend(o.windows(2).map(|w| {
.map(|w| { let a: f64 = self.teams[w[0]].iter().map(|p| p.beta.powi(2)).sum();
let a: f64 = self.teams[w[0]].iter().map(|a| a.beta.powi(2)).sum(); let b: f64 = self.teams[w[1]].iter().map(|p| p.beta.powi(2)).sum();
let b: f64 = self.teams[w[1]].iter().map(|a| a.beta.powi(2)).sum();
compute_margin(self.p_draw, (a + b).sqrt()) compute_margin(self.p_draw, (a + b).sqrt())
}) }));
.collect::<Vec<_>>() }
};
// Use local aliases for the arena slices for readability in the EP loop.
// These are references into the arena, not copies.
let team = &mut arena.teams;
let diff = &mut arena.diffs;
let tie = &arena.ties;
let margin = &arena.margins;
self.evidence = 1.0; self.evidence = 1.0;
@@ -204,7 +213,7 @@ mod tests {
use ::approx::assert_ulps_eq; use ::approx::assert_ulps_eq;
use super::*; use super::*;
use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player}; use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player, arena::ScratchArena};
#[test] #[test]
fn test_1vs1() { fn test_1vs1() {
@@ -220,7 +229,13 @@ mod tests {
); );
let w = [vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0); let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -241,7 +256,13 @@ mod tests {
); );
let w = [vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0); let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -254,7 +275,13 @@ mod tests {
let t_b = Player::new(Gaussian::from_ms(15.568, 0.51), 1.0, ConstantDrift(0.2125)); let t_b = Player::new(Gaussian::from_ms(15.568, 0.51), 1.0, ConstantDrift(0.2125));
let w = [vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 1.0], &w, 0.0); let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 1.0],
&w,
0.0,
&mut ScratchArena::new(),
);
assert_eq!(g.likelihoods[0][0], N_INF); assert_eq!(g.likelihoods[0][0], N_INF);
assert_eq!(g.likelihoods[1][0], N_INF); assert_eq!(g.likelihoods[1][0], N_INF);
@@ -281,7 +308,13 @@ mod tests {
]; ];
let w = [vec![1.0], vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams.clone(), &[1.0, 2.0, 0.0], &w, 0.0); let g = Game::new(
teams.clone(),
&[1.0, 2.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -291,7 +324,13 @@ mod tests {
assert_ulps_eq!(b, Gaussian::from_ms(31.311358, 6.698818), epsilon = 1e-6); assert_ulps_eq!(b, Gaussian::from_ms(31.311358, 6.698818), epsilon = 1e-6);
let w = [vec![1.0], vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams.clone(), &[2.0, 1.0, 0.0], &w, 0.0); let g = Game::new(
teams.clone(),
&[2.0, 1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -301,7 +340,7 @@ mod tests {
assert_ulps_eq!(b, Gaussian::from_ms(25.000000, 6.238469), epsilon = 1e-6); assert_ulps_eq!(b, Gaussian::from_ms(25.000000, 6.238469), epsilon = 1e-6);
let w = [vec![1.0], vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0], vec![1.0]];
let g = Game::new(teams, &[1.0, 2.0, 0.0], &w, 0.5); let g = Game::new(teams, &[1.0, 2.0, 0.0], &w, 0.5, &mut ScratchArena::new());
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -327,7 +366,13 @@ mod tests {
); );
let w = [vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 0.0], &w, 0.25); let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -348,7 +393,13 @@ mod tests {
); );
let w = [vec![1.0], vec![1.0]]; let w = [vec![1.0], vec![1.0]];
let g = Game::new(vec![vec![t_a], vec![t_b]], &[0.0, 0.0], &w, 0.25); let g = Game::new(
vec![vec![t_a], vec![t_b]],
&[0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
let a = p[0][0]; let a = p[0][0];
@@ -382,6 +433,7 @@ mod tests {
&[0.0, 0.0, 0.0], &[0.0, 0.0, 0.0],
&w, &w,
0.25, 0.25,
&mut ScratchArena::new(),
); );
let p = g.posteriors(); let p = g.posteriors();
@@ -417,6 +469,7 @@ mod tests {
&[0.0, 0.0, 0.0], &[0.0, 0.0, 0.0],
&w, &w,
0.25, 0.25,
&mut ScratchArena::new(),
); );
let p = g.posteriors(); let p = g.posteriors();
@@ -462,7 +515,13 @@ mod tests {
]; ];
let w = [vec![1.0, 1.0], vec![1.0], vec![1.0, 1.0]]; let w = [vec![1.0, 1.0], vec![1.0], vec![1.0, 1.0]];
let g = Game::new(vec![t_a, t_b, t_c], &[1.0, 0.0, 0.0], &w, 0.25); let g = Game::new(
vec![t_a, t_b, t_c],
&[1.0, 0.0, 0.0],
&w,
0.25,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!(p[0][0], Gaussian::from_ms(13.051, 2.864), epsilon = 1e-3); assert_ulps_eq!(p[0][0], Gaussian::from_ms(13.051, 2.864), epsilon = 1e-3);
@@ -489,7 +548,13 @@ mod tests {
)]; )];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -507,7 +572,13 @@ mod tests {
let w_b = vec![0.7]; let w_b = vec![0.7];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -525,7 +596,13 @@ mod tests {
let w_b = vec![0.7]; let w_b = vec![0.7];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -554,7 +631,13 @@ mod tests {
)]; )];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -583,7 +666,13 @@ mod tests {
)]; )];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a, t_b],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!(p[0][0], p[1][0], epsilon = 1e-6); assert_ulps_eq!(p[0][0], p[1][0], epsilon = 1e-6);
@@ -620,7 +709,13 @@ mod tests {
let w_b = vec![0.9, 0.6]; let w_b = vec![0.9, 0.6];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -648,7 +743,13 @@ mod tests {
let w_b = vec![0.7, 0.4]; let w_b = vec![0.7, 0.4];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -676,7 +777,13 @@ mod tests {
let w_b = vec![0.7, 2.4]; let w_b = vec![0.7, 2.4];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a.clone(), t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a.clone(), t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!( assert_ulps_eq!(
@@ -713,6 +820,7 @@ mod tests {
&[1.0, 0.0], &[1.0, 0.0],
&w, &w,
0.0, 0.0,
&mut ScratchArena::new(),
); );
let post_2vs1 = g.posteriors(); let post_2vs1 = g.posteriors();
@@ -720,7 +828,13 @@ mod tests {
let w_b = vec![1.0, 0.0]; let w_b = vec![1.0, 0.0];
let w = [w_a, w_b]; let w = [w_a, w_b];
let g = Game::new(vec![t_a, t_b.clone()], &[1.0, 0.0], &w, 0.0); let g = Game::new(
vec![t_a, t_b.clone()],
&[1.0, 0.0],
&w,
0.0,
&mut ScratchArena::new(),
);
let p = g.posteriors(); let p = g.posteriors();
assert_ulps_eq!(p[0][0], post_2vs1[0][0], epsilon = 1e-6); assert_ulps_eq!(p[0][0], post_2vs1[0][0], epsilon = 1e-6);

View File

@@ -436,7 +436,10 @@ mod tests {
use approx::assert_ulps_eq; use approx::assert_ulps_eq;
use super::*; use super::*;
use crate::{ConstantDrift, EPSILON, Game, Gaussian, ITERATIONS, IndexMap, P_DRAW, Player}; use crate::{
ConstantDrift, EPSILON, Game, Gaussian, ITERATIONS, IndexMap, P_DRAW, Player,
arena::ScratchArena,
};
#[test] #[test]
fn test_init() { fn test_init() {
@@ -500,6 +503,7 @@ mod tests {
&[0.0, 1.0], &[0.0, 1.0],
&w, &w,
P_DRAW, P_DRAW,
&mut ScratchArena::new(),
) )
.posteriors(); .posteriors();
let expected = p[0][0]; let expected = p[0][0];

View File

@@ -9,6 +9,7 @@ use std::{
pub mod agent; pub mod agent;
#[cfg(feature = "approx")] #[cfg(feature = "approx")]
mod approx; mod approx;
pub(crate) mod arena;
pub mod batch; pub mod batch;
pub mod drift; pub mod drift;
mod error; mod error;

View File

@@ -1,5 +1,6 @@
use crate::{N_INF, gaussian::Gaussian}; use crate::{N_INF, gaussian::Gaussian};
#[derive(Debug)]
pub(crate) struct TeamMessage { pub(crate) struct TeamMessage {
pub(crate) prior: Gaussian, pub(crate) prior: Gaussian,
pub(crate) likelihood_lose: Gaussian, pub(crate) likelihood_lose: Gaussian,
@@ -67,6 +68,7 @@ impl DrawMessage {
} }
} }
*/ */
#[derive(Debug)]
pub(crate) struct DiffMessage { pub(crate) struct DiffMessage {
pub(crate) prior: Gaussian, pub(crate) prior: Gaussian,
pub(crate) likelihood: Gaussian, pub(crate) likelihood: Gaussian,