embedded_nn/
support.rs

1//! Fixed-point math and support helper functions for `embedded-nn`.
2
3/// Clamps a scalar `x` between `min` and `max`.
4#[inline]
5pub const fn clamp(x: i32, min: i32, max: i32) -> i32 {
6    if x < min {
7        min
8    } else if x > max {
9        max
10    } else {
11        x
12    }
13}
14
15/// Doubling high multiply without saturation.
16///
17/// Computes `(2 * m1 * m2 + 0x4000_0000) >> 31`.
18/// Equivalent to Arm CMSIS-NN `arm_nn_doubling_high_mult_no_sat`.
19#[inline]
20pub fn doubling_high_mult_no_sat(m1: i32, m2: i32) -> i32 {
21    let mult = (m1 as i64) * (m2 as i64) + (1i64 << 30);
22    (mult >> 31) as i32
23}
24
25/// Rounding divide by power of two (midpoint away from zero).
26///
27/// Equivalent to Arm CMSIS-NN `arm_nn_divide_by_power_of_two`.
28#[inline]
29pub fn divide_by_power_of_two(dividend: i32, exponent: i32) -> i32 {
30    if exponent <= 0 {
31        return dividend;
32    }
33    if exponent >= 31 {
34        return 0;
35    }
36
37    let remainder_mask = (1i32 << exponent) - 1;
38    let remainder = remainder_mask & dividend;
39
40    let mut result = dividend >> exponent;
41    let mut threshold = remainder_mask >> 1;
42    if dividend < 0 {
43        threshold += 1;
44    }
45    if remainder > threshold {
46        result += 1;
47    }
48    result
49}
50
51/// Requantizes a 32-bit integer value given a multiplier and shift.
52///
53/// Equivalent to Arm CMSIS-NN `arm_nn_requantize`.
54#[inline]
55pub fn requantize(val: i32, multiplier: i32, shift: i32) -> i32 {
56    if shift >= 0 {
57        let val_shifted = val.wrapping_shl(shift as u32);
58        doubling_high_mult_no_sat(val_shifted, multiplier)
59    } else {
60        let right_shift = -shift;
61        let mult = doubling_high_mult_no_sat(val, multiplier);
62        divide_by_power_of_two(mult, right_shift)
63    }
64}
65
66/// Requantizes a 64-bit accumulator value.
67///
68/// Equivalent to Arm CMSIS-NN `arm_nn_requantize_s64`.
69#[inline]
70pub fn requantize_s64(val: i64, reduced_multiplier: i32, shift: i32) -> i32 {
71    let new_val = val * (reduced_multiplier as i64);
72    let shift_amt = 14 - shift;
73    let mut result = (new_val >> (shift_amt - 1)) as i32;
74    result = (result + 1) >> 1;
75    result
76}
77
78/// Packs four 8-bit signed integers into a single 32-bit integer.
79#[inline]
80pub const fn pack_s8x4_32x1(v0: i8, v1: i8, v2: i8, v3: i8) -> i32 {
81    ((v0 as u8 as i32) & 0xFF)
82        | (((v1 as u8 as i32) << 8) & 0xFF00)
83        | (((v2 as u8 as i32) << 16) & 0xFF0000)
84        | (((v3 as u8 as i32) << 24) as i32)
85}
86
87/// Packs two 16-bit signed integers into a single 32-bit integer.
88#[inline]
89pub const fn pack_q15x2_32x1(v0: i16, v1: i16) -> i32 {
90    ((v0 as u16 as i32) & 0xFFFF) | ((v1 as u16 as i32) << 16)
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn test_clamp() {
99        assert_eq!(clamp(10, 0, 5), 5);
100        assert_eq!(clamp(-5, 0, 5), 0);
101        assert_eq!(clamp(3, 0, 5), 3);
102    }
103
104    #[test]
105    fn test_doubling_high_mult() {
106        let res = doubling_high_mult_no_sat(1073741824, 1073741824); // 0.5 * 0.5 in Q31
107        assert!((res - 536870912).abs() <= 1);
108    }
109
110    #[test]
111    fn test_divide_by_power_of_two() {
112        assert_eq!(divide_by_power_of_two(16, 2), 4);
113        assert_eq!(divide_by_power_of_two(18, 2), 5); // 18 / 4 = 4.5 -> rounds to 5
114    }
115
116    #[test]
117    fn test_requantize() {
118        let val = 1000;
119        let mult = 1073741824; // 0.5
120        let shift = -1;
121        let res = requantize(val, mult, shift);
122        assert_eq!(res, 250);
123    }
124}
embedded_nn/support.rs

embedded_nn/
support.rs
