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:
44
src/arena.rs
Normal file
44
src/arena.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
10
src/batch.rs
10
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<D>,
|
||||
) -> 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()
|
||||
|
||||
232
src/game.rs
232
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<f64>],
|
||||
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::<Vec<_>>();
|
||||
// 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::<Vec<_>>();
|
||||
// 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::<Vec<_>>();
|
||||
|
||||
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::<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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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];
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user