needle-core 0.2.0

no_std compute kernels for Needle: Cactus-Quants matmul, fast Walsh-Hadamard transform, attention, RoPE
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Thin wrappers around libm for no_std float math.

#[inline(always)]
pub fn exp(x: f32) -> f32 {
    libm::expf(x)
}

#[inline(always)]
pub fn ln(x: f32) -> f32 {
    libm::logf(x)
}

#[inline(always)]
pub fn sin(x: f32) -> f32 {
    libm::sinf(x)
}

#[inline(always)]
pub fn cos(x: f32) -> f32 {
    libm::cosf(x)
}

#[inline(always)]
pub fn sqrt(x: f32) -> f32 {
    libm::sqrtf(x)
}

#[inline(always)]
pub fn tanh(x: f32) -> f32 {
    libm::tanhf(x)
}

#[inline(always)]
pub fn powf(base: f32, exp: f32) -> f32 {
    libm::powf(base, exp)
}

#[inline(always)]
pub fn round(x: f32) -> f32 {
    libm::roundf(x)
}

#[inline(always)]
pub fn abs(x: f32) -> f32 {
    libm::fabsf(x)
}

/// Convert IEEE 754 half-precision (binary16) bits to `f32`.
///
/// Exact for every input, including subnormals, infinities and NaN payloads.
/// Lives here so the `no_std` v2 loader and `cq` can share it.
#[inline]
pub fn f16_to_f32(bits: u16) -> f32 {
    let sign = ((bits >> 15) as u32) << 31;
    let exp = ((bits >> 10) & 0x1F) as u32;
    let mant = (bits & 0x3FF) as u32;

    if exp == 0 {
        // Subnormal (or zero): value = mant * 2^-24.
        let val = mant as f32 / (1u32 << 24) as f32;
        return if sign != 0 { -val } else { val };
    }
    if exp == 31 {
        // Inf or NaN — preserve the mantissa payload.
        return f32::from_bits(sign | 0x7F80_0000 | (mant << 13));
    }
    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mant << 13))
}

#[cfg(test)]
mod f16_tests {
    use super::f16_to_f32;

    #[test]
    fn known_values() {
        assert_eq!(f16_to_f32(0x0000), 0.0);
        assert_eq!(f16_to_f32(0x8000), -0.0);
        assert_eq!(f16_to_f32(0x3C00), 1.0);
        assert_eq!(f16_to_f32(0xBC00), -1.0);
        assert_eq!(f16_to_f32(0x4000), 2.0);
        assert_eq!(f16_to_f32(0x3555), 0.33325195); // nearest half to 1/3
        assert_eq!(f16_to_f32(0x7BFF), 65504.0); // largest finite half
        assert_eq!(f16_to_f32(0x0001), 5.9604645e-8); // smallest subnormal
        assert_eq!(f16_to_f32(0x03FF), 6.0975552e-5); // largest subnormal
        assert_eq!(f16_to_f32(0x0400), 6.1035156e-5); // smallest normal
        assert!(f16_to_f32(0x7C00).is_infinite() && f16_to_f32(0x7C00) > 0.0);
        assert!(f16_to_f32(0xFC00).is_infinite() && f16_to_f32(0xFC00) < 0.0);
        assert!(f16_to_f32(0x7E00).is_nan());
    }

    /// Round-trip against the f32 -> f16 direction for every representable half.
    #[test]
    fn every_finite_half_round_trips() {
        for bits in 0u32..=0xFFFF {
            let b = bits as u16;
            let v = f16_to_f32(b);
            if v.is_nan() {
                continue;
            }
            // f16 -> f32 is lossless, so re-narrowing must return the same bits.
            let back = f32_to_f16_bits(v);
            assert_eq!(back, b, "bits={b:#06x} value={v}");
        }
    }

    /// Round-to-nearest-even narrowing, test-only reference.
    fn f32_to_f16_bits(v: f32) -> u16 {
        let x = v.to_bits();
        let sign = ((x >> 31) as u16) << 15;
        let exp = ((x >> 23) & 0xFF) as i32;
        let mant = x & 0x7F_FFFF;
        if exp == 0xFF {
            return sign | 0x7C00 | ((mant >> 13) as u16);
        }
        let unbiased = exp - 127;
        if unbiased < -24 {
            return sign;
        }
        if unbiased < -14 {
            // Subnormal half: shift the implicit 1 back in.
            let shift = (-14 - unbiased) as u32;
            let m = (mant | 0x80_0000) >> (shift + 13);
            return sign | m as u16;
        }
        if unbiased > 15 {
            return sign | 0x7C00;
        }
        sign | (((unbiased + 15) as u16) << 10) | ((mant >> 13) as u16)
    }
}