From b1e0fcb817631c51f0e13bdbd8eef766839cdbfa Mon Sep 17 00:00:00 2001 From: Anders Olsson Date: Fri, 24 Apr 2026 07:24:29 +0200 Subject: [PATCH] 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. --- src/arena.rs | 44 ++++++++++ src/batch.rs | 10 ++- src/game.rs | 232 ++++++++++++++++++++++++++++++++++++------------- src/history.rs | 6 +- src/lib.rs | 1 + src/message.rs | 2 + 6 files changed, 234 insertions(+), 61 deletions(-) create mode 100644 src/arena.rs diff --git a/src/arena.rs b/src/arena.rs new file mode 100644 index 0000000..bd2edad --- /dev/null +++ b/src/arena.rs @@ -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, + pub(crate) diffs: Vec, + pub(crate) ties: Vec, + pub(crate) margins: Vec, +} + +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); + } +} diff --git a/src/batch.rs b/src/batch.rs index 8637251..8f350f3 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use crate::{ Index, N_INF, agent::Agent, + arena::ScratchArena, drift::Drift, game::Game, gaussian::Gaussian, @@ -111,6 +112,7 @@ pub struct Batch { pub(crate) skills: SkillStore, pub(crate) time: i64, p_draw: f64, + arena: ScratchArena, } impl Batch { @@ -120,6 +122,7 @@ impl Batch { skills: SkillStore::new(), time, p_draw, + arena: ScratchArena::new(), } } @@ -219,7 +222,7 @@ impl Batch { let teams = event.within_priors(false, false, &self.skills, agents); 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 (i, item) in team.items.iter_mut().enumerate() { @@ -295,6 +298,9 @@ impl Batch { forward: bool, agents: &AgentStore, ) -> f64 { + // log_evidence is infrequent; a local arena avoids needing &mut self. + let mut arena = ScratchArena::new(); + if targets.is_empty() { if online || forward { self.events @@ -306,6 +312,7 @@ impl Batch { &event.outputs(), &event.weights, self.p_draw, + &mut arena, ) .evidence .ln() @@ -331,6 +338,7 @@ impl Batch { &event.outputs(), &event.weights, self.p_draw, + &mut arena, ) .evidence .ln() diff --git a/src/game.rs b/src/game.rs index 315e6d9..e007011 100644 --- a/src/game.rs +++ b/src/game.rs @@ -1,5 +1,7 @@ use crate::{ - N_INF, N00, approx, compute_margin, + N_INF, N00, approx, + arena::ScratchArena, + compute_margin, drift::Drift, evidence, gaussian::Gaussian, @@ -24,6 +26,7 @@ impl<'a, D: Drift> Game<'a, D> { result: &'a [f64], weights: &'a [Vec], p_draw: f64, + arena: &mut ScratchArena, ) -> Self { debug_assert!( (result.len() == teams.len()), @@ -61,56 +64,62 @@ impl<'a, D: Drift> Game<'a, D> { evidence: 0.0, }; - this.likelihoods(); + this.likelihoods(arena); this } - fn likelihoods(&mut self) { + fn likelihoods(&mut self, arena: &mut ScratchArena) { + arena.reset(); let o = sort_perm(self.result, true); + let n_teams = o.len(); - let mut team = o - .iter() - .map(|&e| { - let performance = self.teams[e] - .iter() - .zip(self.weights[e].iter()) - .fold(N00, |p, (player, &weight)| { - p + (player.performance() * weight) - }); + // Phase 1: team messages into arena (avoids per-call allocation) + arena.teams.extend(o.iter().map(|&e| { + let performance = self.teams[e] + .iter() + .zip(self.weights[e].iter()) + .fold(N00, |p, (player, &weight)| { + p + (player.performance() * weight) + }); + TeamMessage { + prior: performance, + ..Default::default() + } + })); - TeamMessage { - prior: performance, - ..Default::default() - } - }) - .collect::>(); + // Phase 2: diff messages (split-borrow: teams immut, diffs mut) + { + let (teams, diffs) = (&arena.teams, &mut arena.diffs); + for i in 0..n_teams.saturating_sub(1) { + diffs.push(DiffMessage { + prior: teams[i].prior - teams[i + 1].prior, + likelihood: N_INF, + }); + } + } - let mut diff = team - .windows(2) - .map(|w| DiffMessage { - prior: w[0].prior - w[1].prior, - likelihood: N_INF, - }) - .collect::>(); + // Phase 3: tie and margin + arena + .ties + .extend(o.windows(2).map(|e| self.result[e[0]] == self.result[e[1]])); - let tie = o - .windows(2) - .map(|e| self.result[e[0]] == self.result[e[1]]) - .collect::>(); - - let margin = if self.p_draw == 0.0 { - vec![0.0; o.len() - 1] + if self.p_draw == 0.0 { + arena.margins.resize(n_teams.saturating_sub(1), 0.0); } else { - o.windows(2) - .map(|w| { - let a: f64 = self.teams[w[0]].iter().map(|a| a.beta.powi(2)).sum(); - let b: f64 = self.teams[w[1]].iter().map(|a| a.beta.powi(2)).sum(); + arena.margins.extend(o.windows(2).map(|w| { + let a: f64 = self.teams[w[0]].iter().map(|p| p.beta.powi(2)).sum(); + let b: f64 = self.teams[w[1]].iter().map(|p| p.beta.powi(2)).sum(); + compute_margin(self.p_draw, (a + b).sqrt()) + })); + } - compute_margin(self.p_draw, (a + b).sqrt()) - }) - .collect::>() - }; + // 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; @@ -204,7 +213,7 @@ mod tests { use ::approx::assert_ulps_eq; use super::*; - use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player}; + use crate::{ConstantDrift, GAMMA, Gaussian, N_INF, Player, arena::ScratchArena}; #[test] fn test_1vs1() { @@ -220,7 +229,13 @@ mod tests { ); 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 a = p[0][0]; @@ -241,7 +256,13 @@ mod tests { ); 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 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 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[1][0], N_INF); @@ -281,7 +308,13 @@ mod tests { ]; 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 a = p[0][0]; @@ -291,7 +324,13 @@ mod tests { 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 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 a = p[0][0]; @@ -301,7 +340,7 @@ mod tests { 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 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 a = p[0][0]; @@ -327,7 +366,13 @@ mod tests { ); 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 a = p[0][0]; @@ -348,7 +393,13 @@ mod tests { ); 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 a = p[0][0]; @@ -382,6 +433,7 @@ mod tests { &[0.0, 0.0, 0.0], &w, 0.25, + &mut ScratchArena::new(), ); let p = g.posteriors(); @@ -417,6 +469,7 @@ mod tests { &[0.0, 0.0, 0.0], &w, 0.25, + &mut ScratchArena::new(), ); 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 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(); 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 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(); assert_ulps_eq!( @@ -507,7 +572,13 @@ mod tests { let w_b = vec![0.7]; 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(); assert_ulps_eq!( @@ -525,7 +596,13 @@ mod tests { let w_b = vec![0.7]; 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(); assert_ulps_eq!( @@ -554,7 +631,13 @@ mod tests { )]; 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(); assert_ulps_eq!( @@ -583,7 +666,13 @@ mod tests { )]; 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(); 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 = [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(); assert_ulps_eq!( @@ -648,7 +743,13 @@ mod tests { let w_b = vec![0.7, 0.4]; 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(); assert_ulps_eq!( @@ -676,7 +777,13 @@ mod tests { let w_b = vec![0.7, 2.4]; 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(); assert_ulps_eq!( @@ -713,6 +820,7 @@ mod tests { &[1.0, 0.0], &w, 0.0, + &mut ScratchArena::new(), ); let post_2vs1 = g.posteriors(); @@ -720,7 +828,13 @@ mod tests { let w_b = vec![1.0, 0.0]; 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(); assert_ulps_eq!(p[0][0], post_2vs1[0][0], epsilon = 1e-6); diff --git a/src/history.rs b/src/history.rs index f2283d0..c34b743 100644 --- a/src/history.rs +++ b/src/history.rs @@ -436,7 +436,10 @@ mod tests { use approx::assert_ulps_eq; 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] fn test_init() { @@ -500,6 +503,7 @@ mod tests { &[0.0, 1.0], &w, P_DRAW, + &mut ScratchArena::new(), ) .posteriors(); let expected = p[0][0]; diff --git a/src/lib.rs b/src/lib.rs index b6c2924..ca0ea06 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ use std::{ pub mod agent; #[cfg(feature = "approx")] mod approx; +pub(crate) mod arena; pub mod batch; pub mod drift; mod error; diff --git a/src/message.rs b/src/message.rs index c6fd9bc..c91968e 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1,5 +1,6 @@ use crate::{N_INF, gaussian::Gaussian}; +#[derive(Debug)] pub(crate) struct TeamMessage { pub(crate) prior: Gaussian, pub(crate) likelihood_lose: Gaussian, @@ -67,6 +68,7 @@ impl DrawMessage { } } */ +#[derive(Debug)] pub(crate) struct DiffMessage { pub(crate) prior: Gaussian, pub(crate) likelihood: Gaussian,