silvic_core/
quantize.rs

1//! Gradient quantization (LightGBM 4's `GradientDiscretizer`, `use_quantized_grad`) and the counter-based random
2//! numbers behind stochastic rounding and bagging.
3//!
4//! Every random draw is a pure function of (seed, stream, iteration, global row id), so it is the same on every
5//! backend and on whichever rank holds the row. The GPU kernels (`silvic-gpu/kernels/kernels.kern`) implement the
6//! same functions with the same operation order; [`mix32`], [`uniform`] and [`quantize_row`] are the contract they
7//! are tested against.
8
9/// Random-number streams; one per use so the draws for one row never correlate.
10pub const STREAM_GRADIENT: u32 = 0;
11pub const STREAM_HESSIAN: u32 = 1;
12pub const STREAM_BAGGING: u32 = 2;
13
14/// `lowbias32` (Chris Wellons' hash-prospector): a bijective 32-bit mixer with low bias.
15#[inline]
16pub fn mix32(mut x: u32) -> u32 {
17    x ^= x >> 16;
18    x = x.wrapping_mul(0x7feb_352d);
19    x ^= x >> 15;
20    x = x.wrapping_mul(0x846c_a68b);
21    x ^= x >> 16;
22    x
23}
24
25/// A uniform draw in `[0, 1)` with 24 bits of precision, exact in f32.
26#[inline]
27pub fn uniform(seed: u32, stream: u32, iteration: u32, row: u64) -> f32 {
28    let row_hash = mix32((row as u32) ^ mix32((row >> 32) as u32));
29    let h = mix32(seed ^ mix32(stream ^ mix32(iteration ^ row_hash)));
30    (h >> 8) as f32 * (1.0 / 16_777_216.0)
31}
32
33/// What one iteration's quantization needs; mirrored by `QuantParams` in the kernels.
34#[derive(Debug, Clone, Copy, PartialEq)]
35pub struct QuantSpec {
36    /// `1 / grad_scale`, where `grad_scale = max |g| / (num_grad_quant_bins / 2)` over every rank.
37    pub inv_grad_scale: f32,
38    /// `1 / hess_scale`, where `hess_scale = max h / num_grad_quant_bins`; unused with a constant hessian.
39    pub inv_hess_scale: f32,
40    pub grad_scale: f64,
41    pub hess_scale: f64,
42    /// `num_grad_quant_bins / 2`.
43    pub grad_max_q: i32,
44    /// `num_grad_quant_bins`.
45    pub hess_max_q: i32,
46    /// Every hessian equals `hess_scale`, so `h_q` is 1 and hessian sums are exact row counts.
47    pub constant_hessian: bool,
48    pub stochastic: bool,
49    pub seed: u32,
50    pub iteration: u32,
51}
52
53impl QuantSpec {
54    /// LightGBM `GradientDiscretizer::DiscretizeGradients` scales from the global maxima. All-zero gradients get
55    /// scale 1 so the inverse stays finite.
56    pub fn new(
57        max_abs_grad: f64,
58        max_hess: f64,
59        constant_hessian: Option<f64>,
60        num_bins: u32,
61        stochastic: bool,
62        seed: u32,
63        iteration: u32,
64    ) -> QuantSpec {
65        let grad_max_q = (num_bins / 2) as i32;
66        let hess_max_q = num_bins as i32;
67        let grad_scale = positive_or_one(max_abs_grad / f64::from(grad_max_q));
68        let hess_scale = match constant_hessian {
69            Some(value) => positive_or_one(value),
70            None => positive_or_one(max_hess / f64::from(hess_max_q)),
71        };
72        QuantSpec {
73            inv_grad_scale: (1.0 / grad_scale) as f32,
74            inv_hess_scale: (1.0 / hess_scale) as f32,
75            grad_scale,
76            hess_scale,
77            grad_max_q,
78            hess_max_q,
79            constant_hessian: constant_hessian.is_some(),
80            stochastic,
81            seed,
82            iteration,
83        }
84    }
85
86    /// Rows per GPU threadgroup so a packed i32 (two i16 halves) histogram bin never overflows.
87    pub fn max_rows_per_packed_block(&self) -> u32 {
88        (32_767 / self.grad_max_q.max(self.hess_max_q)) as u32
89    }
90}
91
92fn positive_or_one(value: f64) -> f64 {
93    if value > 0.0 && value.is_finite() {
94        value
95    } else {
96        1.0
97    }
98}
99
100/// Quantizes one row's gradient and hessian and packs them as `h_q * 65536 + g_q`, the layout the histogram
101/// kernels add in one 32-bit atomic.
102#[inline]
103pub fn quantize_row(spec: &QuantSpec, grad: f32, hess: f32, row: u64) -> i32 {
104    let g_q = quantize_grad(spec, grad, row);
105    let h_q = quantize_hess(spec, hess, row);
106    h_q * 65_536 + g_q
107}
108
109#[inline]
110fn quantize_grad(spec: &QuantSpec, grad: f32, row: u64) -> i32 {
111    let q = if spec.stochastic {
112        let u = uniform(spec.seed, STREAM_GRADIENT, spec.iteration, row);
113        if grad >= 0.0 {
114            grad.mul_add(spec.inv_grad_scale, u).trunc()
115        } else {
116            grad.mul_add(spec.inv_grad_scale, -u).trunc()
117        }
118    } else {
119        (grad * spec.inv_grad_scale).round_ties_even()
120    };
121    (q as i32).clamp(-spec.grad_max_q, spec.grad_max_q)
122}
123
124#[inline]
125fn quantize_hess(spec: &QuantSpec, hess: f32, row: u64) -> i32 {
126    if spec.constant_hessian {
127        return 1;
128    }
129    let q = if spec.stochastic {
130        let u = uniform(spec.seed, STREAM_HESSIAN, spec.iteration, row);
131        hess.mul_add(spec.inv_hess_scale, u).trunc()
132    } else {
133        (hess * spec.inv_hess_scale).round_ties_even()
134    };
135    (q as i32).clamp(0, spec.hess_max_q)
136}
137
138/// Splits a packed value back into `(g_q, h_q)`.
139#[inline]
140pub fn unpack(packed: i32) -> (i32, i32) {
141    let g = i32::from(packed as u32 as u16 as i16);
142    (g, (packed - g) >> 16)
143}
144
145/// Whether `row` is in the bag of `bag_round` (LightGBM bags rows independently, also for lambdarank, in 3.3.5).
146#[inline]
147pub fn in_bag(bagging_seed: u32, bag_round: u32, row: u64, fraction: f64) -> bool {
148    f64::from(uniform(bagging_seed, STREAM_BAGGING, bag_round, row)) < fraction
149}
150
151/// The power-of-two scale for exact fixed-point leaf sums: `value * scale` is exact in f32 and at most `2^30` in
152/// magnitude for every `|value| <= max_abs`, so 2^33 rows still sum inside i64.
153pub fn fixed_point_scale(max_abs: f64) -> f32 {
154    if !max_abs.is_finite() || max_abs <= 0.0 {
155        return 1.0;
156    }
157    let exponent = 30 - max_abs.log2().ceil() as i32;
158    2f32.powi(exponent.clamp(-126, 127))
159}
160
161/// `rint(value * scale)` as an i64, matching the kernels' `(i64)F32_RINT(v * scale)`.
162#[inline]
163pub fn to_fixed(value: f32, scale: f32) -> i64 {
164    (value * scale).round_ties_even() as i64
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn uniform_is_in_unit_interval_and_row_keyed() {
173        let draws: Vec<f32> = (0..10_000).map(|row| uniform(7, STREAM_GRADIENT, 3, row)).collect();
174        assert!(draws.iter().all(|u| (0.0..1.0).contains(u)));
175        let mean = draws.iter().map(|&u| f64::from(u)).sum::<f64>() / draws.len() as f64;
176        assert!((mean - 0.5).abs() < 0.02, "mean {mean}");
177        assert_ne!(
178            uniform(7, STREAM_GRADIENT, 3, 1 << 32),
179            uniform(7, STREAM_GRADIENT, 3, 0)
180        );
181    }
182
183    #[test]
184    fn stochastic_rounding_is_unbiased() {
185        let spec = QuantSpec::new(2.0, 1.0, None, 4, true, 11, 0);
186        let grad = 0.3f32;
187        let sum: i64 = (0..200_000u64)
188            .map(|row| i64::from(unpack(quantize_row(&spec, grad, 0.5, row)).0))
189            .sum();
190        let mean = sum as f64 / 200_000.0 * spec.grad_scale;
191        assert!((mean - 0.3).abs() < 0.01, "mean {mean}");
192    }
193
194    #[test]
195    fn pack_round_trips_negative_gradients() {
196        for g in -32..=32 {
197            for h in 0..=64 {
198                assert_eq!(unpack(h * 65_536 + g), (g, h));
199            }
200        }
201    }
202
203    #[test]
204    fn constant_hessian_counts_rows() {
205        let spec = QuantSpec::new(1.0, 1.0, Some(1.0), 4, true, 1, 0);
206        assert_eq!(unpack(quantize_row(&spec, -0.7, 1.0, 5)).1, 1);
207    }
208
209    #[test]
210    fn fixed_point_scale_bounds_magnitude() {
211        let scale = fixed_point_scale(3.7);
212        assert!(f64::from(3.7f32 * scale) <= f64::from(1u32 << 30));
213        assert!(f64::from(3.7f32 * scale) > f64::from(1u32 << 29));
214        assert_eq!(fixed_point_scale(0.0), 1.0);
215    }
216}
silvic_core/quantize.rs

silvic_core/
quantize.rs
