gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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:34:07 +01:00
parent baebbae64c
commit 74cf0643eb
5 changed files with 127 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/rng_util.rs
+53
View File
@@ -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
}
}
}
src/sampler/motpe.rs
+20 -15
View File
@@ -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<f64>,
/// Thread-safe RNG for sampling.
rng: Mutex<fastrand::Rng>,
/// 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<MultiObjectiveTrial> = 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 {
src/sampler/random.rs
+32 -23
View File
@@ -27,7 +27,7 @@
//! 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::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial};
@@ -58,7 +58,9 @@ use crate::types::Direction;
/// let sampler = RandomSampler::with_seed(42);
/// ```
pub struct RandomSampler {
rng: Mutex<fastrand::Rng>,
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);
}
}
src/sampler/tpe/sampler.rs
+20 -21
View File
@@ -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<f64>,
/// Thread-safe RNG for sampling.
rng: Mutex<fastrand::Rng>,
/// 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<CompletedTrial> = 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 {
tests/sampler/tpe.rs
+2 -2
View File
@@ -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)
})