kevy_vector/
dist.rs

1//! Distance metrics + the wire vector format (RFC D1/D3).
2
3/// Distance metric. Scores are "smaller = closer" for every variant
4/// (cosine → `1 - cos`, ip → `-dot`), so one ascending merge works.
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
6pub enum Distance {
7    /// Cosine distance (vectors pre-normalized at insert).
8    #[default]
9    Cosine,
10    /// Squared euclidean.
11    L2,
12    /// Negative inner product.
13    Ip,
14}
15
16impl Distance {
17    /// Tag for sidecar round-trip.
18    pub fn tag(self) -> &'static str {
19        match self {
20            Distance::Cosine => "cosine",
21            Distance::L2 => "l2",
22            Distance::Ip => "ip",
23        }
24    }
25
26    /// Parse a tag (ASCII case-insensitive).
27    pub fn parse(raw: &[u8]) -> Option<Distance> {
28        if raw.eq_ignore_ascii_case(b"cosine") {
29            Some(Distance::Cosine)
30        } else if raw.eq_ignore_ascii_case(b"l2") {
31            Some(Distance::L2)
32        } else if raw.eq_ignore_ascii_case(b"ip") {
33            Some(Distance::Ip)
34        } else {
35            None
36        }
37    }
38
39    /// Distance between two prepared vectors (see [`prepare`]).
40    #[inline]
41    pub(crate) fn eval(self, a: &[f32], b: &[f32]) -> f32 {
42        match self {
43            // prepared cosine vectors are unit length → 1 - dot
44            Distance::Cosine => 1.0 - dot(a, b),
45            Distance::L2 => l2(a, b),
46            Distance::Ip => -dot(a, b),
47        }
48    }
49
50    /// Normalize a vector into its stored/query form (cosine only).
51    pub(crate) fn prepare(self, v: &mut [f32]) {
52        if self == Distance::Cosine {
53            let norm = dot(v, v).sqrt();
54            if norm > 0.0 {
55                for x in v.iter_mut() {
56                    *x /= norm;
57                }
58            }
59        }
60    }
61}
62
63// A single-accumulator float reduction is a loop-carried dependency
64// the compiler may NOT reassociate — it compiles to a SCALAR add
65// chain. Eight independent lanes let it auto-vectorize at whatever
66// SIMD width the target enables (perf-record put the distance kernel
67// inside the 65% beam-search dominator on the KNN shape).
68//
69// The portable baseline vectorizes at the DEFAULT target width
70// (SSE2 on x86_64) — still 48% of the KNN shape. On x86_64 a
71// runtime-dispatched AVX2+FMA clone of the SAME loops doubles the
72// lane width (the hnswlib/VecSim pattern): `#[target_feature]`
73// recompiles the body at the wider ISA, and the one-time
74// `is_x86_feature_detected!` gate makes the call sound.
75#[cfg(target_arch = "x86_64")]
76mod avx {
77    #[target_feature(enable = "avx2,fma")]
78    pub(super) unsafe fn dot(a: &[f32], b: &[f32]) -> f32 {
79        super::dot_lanes(a, b)
80    }
81
82    #[target_feature(enable = "avx2,fma")]
83    pub(super) unsafe fn l2(a: &[f32], b: &[f32]) -> f32 {
84        super::l2_lanes(a, b)
85    }
86}
87
88#[cfg(target_arch = "x86_64")]
89#[inline]
90fn has_avx2() -> bool {
91    use std::sync::atomic::{AtomicU8, Ordering};
92    static STATE: AtomicU8 = AtomicU8::new(0);
93    match STATE.load(Ordering::Relaxed) {
94        0 => {
95            let yes = std::arch::is_x86_feature_detected!("avx2")
96                && std::arch::is_x86_feature_detected!("fma");
97            STATE.store(if yes { 2 } else { 1 }, Ordering::Relaxed);
98            yes
99        }
100        s => s == 2,
101    }
102}
103
104#[inline]
105fn dot(a: &[f32], b: &[f32]) -> f32 {
106    #[cfg(target_arch = "x86_64")]
107    if has_avx2() {
108        // SAFETY: gated on runtime AVX2+FMA detection.
109        return unsafe { avx::dot(a, b) };
110    }
111    dot_lanes(a, b)
112}
113
114#[inline]
115fn l2(a: &[f32], b: &[f32]) -> f32 {
116    #[cfg(target_arch = "x86_64")]
117    if has_avx2() {
118        // SAFETY: gated on runtime AVX2+FMA detection.
119        return unsafe { avx::l2(a, b) };
120    }
121    l2_lanes(a, b)
122}
123
124#[inline(always)]
125fn dot_lanes(a: &[f32], b: &[f32]) -> f32 {
126    let mut acc = [0.0f32; 8];
127    let (ac, ar) = a.split_at(a.len() - a.len() % 8);
128    let (bc, br) = b.split_at(ac.len().min(b.len()));
129    for (ca, cb) in ac.chunks_exact(8).zip(bc.chunks_exact(8)) {
130        for i in 0..8 {
131            acc[i] += ca[i] * cb[i];
132        }
133    }
134    let mut s: f32 = acc.iter().sum();
135    for (x, y) in ar.iter().zip(br) {
136        s += x * y;
137    }
138    s
139}
140
141#[inline(always)]
142fn l2_lanes(a: &[f32], b: &[f32]) -> f32 {
143    let mut acc = [0.0f32; 8];
144    let (ac, ar) = a.split_at(a.len() - a.len() % 8);
145    let (bc, br) = b.split_at(ac.len().min(b.len()));
146    for (ca, cb) in ac.chunks_exact(8).zip(bc.chunks_exact(8)) {
147        for i in 0..8 {
148            let d = ca[i] - cb[i];
149            acc[i] += d * d;
150        }
151    }
152    let mut s: f32 = acc.iter().sum();
153    for (x, y) in ar.iter().zip(br) {
154        let d = x - y;
155        s += d * d;
156    }
157    s
158}
159
160/// Decode a wire vector: raw f32 LE bytes (`len == dim*4`), or the
161/// debug form `csv:1.0,2.5,…`. `None` on any mismatch.
162pub fn parse_vector(raw: &[u8], dim: usize) -> Option<Vec<f32>> {
163    if let Some(csv) = raw.strip_prefix(b"csv:") {
164        let vals: Option<Vec<f32>> = std::str::from_utf8(csv)
165            .ok()?
166            .split(',')
167            .map(|s| s.trim().parse::<f32>().ok())
168            .collect();
169        let vals = vals?;
170        return (vals.len() == dim && vals.iter().all(|x| x.is_finite())).then_some(vals);
171    }
172    if raw.len() != dim * 4 {
173        return None;
174    }
175    let mut out = Vec::with_capacity(dim);
176    for chunk in raw.chunks_exact(4) {
177        let x = f32::from_le_bytes(chunk.try_into().expect("4 bytes"));
178        if !x.is_finite() {
179            return None;
180        }
181        out.push(x);
182    }
183    Some(out)
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn metrics_smaller_is_closer() {
192        let mut a = vec![1.0, 0.0];
193        let mut b = vec![0.9, 0.1];
194        let mut c = vec![-1.0, 0.0];
195        for v in [&mut a, &mut b, &mut c] {
196            Distance::Cosine.prepare(v);
197        }
198        assert!(Distance::Cosine.eval(&a, &b) < Distance::Cosine.eval(&a, &c));
199        assert!(Distance::L2.eval(&[0.0, 0.0], &[1.0, 1.0]) > Distance::L2.eval(&[0.0, 0.0], &[0.5, 0.5]));
200        assert!(Distance::Ip.eval(&[1.0, 1.0], &[2.0, 2.0]) < Distance::Ip.eval(&[1.0, 1.0], &[0.1, 0.1]));
201    }
202
203    #[test]
204    fn wire_formats() {
205        let mut raw = Vec::new();
206        for x in [1.0f32, -2.5, 3.25] {
207            raw.extend_from_slice(&x.to_le_bytes());
208        }
209        assert_eq!(parse_vector(&raw, 3), Some(vec![1.0, -2.5, 3.25]));
210        assert_eq!(parse_vector(&raw, 4), None, "dim mismatch");
211        assert_eq!(parse_vector(b"csv:1, 2.5, -3", 3), Some(vec![1.0, 2.5, -3.0]));
212        assert_eq!(parse_vector(b"csv:1,x,3", 3), None);
213        let mut nan = Vec::new();
214        for x in [1.0f32, f32::NAN] {
215            nan.extend_from_slice(&x.to_le_bytes());
216        }
217        assert_eq!(parse_vector(&nan, 2), None, "non-finite rejected");
218        assert!(Distance::parse(b"COSINE") == Some(Distance::Cosine));
219        assert!(Distance::parse(b"nope").is_none());
220    }
221}
kevy_vector/dist.rs

kevy_vector/
dist.rs
