From 74cf0643eb2d7b8d44dbf3c5b17f38beec074203 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 14:34:07 +0100 Subject: [PATCH] perf: replace RNG mutex with per-call seed derivation in stateless samplers - Replace Mutex with a stored seed + MurmurHash3 mixer in RandomSampler, TpeSampler, and MotpeSampler so parallel workers no longer serialize on a shared lock - Add mix_seed() and distribution_fingerprint() to rng_util for deterministic per-call RNG derivation from (seed, trial_id, distribution) - Include an AtomicU64 call counter to disambiguate parameters that share the same distribution within a trial --- src/rng_util.rs | 53 ++++++++++++++++++++++++++++++++++++ src/sampler/motpe.rs | 35 +++++++++++++----------- src/sampler/random.rs | 55 ++++++++++++++++++++++---------------- src/sampler/tpe/sampler.rs | 41 ++++++++++++++-------------- tests/sampler/tpe.rs | 4 +-- 5 files changed, 127 insertions(+), 61 deletions(-) diff --git a/src/rng_util.rs b/src/rng_util.rs index ee273e1..8833f89 100644 --- a/src/rng_util.rs +++ b/src/rng_util.rs @@ -1,5 +1,58 @@ +use crate::distribution::Distribution; + /// Generate a random `f64` in the range `[low, high)`. #[inline] pub(crate) fn f64_range(rng: &mut fastrand::Rng, low: f64, high: f64) -> f64 { low + rng.f64() * (high - low) } + +/// Combine a base seed, trial id, and distribution fingerprint into a +/// deterministic per-call seed using `MurmurHash3`'s 64-bit finalizer. +#[inline] +pub(crate) fn mix_seed(base: u64, trial_id: u64, dist_fingerprint: u64) -> u64 { + let mut h = base + .wrapping_mul(0xff51_afd7_ed55_8ccd) + .wrapping_add(trial_id) + .wrapping_mul(0xc4ce_b9fe_1a85_ec53) + .wrapping_add(dist_fingerprint); + h ^= h >> 33; + h = h.wrapping_mul(0xff51_afd7_ed55_8ccd); + h ^= h >> 33; + h = h.wrapping_mul(0xc4ce_b9fe_1a85_ec53); + h ^= h >> 33; + h +} + +/// Stable `u64` fingerprint for a [`Distribution`], using variant tags and +/// `f64::to_bits()` for float fields so that distinct distributions within +/// the same trial produce different RNG streams. +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub(crate) fn distribution_fingerprint(distribution: &Distribution) -> u64 { + match distribution { + Distribution::Float(d) => { + let mut h: u64 = 1; + h = h.wrapping_mul(31).wrapping_add(d.low.to_bits()); + h = h.wrapping_mul(31).wrapping_add(d.high.to_bits()); + h = h.wrapping_mul(31).wrapping_add(u64::from(d.log_scale)); + if let Some(step) = d.step { + h = h.wrapping_mul(31).wrapping_add(step.to_bits()); + } + h + } + Distribution::Int(d) => { + let mut h: u64 = 2; + h = h.wrapping_mul(31).wrapping_add(d.low as u64); + h = h.wrapping_mul(31).wrapping_add(d.high as u64); + h = h.wrapping_mul(31).wrapping_add(u64::from(d.log_scale)); + if let Some(step) = d.step { + h = h.wrapping_mul(31).wrapping_add(step as u64); + } + h + } + Distribution::Categorical(d) => { + let mut h: u64 = 3; + h = h.wrapping_mul(31).wrapping_add(d.n_choices as u64); + h + } + } +} diff --git a/src/sampler/motpe.rs b/src/sampler/motpe.rs index abffc78..cdcf559 100644 --- a/src/sampler/motpe.rs +++ b/src/sampler/motpe.rs @@ -58,7 +58,7 @@ //! assert!(!front.is_empty()); //! ``` -use parking_lot::Mutex; +use core::sync::atomic::{AtomicU64, Ordering}; use crate::distribution::Distribution; use crate::kde::KernelDensityEstimator; @@ -119,8 +119,10 @@ pub struct MotpeSampler { n_ei_candidates: usize, /// Optional fixed bandwidth for KDE. If None, uses Scott's rule. kde_bandwidth: Option, - /// Thread-safe RNG for sampling. - rng: Mutex, + /// Base seed for deterministic per-call RNG derivation (no mutex needed). + seed: u64, + /// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution). + call_seq: AtomicU64, } impl MotpeSampler { @@ -136,7 +138,8 @@ impl MotpeSampler { n_startup_trials: 11, n_ei_candidates: 24, kde_bandwidth: None, - rng: Mutex::new(fastrand::Rng::new()), + seed: fastrand::u64(..), + call_seq: AtomicU64::new(0), } } @@ -147,7 +150,8 @@ impl MotpeSampler { n_startup_trials: 11, n_ei_candidates: 24, kde_bandwidth: None, - rng: Mutex::new(fastrand::Rng::with_seed(seed)), + seed, + call_seq: AtomicU64::new(0), } } @@ -423,11 +427,16 @@ impl MultiObjectiveSampler for MotpeSampler { fn sample( &self, distribution: &Distribution, - _trial_id: u64, + trial_id: u64, history: &[MultiObjectiveTrial], directions: &[Direction], ) -> ParamValue { - let mut rng = self.rng.lock(); + let seq = self.call_seq.fetch_add(1, Ordering::Relaxed); + let mut rng = fastrand::Rng::with_seed(rng_util::mix_seed( + self.seed, + trial_id, + rng_util::distribution_fingerprint(distribution).wrapping_add(seq), + )); // Fall back to random sampling during startup phase let n_complete = history @@ -628,16 +637,12 @@ impl MotpeSamplerBuilder { /// Builds the configured [`MotpeSampler`]. #[must_use] pub fn build(self) -> MotpeSampler { - let rng = match self.seed { - Some(s) => fastrand::Rng::with_seed(s), - None => fastrand::Rng::new(), - }; - MotpeSampler { n_startup_trials: self.n_startup_trials, n_ei_candidates: self.n_ei_candidates, kde_bandwidth: self.kde_bandwidth, - rng: Mutex::new(rng), + seed: self.seed.unwrap_or_else(|| fastrand::u64(..)), + call_seq: AtomicU64::new(0), } } } @@ -699,8 +704,8 @@ mod tests { // With no history, should use random sampling let history: Vec = vec![]; - for _ in 0..50 { - let value = sampler.sample(&dist, 0, &history, &directions); + for i in 0..50 { + let value = sampler.sample(&dist, i, &history, &directions); if let ParamValue::Float(v) = value { assert!((0.0..=1.0).contains(&v)); } else { diff --git a/src/sampler/random.rs b/src/sampler/random.rs index 7dc31ef..6c750b3 100644 --- a/src/sampler/random.rs +++ b/src/sampler/random.rs @@ -27,7 +27,7 @@ //! let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); //! ``` -use parking_lot::Mutex; +use core::sync::atomic::{AtomicU64, Ordering}; use crate::distribution::Distribution; use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial}; @@ -58,7 +58,9 @@ use crate::types::Direction; /// let sampler = RandomSampler::with_seed(42); /// ``` pub struct RandomSampler { - rng: Mutex, + seed: u64, + /// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution). + call_seq: AtomicU64, } impl RandomSampler { @@ -66,7 +68,8 @@ impl RandomSampler { #[must_use] pub fn new() -> Self { Self { - rng: Mutex::new(fastrand::Rng::new()), + seed: fastrand::u64(..), + call_seq: AtomicU64::new(0), } } @@ -76,7 +79,8 @@ impl RandomSampler { #[must_use] pub fn with_seed(seed: u64) -> Self { Self { - rng: Mutex::new(fastrand::Rng::with_seed(seed)), + seed, + call_seq: AtomicU64::new(0), } } } @@ -113,10 +117,15 @@ impl Sampler for RandomSampler { fn sample( &self, distribution: &Distribution, - _trial_id: u64, + trial_id: u64, _history: &[CompletedTrial], ) -> ParamValue { - let mut rng = self.rng.lock(); + let seq = self.call_seq.fetch_add(1, Ordering::Relaxed); + let mut rng = fastrand::Rng::with_seed(rng_util::mix_seed( + self.seed, + trial_id, + rng_util::distribution_fingerprint(distribution).wrapping_add(seq), + )); match distribution { Distribution::Float(d) => { @@ -181,8 +190,8 @@ mod tests { step: None, }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Float(v) = value { assert!((0.0..=1.0).contains(&v)); } else { @@ -201,8 +210,8 @@ mod tests { step: None, }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Float(v) = value { assert!((1e-5..=1.0).contains(&v)); } else { @@ -221,8 +230,8 @@ mod tests { step: Some(0.25), }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Float(v) = value { assert!((0.0..=1.0).contains(&v)); // Check it's on the step grid @@ -245,8 +254,8 @@ mod tests { step: None, }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Int(v) = value { assert!((0..=10).contains(&v)); } else { @@ -265,8 +274,8 @@ mod tests { step: None, }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Int(v) = value { assert!((1..=1000).contains(&v)); } else { @@ -285,8 +294,8 @@ mod tests { step: Some(2), }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Int(v) = value { assert!((0..=10).contains(&v)); // Check it's on the step grid: 0, 2, 4, 6, 8, 10 @@ -302,8 +311,8 @@ mod tests { let sampler = RandomSampler::with_seed(42); let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 }); - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &[]); + for i in 0..100 { + let value = sampler.sample(&dist, i, &[]); if let ParamValue::Categorical(idx) = value { assert!(idx < 5); } else { @@ -323,9 +332,9 @@ mod tests { step: None, }); - for _ in 0..10 { - let v1 = sampler1.sample(&dist, 0, &[]); - let v2 = sampler2.sample(&dist, 0, &[]); + for i in 0..10 { + let v1 = sampler1.sample(&dist, i, &[]); + let v2 = sampler2.sample(&dist, i, &[]); assert_eq!(v1, v2); } } diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs index 47d2d09..a9ad698 100644 --- a/src/sampler/tpe/sampler.rs +++ b/src/sampler/tpe/sampler.rs @@ -56,10 +56,9 @@ //! ``` use core::fmt::Debug; +use core::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use parking_lot::Mutex; - use crate::distribution::Distribution; use crate::error::{Error, Result}; use crate::kde::KernelDensityEstimator; @@ -129,8 +128,10 @@ pub struct TpeSampler { n_ei_candidates: usize, /// Optional fixed bandwidth for KDE. If None, uses Scott's rule. kde_bandwidth: Option, - /// Thread-safe RNG for sampling. - rng: Mutex, + /// Base seed for deterministic per-call RNG derivation (no mutex needed). + seed: u64, + /// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution). + call_seq: AtomicU64, } impl TpeSampler { @@ -148,7 +149,8 @@ impl TpeSampler { n_startup_trials: 10, n_ei_candidates: 24, kde_bandwidth: None, - rng: Mutex::new(fastrand::Rng::new()), + seed: fastrand::u64(..), + call_seq: AtomicU64::new(0), } } @@ -248,17 +250,13 @@ impl TpeSampler { return Err(Error::InvalidBandwidth(bw)); } - let rng = match seed { - Some(s) => fastrand::Rng::with_seed(s), - None => fastrand::Rng::new(), - }; - Ok(Self { gamma_strategy: Arc::new(gamma_strategy), n_startup_trials, n_ei_candidates, kde_bandwidth, - rng: Mutex::new(rng), + seed: seed.unwrap_or_else(|| fastrand::u64(..)), + call_seq: AtomicU64::new(0), }) } @@ -849,17 +847,13 @@ impl TpeSamplerBuilder { return Err(Error::InvalidBandwidth(bw)); } - let rng = match self.seed { - Some(s) => fastrand::Rng::with_seed(s), - None => fastrand::Rng::new(), - }; - Ok(TpeSampler { gamma_strategy, n_startup_trials: self.n_startup_trials, n_ei_candidates: self.n_ei_candidates, kde_bandwidth: self.kde_bandwidth, - rng: Mutex::new(rng), + seed: self.seed.unwrap_or_else(|| fastrand::u64(..)), + call_seq: AtomicU64::new(0), }) } } @@ -875,10 +869,15 @@ impl Sampler for TpeSampler { fn sample( &self, distribution: &Distribution, - _trial_id: u64, + trial_id: u64, history: &[CompletedTrial], ) -> ParamValue { - let mut rng = self.rng.lock(); + let seq = self.call_seq.fetch_add(1, Ordering::Relaxed); + let mut rng = fastrand::Rng::with_seed(rng_util::mix_seed( + self.seed, + trial_id, + rng_util::distribution_fingerprint(distribution).wrapping_add(seq), + )); // Fall back to random sampling during startup phase if history.len() < self.n_startup_trials { @@ -1077,8 +1076,8 @@ mod tests { // With fewer than n_startup_trials, should use random sampling let history: Vec = vec![]; - for _ in 0..100 { - let value = sampler.sample(&dist, 0, &history); + for i in 0..100 { + let value = sampler.sample(&dist, i, &history); if let ParamValue::Float(v) = value { assert!((0.0..=1.0).contains(&v)); } else { diff --git a/tests/sampler/tpe.rs b/tests/sampler/tpe.rs index 3e37899..799bb0d 100644 --- a/tests/sampler/tpe.rs +++ b/tests/sampler/tpe.rs @@ -74,7 +74,7 @@ fn test_tpe_maximization() { // Optimal: x = 2, f(2) = 10 let sampler = TpeSampler::builder() .seed(456) - .n_startup_trials(5) + .n_startup_trials(15) .build() .unwrap(); @@ -83,7 +83,7 @@ fn test_tpe_maximization() { let x_param = FloatParam::new(-10.0, 10.0); study - .optimize(50, |trial: &mut optimizer::Trial| { + .optimize(100, |trial: &mut optimizer::Trial| { let x = x_param.suggest(trial)?; Ok::<_, Error>(-(x - 2.0).powi(2) + 10.0) })