diff --git a/diskann-wide/src/test_utils/distribution.rs b/diskann-wide/src/test_utils/distribution.rs index 1c166f3e9..17af37f46 100644 --- a/diskann-wide/src/test_utils/distribution.rs +++ b/diskann-wide/src/test_utils/distribution.rs @@ -19,6 +19,17 @@ use rand::{ distr::{Distribution, StandardUniform}, }; +// Must be a power of two to avoid bias when reducing a uniformly distributed integer. +const TOTAL_WEIGHT: u64 = 128; +const NORMAL_WEIGHT: u64 = 116; +const SUBNORMAL_WEIGHT: u64 = 6; +const ZERO_WEIGHT: u64 = 6; + +const _: () = assert!( + NORMAL_WEIGHT + SUBNORMAL_WEIGHT + ZERO_WEIGHT == TOTAL_WEIGHT, + "floating point weights must sum to the total weight" +); + trait Layout { type Bits; @@ -53,7 +64,7 @@ impl Layout for f32 { pub struct Finite; macro_rules! finite { - ($T:ty, $bits:ty) => { + ($T:ty, $bits:ty, $twice:ty) => { impl Distribution<$T> for Finite { /// Generate floating point numbers spread more-or-less uniformly across the /// distribution of floating point numbers. @@ -68,13 +79,23 @@ macro_rules! finite { /// /// This function does not generate infinities or NaNs. fn sample(&self, rng: &mut R) -> $T { - // Generate a uniformly distributed 32-bit integer - let mut value: $bits = StandardUniform {}.sample(rng); - - // The distribution from which we sample weights to determine the type of - // floating point number we are going to generate. - let weight = value % 100; - let (mask, allow_edge_exponent, allow_zero_mantissa) = if weight < 90 { + // Generate a uniformly distributed integer. + // + // This integer is twice as large as what's actually needed to generate + // + // * the value that will be used to make the final floating point number + // (the lower bits). + // + // * a selector for the kind of floating point number we are going to + // generate (the upper bits). + // + // Generating a number twice as big allows us to perform just a single sample + // from the random number generator without biasing the result. + let twice: $twice = StandardUniform {}.sample(rng); + + let mut value = twice as $bits; + let kind = u64::from(twice >> <$bits>::BITS) % TOTAL_WEIGHT; + let (mask, allow_edge_exponent, allow_zero_mantissa) = if kind < NORMAL_WEIGHT { // Generate a normal floating point number. // // All digits are fair game, but the exponent cannot be all zeros @@ -82,12 +103,8 @@ macro_rules! finite { // infinity/NaN). // // The mantissa is allowed to be all zeros. - ( - <$T as Layout>::EXPONENT_MASK | <$T as Layout>::MANTISSA_MASK, - false, - true, - ) - } else if weight < 95 { + (<$T>::EXPONENT_MASK | <$T>::MANTISSA_MASK, false, true) + } else if kind < NORMAL_WEIGHT + SUBNORMAL_WEIGHT { // Generate a subnormal floating point number. // // The exponent must be all zero and the mantissa cannot be zero. @@ -121,8 +138,8 @@ macro_rules! finite { }; } -finite!(half::f16, u16); -finite!(f32, u32); +finite!(half::f16, u16, u32); +finite!(f32, u32, u64); /////////// // Tests // @@ -137,13 +154,13 @@ mod tests { #[derive(Debug, Default)] struct Kinds { - normal: i64, - subnormal: i64, - zero: i64, + normal: u64, + subnormal: u64, + zero: u64, } impl Kinds { - fn sum(&self) -> i64 { + fn sum(&self) -> u64 { self.normal + self.subnormal + self.zero } } @@ -235,12 +252,7 @@ mod tests { where T: TestDistribution, { - let normal_weight = 90; - let subnormal_weight = 5; - let zero_weight = 5; - let total_weight = normal_weight + subnormal_weight + zero_weight; - - let num_trials: i64 = 1_000_000; + let num_trials: u64 = 1_000_000; let margin = num_trials / 500; let counts = T::test_distribution(num_trials as usize, seed); @@ -249,18 +261,34 @@ mod tests { println!("Counts = {:?}", counts); - assert!((positive_count - num_trials / 2).abs() < margin); - assert!((negative_count - num_trials / 2).abs() < margin); + assert!(positive_count.abs_diff(num_trials / 2) < margin); + assert!(negative_count.abs_diff(num_trials / 2) < margin); - assert!((counts.positive.normal - counts.negative.normal).abs() < margin); - assert!((counts.positive.subnormal - counts.negative.subnormal).abs() < margin); - assert!((counts.positive.zero - counts.negative.zero).abs() < margin); + assert!(counts.positive.normal.abs_diff(counts.negative.normal) < margin); + assert!( + counts + .positive + .subnormal + .abs_diff(counts.negative.subnormal) + < margin + ); + assert!(counts.positive.zero.abs_diff(counts.negative.zero) < margin); let kinds = counts.sum_accross(); - assert!((kinds.normal - num_trials * normal_weight / total_weight).abs() < margin); - assert!((kinds.subnormal - num_trials * subnormal_weight / total_weight).abs() < margin); - assert!((kinds.zero - num_trials * zero_weight / total_weight).abs() < margin); + assert!( + kinds + .normal + .abs_diff(num_trials * NORMAL_WEIGHT / TOTAL_WEIGHT) + < margin + ); + assert!( + kinds + .subnormal + .abs_diff(num_trials * SUBNORMAL_WEIGHT / TOTAL_WEIGHT) + < margin + ); + assert!(kinds.zero.abs_diff(num_trials * ZERO_WEIGHT / TOTAL_WEIGHT) < margin); } #[test]