From c50e3909691e7a0418e228416eadbe87e995b315 Mon Sep 17 00:00:00 2001 From: Stephen Waits Date: Thu, 14 May 2026 07:02:13 -0600 Subject: [PATCH] perf(bayesian_opt): reuse scratch buffers in the EI acquisition loop (2.40M -> 2.26M instr) The acquisition loop ran acquisition_samples GP predictions per BO iteration, each allocating three short-lived Vecs: the candidate point, the k_star kernel vector, and solve_lower's output. Threading reused buffers through new sample_uniform_in_bounds_into / predict_into / solve_lower_into entry points removes ~3000 alloc/free pairs from bayesian_opt_short. bayesian_opt_short: 2_398_972 -> 2_255_904 (-6%). This is a structural (allocation) win, not an algorithmic one -- the GP fit (Cholesky) and EI prediction are inherently O(n^2)/O(n^3) with transcendental kernels, and that work is unchanged. Output bit-identical -- all 606 tests pass, including the run() snapshot; async builds clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/algorithms/bayesian_opt.rs | 74 ++++++++++++++++++++-------------- src/internal/cholesky.rs | 18 ++++++--- 2 files changed, 56 insertions(+), 36 deletions(-) diff --git a/src/algorithms/bayesian_opt.rs b/src/algorithms/bayesian_opt.rs index 76f7506..0318b7c 100644 --- a/src/algorithms/bayesian_opt.rs +++ b/src/algorithms/bayesian_opt.rs @@ -194,16 +194,22 @@ where let best_target = targets.iter().cloned().fold(f64::INFINITY, f64::min); - // Maximize EI by best-of-N random sampling. + // Maximize EI by best-of-N random sampling. `cand` and the two + // GP-prediction scratch buffers are reused across all samples + // so the inner loop allocates nothing. let mut best_x = sample_uniform_in_bounds(&self.bounds, &mut rng); let mut best_ei = -f64::INFINITY; + let mut cand: Vec = Vec::with_capacity(dim); + let mut k_star_buf: Vec = Vec::new(); + let mut v_temp_buf: Vec = Vec::new(); for _ in 0..self.config.acquisition_samples { - let cand = sample_uniform_in_bounds(&self.bounds, &mut rng); - let (mu, sigma) = posterior.predict(&cand); + sample_uniform_in_bounds_into(&self.bounds, &mut rng, &mut cand); + let (mu, sigma) = posterior.predict_into(&cand, &mut k_star_buf, &mut v_temp_buf); let ei = expected_improvement(mu, sigma, best_target); if ei > best_ei { best_ei = ei; - best_x = cand; + best_x.clear(); + best_x.extend_from_slice(&cand); } } @@ -270,18 +276,23 @@ fn better(a: &Evaluation, b: &Evaluation, direction: Direction) -> bool { } } +/// Sample a uniform-in-bounds point into `out` (reused across calls). +fn sample_uniform_in_bounds_into(bounds: &RealBounds, rng: &mut Rng, out: &mut Vec) { + out.clear(); + for &(lo, hi) in &bounds.bounds { + let v = if lo == hi { + lo + } else { + lo + (hi - lo) * rng.random::() + }; + out.push(v); + } +} + fn sample_uniform_in_bounds(bounds: &RealBounds, rng: &mut Rng) -> Vec { - bounds - .bounds - .iter() - .map(|&(lo, hi)| { - if lo == hi { - lo - } else { - lo + (hi - lo) * rng.random::() - } - }) - .collect() + let mut out = Vec::new(); + sample_uniform_in_bounds_into(bounds, rng, &mut out); + out } /// Anisotropic RBF kernel: `k(x, y) = σ² · exp(-0.5 · Σ ((x_i - y_i)/ℓ_i)²)`. @@ -333,26 +344,23 @@ impl GpPosterior { }) } - fn predict(&self, x: &[f64]) -> (f64, f64) { + /// Predict `(mean, std)` at `x`, using caller-owned scratch buffers + /// (`k_star`, `v_temp`) so the hot acquisition loop allocates nothing. + fn predict_into(&self, x: &[f64], k_star: &mut Vec, v_temp: &mut Vec) -> (f64, f64) { let n = self.decisions.len(); - let mut k_star = vec![0.0_f64; n]; - for (i, k_star_i) in k_star.iter_mut().enumerate() { - *k_star_i = rbf_kernel( - x, - &self.decisions[i], - &self.length_scales, - self.signal_variance, - ); + k_star.clear(); + k_star.reserve(n); + for d in &self.decisions { + k_star.push(rbf_kernel(x, d, &self.length_scales, self.signal_variance)); } - let _ = n; let mu: f64 = k_star .iter() .zip(self.alpha.iter()) .map(|(a, b)| a * b) .sum(); - // Var = k(x,x) - k_star^T · K^{-1} · k_star - // Compute K^{-1}·k_star = solve_upper_transpose(L, solve_lower(L, k_star)) - let v_temp = crate::internal::cholesky::solve_lower(&self.chol_l, &k_star); + // Var = k(x,x) - k_star^T · K^{-1} · k_star; the squared norm of + // `solve_lower(L, k_star)` is exactly `k_star^T · K^{-1} · k_star`. + crate::internal::cholesky::solve_lower_into(&self.chol_l, k_star, v_temp); let v: f64 = v_temp.iter().map(|x| x * x).sum(); let var = (self.signal_variance - v).max(0.0); (mu, var.sqrt()) @@ -496,13 +504,17 @@ impl BayesianOpt { let mut best_x = sample_uniform_in_bounds(&self.bounds, &mut rng); let mut best_ei = -f64::INFINITY; + let mut cand: Vec = Vec::with_capacity(dim); + let mut k_star_buf: Vec = Vec::new(); + let mut v_temp_buf: Vec = Vec::new(); for _ in 0..self.config.acquisition_samples { - let cand = sample_uniform_in_bounds(&self.bounds, &mut rng); - let (mu, sigma) = posterior.predict(&cand); + sample_uniform_in_bounds_into(&self.bounds, &mut rng, &mut cand); + let (mu, sigma) = posterior.predict_into(&cand, &mut k_star_buf, &mut v_temp_buf); let ei = expected_improvement(mu, sigma, best_target); if ei > best_ei { best_ei = ei; - best_x = cand; + best_x.clear(); + best_x.extend_from_slice(&cand); } } diff --git a/src/internal/cholesky.rs b/src/internal/cholesky.rs index f4b2f05..7af766a 100644 --- a/src/internal/cholesky.rs +++ b/src/internal/cholesky.rs @@ -39,17 +39,25 @@ pub(crate) fn cholesky(a: &[Vec]) -> Result>, &'static str> { Ok(l) } -/// Solve `L · y = b` (forward substitution) for lower-triangular `L`. -pub(crate) fn solve_lower(l: &[Vec], b: &[f64]) -> Vec { +/// Solve `L · y = b` (forward substitution) for lower-triangular `L`, +/// writing the result into `out` (reused across calls to avoid allocating). +pub(crate) fn solve_lower_into(l: &[Vec], b: &[f64], out: &mut Vec) { let n = l.len(); - let mut y = vec![0.0_f64; n]; + out.clear(); + out.resize(n, 0.0); for i in 0..n { let mut sum = b[i]; for k in 0..i { - sum -= l[i][k] * y[k]; + sum -= l[i][k] * out[k]; } - y[i] = sum / l[i][i]; + out[i] = sum / l[i][i]; } +} + +/// Solve `L · y = b` (forward substitution) for lower-triangular `L`. +pub(crate) fn solve_lower(l: &[Vec], b: &[f64]) -> Vec { + let mut y = Vec::new(); + solve_lower_into(l, b, &mut y); y }