From 9ab94002d951866a7a60261162baf3bb2f08baf4 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 14:51:03 +0100 Subject: [PATCH] perf(tpe): eliminate unnecessary Vec allocations in KDE/TPE sampling - In-place log transformation in sample_tpe_float (saves 2 allocs per log-scale float sample) - Pass owned Vec to sample_tpe_int to reduce peak memory during int-to-float conversion - Stack-allocate categorical count arrays for <=32 choices in sample_tpe_categorical --- src/sampler/motpe.rs | 55 ++++++++++++++++++++++++++------- src/sampler/tpe/multivariate.rs | 55 ++++++++++++++++++++++++++------- src/sampler/tpe/sampler.rs | 55 ++++++++++++++++++++++++++------- 3 files changed, 130 insertions(+), 35 deletions(-) diff --git a/src/sampler/motpe.rs b/src/sampler/motpe.rs index cdcf559..883cfe0 100644 --- a/src/sampler/motpe.rs +++ b/src/sampler/motpe.rs @@ -263,8 +263,20 @@ impl MotpeSampler { let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { let i_low = low.ln(); let i_high = high.ln(); - let g: Vec = good_values.iter().map(|&v| v.ln()).collect(); - let b: Vec = bad_values.iter().map(|&v| v.ln()).collect(); + let g = { + let mut v = good_values; + for x in &mut v { + *x = x.ln(); + } + v + }; + let b = { + let mut v = bad_values; + for x in &mut v { + *x = x.ln(); + } + v + }; (i_low, i_high, g, b) } else { (low, high, good_values, bad_values) @@ -339,12 +351,12 @@ impl MotpeSampler { high: i64, log_scale: bool, step: Option, - good_values: &[i64], - bad_values: &[i64], + good_values: Vec, + bad_values: Vec, rng: &mut fastrand::Rng, ) -> i64 { - let good_floats: Vec = good_values.iter().map(|&v| v as f64).collect(); - let bad_floats: Vec = bad_values.iter().map(|&v| v as f64).collect(); + let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); + let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); let float_value = self.sample_tpe_float( low as f64, @@ -375,9 +387,30 @@ impl MotpeSampler { bad_indices: &[usize], rng: &mut fastrand::Rng, ) -> usize { - let mut good_counts = vec![0usize; n_choices]; - let mut bad_counts = vec![0usize; n_choices]; + // Stack-allocate for the common case (<=32 choices), heap for rare large cases + let mut good_buf = [0usize; 32]; + let mut bad_buf = [0usize; 32]; + let mut weight_buf = [0.0f64; 32]; + let mut good_vec; + let mut bad_vec; + let mut weight_vec; + + let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = + if n_choices <= 32 { + ( + &mut good_buf[..n_choices], + &mut bad_buf[..n_choices], + &mut weight_buf[..n_choices], + ) + } else { + good_vec = vec![0usize; n_choices]; + bad_vec = vec![0usize; n_choices]; + weight_vec = vec![0.0f64; n_choices]; + (&mut good_vec, &mut bad_vec, &mut weight_vec) + }; + + // Count occurrences in good and bad groups for &idx in good_indices { if idx < n_choices { good_counts[idx] += 1; @@ -393,7 +426,7 @@ impl MotpeSampler { let good_total = good_indices.len() as f64 + n_choices as f64; let bad_total = bad_indices.len() as f64 + n_choices as f64; - let mut weights = vec![0.0f64; n_choices]; + // Calculate l(x)/g(x) ratio for each category for i in 0..n_choices { let l_prob = (good_counts[i] as f64 + 1.0) / good_total; let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; @@ -521,8 +554,8 @@ impl MultiObjectiveSampler for MotpeSampler { d.high, d.log_scale, d.step, - &good_values, - &bad_values, + good_values, + bad_values, &mut rng, ); ParamValue::Int(value) diff --git a/src/sampler/tpe/multivariate.rs b/src/sampler/tpe/multivariate.rs index cb9c4e3..731fc44 100644 --- a/src/sampler/tpe/multivariate.rs +++ b/src/sampler/tpe/multivariate.rs @@ -1356,8 +1356,8 @@ impl MultivariateTpeSampler { d.high, d.log_scale, d.step, - &good_values, - &bad_values, + good_values, + bad_values, rng, ); ParamValue::Int(value) @@ -1412,8 +1412,20 @@ impl MultivariateTpeSampler { let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { let i_low = low.ln(); let i_high = high.ln(); - let g: Vec = good_values.iter().map(|&v| v.ln()).collect(); - let b: Vec = bad_values.iter().map(|&v| v.ln()).collect(); + let g = { + let mut v = good_values; + for x in &mut v { + *x = x.ln(); + } + v + }; + let b = { + let mut v = bad_values; + for x in &mut v { + *x = x.ln(); + } + v + }; (i_low, i_high, g, b) } else { (low, high, good_values, bad_values) @@ -1487,13 +1499,13 @@ impl MultivariateTpeSampler { high: i64, log_scale: bool, step: Option, - good_values: &[i64], - bad_values: &[i64], + good_values: Vec, + bad_values: Vec, rng: &mut fastrand::Rng, ) -> i64 { // Convert to floats for KDE - let good_floats: Vec = good_values.iter().map(|&v| v as f64).collect(); - let bad_floats: Vec = bad_values.iter().map(|&v| v as f64).collect(); + let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); + let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); // Use float TPE sampling let float_value = self.sample_tpe_float( @@ -1598,10 +1610,30 @@ impl MultivariateTpeSampler { bad_indices: &[usize], rng: &mut fastrand::Rng, ) -> usize { - // Count occurrences in good and bad groups - let mut good_counts = vec![0usize; n_choices]; - let mut bad_counts = vec![0usize; n_choices]; + // Stack-allocate for the common case (<=32 choices), heap for rare large cases + let mut good_buf = [0usize; 32]; + let mut bad_buf = [0usize; 32]; + let mut weight_buf = [0.0f64; 32]; + let mut good_vec; + let mut bad_vec; + let mut weight_vec; + + let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = + if n_choices <= 32 { + ( + &mut good_buf[..n_choices], + &mut bad_buf[..n_choices], + &mut weight_buf[..n_choices], + ) + } else { + good_vec = vec![0usize; n_choices]; + bad_vec = vec![0usize; n_choices]; + weight_vec = vec![0.0f64; n_choices]; + (&mut good_vec, &mut bad_vec, &mut weight_vec) + }; + + // Count occurrences in good and bad groups for &idx in good_indices { if idx < n_choices { good_counts[idx] += 1; @@ -1618,7 +1650,6 @@ impl MultivariateTpeSampler { let bad_total = bad_indices.len() as f64 + n_choices as f64; // Calculate l(x)/g(x) ratio for each category - let mut weights = vec![0.0f64; n_choices]; for i in 0..n_choices { let l_prob = (good_counts[i] as f64 + 1.0) / good_total; let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs index a9ad698..fb9efe6 100644 --- a/src/sampler/tpe/sampler.rs +++ b/src/sampler/tpe/sampler.rs @@ -373,8 +373,20 @@ impl TpeSampler { let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { let i_low = low.ln(); let i_high = high.ln(); - let g: Vec = good_values.iter().map(|&v| v.ln()).collect(); - let b: Vec = bad_values.iter().map(|&v| v.ln()).collect(); + let g = { + let mut v = good_values; + for x in &mut v { + *x = x.ln(); + } + v + }; + let b = { + let mut v = bad_values; + for x in &mut v { + *x = x.ln(); + } + v + }; (i_low, i_high, g, b) } else { (low, high, good_values, bad_values) @@ -454,13 +466,13 @@ impl TpeSampler { high: i64, log_scale: bool, step: Option, - good_values: &[i64], - bad_values: &[i64], + good_values: Vec, + bad_values: Vec, rng: &mut fastrand::Rng, ) -> i64 { // Convert to floats for KDE - let good_floats: Vec = good_values.iter().map(|&v| v as f64).collect(); - let bad_floats: Vec = bad_values.iter().map(|&v| v as f64).collect(); + let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); + let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); // Use float TPE sampling let float_value = self.sample_tpe_float( @@ -497,10 +509,30 @@ impl TpeSampler { bad_indices: &[usize], rng: &mut fastrand::Rng, ) -> usize { - // Count occurrences in good and bad groups - let mut good_counts = vec![0usize; n_choices]; - let mut bad_counts = vec![0usize; n_choices]; + // Stack-allocate for the common case (<=32 choices), heap for rare large cases + let mut good_buf = [0usize; 32]; + let mut bad_buf = [0usize; 32]; + let mut weight_buf = [0.0f64; 32]; + let mut good_vec; + let mut bad_vec; + let mut weight_vec; + + let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = + if n_choices <= 32 { + ( + &mut good_buf[..n_choices], + &mut bad_buf[..n_choices], + &mut weight_buf[..n_choices], + ) + } else { + good_vec = vec![0usize; n_choices]; + bad_vec = vec![0usize; n_choices]; + weight_vec = vec![0.0f64; n_choices]; + (&mut good_vec, &mut bad_vec, &mut weight_vec) + }; + + // Count occurrences in good and bad groups for &idx in good_indices { if idx < n_choices { good_counts[idx] += 1; @@ -517,7 +549,6 @@ impl TpeSampler { let bad_total = bad_indices.len() as f64 + n_choices as f64; // Calculate l(x)/g(x) ratio for each category - let mut weights = vec![0.0f64; n_choices]; for i in 0..n_choices { let l_prob = (good_counts[i] as f64 + 1.0) / good_total; let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; @@ -967,8 +998,8 @@ impl Sampler for TpeSampler { d.high, d.log_scale, d.step, - &good_values, - &bad_values, + good_values, + bad_values, &mut rng, ); ParamValue::Int(value)