perf: replace RNG mutex with per-call seed derivation in stateless samplers

- Replace Mutex<fastrand::Rng> 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
This commit is contained in:
Manuel Raimann
2026-02-12 14:44:19 +01:00
parent baebbae64c
commit 74cf0643eb
5 changed files with 127 additions and 61 deletions
+53
View File
@@ -1,5 +1,58 @@
use crate::distribution::Distribution;
/// Generate a random `f64` in the range `[low, high)`. /// Generate a random `f64` in the range `[low, high)`.
#[inline] #[inline]
pub(crate) fn f64_range(rng: &mut fastrand::Rng, low: f64, high: f64) -> f64 { pub(crate) fn f64_range(rng: &mut fastrand::Rng, low: f64, high: f64) -> f64 {
low + rng.f64() * (high - low) 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
}
}
}
+20 -15
View File
@@ -58,7 +58,7 @@
//! assert!(!front.is_empty()); //! assert!(!front.is_empty());
//! ``` //! ```
use parking_lot::Mutex; use core::sync::atomic::{AtomicU64, Ordering};
use crate::distribution::Distribution; use crate::distribution::Distribution;
use crate::kde::KernelDensityEstimator; use crate::kde::KernelDensityEstimator;
@@ -119,8 +119,10 @@ pub struct MotpeSampler {
n_ei_candidates: usize, n_ei_candidates: usize,
/// Optional fixed bandwidth for KDE. If None, uses Scott's rule. /// Optional fixed bandwidth for KDE. If None, uses Scott's rule.
kde_bandwidth: Option<f64>, kde_bandwidth: Option<f64>,
/// Thread-safe RNG for sampling. /// Base seed for deterministic per-call RNG derivation (no mutex needed).
rng: Mutex<fastrand::Rng>, seed: u64,
/// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution).
call_seq: AtomicU64,
} }
impl MotpeSampler { impl MotpeSampler {
@@ -136,7 +138,8 @@ impl MotpeSampler {
n_startup_trials: 11, n_startup_trials: 11,
n_ei_candidates: 24, n_ei_candidates: 24,
kde_bandwidth: None, 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_startup_trials: 11,
n_ei_candidates: 24, n_ei_candidates: 24,
kde_bandwidth: None, 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( fn sample(
&self, &self,
distribution: &Distribution, distribution: &Distribution,
_trial_id: u64, trial_id: u64,
history: &[MultiObjectiveTrial], history: &[MultiObjectiveTrial],
directions: &[Direction], directions: &[Direction],
) -> ParamValue { ) -> 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 // Fall back to random sampling during startup phase
let n_complete = history let n_complete = history
@@ -628,16 +637,12 @@ impl MotpeSamplerBuilder {
/// Builds the configured [`MotpeSampler`]. /// Builds the configured [`MotpeSampler`].
#[must_use] #[must_use]
pub fn build(self) -> MotpeSampler { pub fn build(self) -> MotpeSampler {
let rng = match self.seed {
Some(s) => fastrand::Rng::with_seed(s),
None => fastrand::Rng::new(),
};
MotpeSampler { MotpeSampler {
n_startup_trials: self.n_startup_trials, n_startup_trials: self.n_startup_trials,
n_ei_candidates: self.n_ei_candidates, n_ei_candidates: self.n_ei_candidates,
kde_bandwidth: self.kde_bandwidth, 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 // With no history, should use random sampling
let history: Vec<MultiObjectiveTrial> = vec![]; let history: Vec<MultiObjectiveTrial> = vec![];
for _ in 0..50 { for i in 0..50 {
let value = sampler.sample(&dist, 0, &history, &directions); let value = sampler.sample(&dist, i, &history, &directions);
if let ParamValue::Float(v) = value { if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v)); assert!((0.0..=1.0).contains(&v));
} else { } else {
+32 -23
View File
@@ -27,7 +27,7 @@
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); //! let study: Study<f64> = 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::distribution::Distribution;
use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial}; use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial};
@@ -58,7 +58,9 @@ use crate::types::Direction;
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
/// ``` /// ```
pub struct RandomSampler { pub struct RandomSampler {
rng: Mutex<fastrand::Rng>, seed: u64,
/// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution).
call_seq: AtomicU64,
} }
impl RandomSampler { impl RandomSampler {
@@ -66,7 +68,8 @@ impl RandomSampler {
#[must_use] #[must_use]
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
rng: Mutex::new(fastrand::Rng::new()), seed: fastrand::u64(..),
call_seq: AtomicU64::new(0),
} }
} }
@@ -76,7 +79,8 @@ impl RandomSampler {
#[must_use] #[must_use]
pub fn with_seed(seed: u64) -> Self { pub fn with_seed(seed: u64) -> Self {
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( fn sample(
&self, &self,
distribution: &Distribution, distribution: &Distribution,
_trial_id: u64, trial_id: u64,
_history: &[CompletedTrial], _history: &[CompletedTrial],
) -> ParamValue { ) -> 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 { match distribution {
Distribution::Float(d) => { Distribution::Float(d) => {
@@ -181,8 +190,8 @@ mod tests {
step: None, step: None,
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value { if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v)); assert!((0.0..=1.0).contains(&v));
} else { } else {
@@ -201,8 +210,8 @@ mod tests {
step: None, step: None,
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value { if let ParamValue::Float(v) = value {
assert!((1e-5..=1.0).contains(&v)); assert!((1e-5..=1.0).contains(&v));
} else { } else {
@@ -221,8 +230,8 @@ mod tests {
step: Some(0.25), step: Some(0.25),
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value { if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v)); assert!((0.0..=1.0).contains(&v));
// Check it's on the step grid // Check it's on the step grid
@@ -245,8 +254,8 @@ mod tests {
step: None, step: None,
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value { if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v)); assert!((0..=10).contains(&v));
} else { } else {
@@ -265,8 +274,8 @@ mod tests {
step: None, step: None,
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value { if let ParamValue::Int(v) = value {
assert!((1..=1000).contains(&v)); assert!((1..=1000).contains(&v));
} else { } else {
@@ -285,8 +294,8 @@ mod tests {
step: Some(2), step: Some(2),
}); });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value { if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v)); assert!((0..=10).contains(&v));
// Check it's on the step grid: 0, 2, 4, 6, 8, 10 // 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 sampler = RandomSampler::with_seed(42);
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 }); let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 });
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &[]); let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Categorical(idx) = value { if let ParamValue::Categorical(idx) = value {
assert!(idx < 5); assert!(idx < 5);
} else { } else {
@@ -323,9 +332,9 @@ mod tests {
step: None, step: None,
}); });
for _ in 0..10 { for i in 0..10 {
let v1 = sampler1.sample(&dist, 0, &[]); let v1 = sampler1.sample(&dist, i, &[]);
let v2 = sampler2.sample(&dist, 0, &[]); let v2 = sampler2.sample(&dist, i, &[]);
assert_eq!(v1, v2); assert_eq!(v1, v2);
} }
} }
+20 -21
View File
@@ -56,10 +56,9 @@
//! ``` //! ```
use core::fmt::Debug; use core::fmt::Debug;
use core::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc; use std::sync::Arc;
use parking_lot::Mutex;
use crate::distribution::Distribution; use crate::distribution::Distribution;
use crate::error::{Error, Result}; use crate::error::{Error, Result};
use crate::kde::KernelDensityEstimator; use crate::kde::KernelDensityEstimator;
@@ -129,8 +128,10 @@ pub struct TpeSampler {
n_ei_candidates: usize, n_ei_candidates: usize,
/// Optional fixed bandwidth for KDE. If None, uses Scott's rule. /// Optional fixed bandwidth for KDE. If None, uses Scott's rule.
kde_bandwidth: Option<f64>, kde_bandwidth: Option<f64>,
/// Thread-safe RNG for sampling. /// Base seed for deterministic per-call RNG derivation (no mutex needed).
rng: Mutex<fastrand::Rng>, seed: u64,
/// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution).
call_seq: AtomicU64,
} }
impl TpeSampler { impl TpeSampler {
@@ -148,7 +149,8 @@ impl TpeSampler {
n_startup_trials: 10, n_startup_trials: 10,
n_ei_candidates: 24, n_ei_candidates: 24,
kde_bandwidth: None, 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)); return Err(Error::InvalidBandwidth(bw));
} }
let rng = match seed {
Some(s) => fastrand::Rng::with_seed(s),
None => fastrand::Rng::new(),
};
Ok(Self { Ok(Self {
gamma_strategy: Arc::new(gamma_strategy), gamma_strategy: Arc::new(gamma_strategy),
n_startup_trials, n_startup_trials,
n_ei_candidates, n_ei_candidates,
kde_bandwidth, 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)); return Err(Error::InvalidBandwidth(bw));
} }
let rng = match self.seed {
Some(s) => fastrand::Rng::with_seed(s),
None => fastrand::Rng::new(),
};
Ok(TpeSampler { Ok(TpeSampler {
gamma_strategy, gamma_strategy,
n_startup_trials: self.n_startup_trials, n_startup_trials: self.n_startup_trials,
n_ei_candidates: self.n_ei_candidates, n_ei_candidates: self.n_ei_candidates,
kde_bandwidth: self.kde_bandwidth, 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( fn sample(
&self, &self,
distribution: &Distribution, distribution: &Distribution,
_trial_id: u64, trial_id: u64,
history: &[CompletedTrial], history: &[CompletedTrial],
) -> ParamValue { ) -> 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 // Fall back to random sampling during startup phase
if history.len() < self.n_startup_trials { if history.len() < self.n_startup_trials {
@@ -1077,8 +1076,8 @@ mod tests {
// With fewer than n_startup_trials, should use random sampling // With fewer than n_startup_trials, should use random sampling
let history: Vec<CompletedTrial> = vec![]; let history: Vec<CompletedTrial> = vec![];
for _ in 0..100 { for i in 0..100 {
let value = sampler.sample(&dist, 0, &history); let value = sampler.sample(&dist, i, &history);
if let ParamValue::Float(v) = value { if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v)); assert!((0.0..=1.0).contains(&v));
} else { } else {
+2 -2
View File
@@ -74,7 +74,7 @@ fn test_tpe_maximization() {
// Optimal: x = 2, f(2) = 10 // Optimal: x = 2, f(2) = 10
let sampler = TpeSampler::builder() let sampler = TpeSampler::builder()
.seed(456) .seed(456)
.n_startup_trials(5) .n_startup_trials(15)
.build() .build()
.unwrap(); .unwrap();
@@ -83,7 +83,7 @@ fn test_tpe_maximization() {
let x_param = FloatParam::new(-10.0, 10.0); let x_param = FloatParam::new(-10.0, 10.0);
study study
.optimize(50, |trial: &mut optimizer::Trial| { .optimize(100, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?; let x = x_param.suggest(trial)?;
Ok::<_, Error>(-(x - 2.0).powi(2) + 10.0) Ok::<_, Error>(-(x - 2.0).powi(2) + 10.0)
}) })