qdrant-edge 0.7.2

A lightweight, in-process vector search engine designed for embedded devices, autonomous systems, and mobile agents.
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

use std::arch::aarch64::*;

#[target_feature(enable = "neon")]
#[allow(clippy::missing_safety_doc)]
pub unsafe fn neon_cosine_similarity_bytes(v1: &[u8], v2: &[u8]) -> f32 {
    debug_assert!(v1.len() == v2.len());
    let mut ptr1: *const u8 = v1.as_ptr();
    let mut ptr2: *const u8 = v2.as_ptr();

    unsafe {
        let mut mul1 = vdupq_n_u32(0);
        let mut mul2 = vdupq_n_u32(0);
        let mut norm11 = vdupq_n_u32(0);
        let mut norm12 = vdupq_n_u32(0);
        let mut norm21 = vdupq_n_u32(0);
        let mut norm22 = vdupq_n_u32(0);
        let len = v1.len();
        for _ in 0..len / 16 {
            let p1 = vld1q_u8(ptr1);
            let p2 = vld1q_u8(ptr2);
            ptr1 = ptr1.add(16);
            ptr2 = ptr2.add(16);

            let p1_low = vget_low_u8(p1);
            let p1_high = vget_high_u8(p1);
            let p2_low = vget_low_u8(p2);
            let p2_high = vget_high_u8(p2);

            let mul_low = vmull_u8(p1_low, p2_low);
            let mul_high = vmull_u8(p1_high, p2_high);
            mul1 = vpadalq_u16(mul1, mul_low);
            mul2 = vpadalq_u16(mul2, mul_high);

            let mul_low = vmull_u8(p1_low, p1_low);
            let mul_high = vmull_u8(p1_high, p1_high);
            norm11 = vpadalq_u16(norm11, mul_low);
            norm12 = vpadalq_u16(norm12, mul_high);

            let mul_low = vmull_u8(p2_low, p2_low);
            let mul_high = vmull_u8(p2_high, p2_high);
            norm21 = vpadalq_u16(norm21, mul_low);
            norm22 = vpadalq_u16(norm22, mul_high);
        }
        let mut dot_product = vaddvq_u32(vaddq_u32(mul1, mul2)) as f32;
        let mut norm1 = vaddvq_u32(vaddq_u32(norm11, norm12)) as f32;
        let mut norm2 = vaddvq_u32(vaddq_u32(norm21, norm22)) as f32;

        let remainder = len % 16;
        if remainder != 0 {
            let mut remainder_dot_product = 0;
            let mut remainder_norm1 = 0;
            let mut remainder_norm2 = 0;
            for _ in 0..remainder {
                let v1 = *ptr1;
                let v2 = *ptr2;
                ptr1 = ptr1.add(1);
                ptr2 = ptr2.add(1);
                remainder_dot_product += i32::from(v1) * i32::from(v2);
                remainder_norm1 += i32::from(v1) * i32::from(v1);
                remainder_norm2 += i32::from(v2) * i32::from(v2);
            }
            dot_product += remainder_dot_product as f32;
            norm1 += remainder_norm1 as f32;
            norm2 += remainder_norm2 as f32;
        }

        let denominator = norm1 * norm2;
        if denominator == 0.0 {
            return 0.0;
        }

        dot_product / denominator.sqrt()
    }
}

#[cfg(test)]
mod tests {
    use std::arch::is_aarch64_feature_detected;

    use super::*;
    use crate::segment::spaces::metric_uint::simple_cosine::cosine_similarity_bytes;

    #[test]
    fn test_spaces_neon() {
        if is_aarch64_feature_detected!("neon") {
            let v1: Vec<u8> = vec![
                255, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 255, 255,
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 255, 255, 0, 1, 2, 3,
                4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 255, 255, 0, 1, 2, 3, 4, 5, 6, 7,
                8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 255, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
                11, 12, 13, 14, 15, 16, 17,
            ];
            let v2: Vec<u8> = vec![
                255, 255, 0, 254, 253, 252, 251, 250, 249, 248, 247, 246, 245, 244, 243, 242, 241,
                240, 239, 238, 255, 255, 255, 254, 253, 252, 251, 250, 249, 248, 247, 246, 245,
                244, 243, 242, 241, 240, 239, 238, 255, 255, 255, 254, 253, 252, 251, 250, 249,
                248, 247, 246, 245, 244, 243, 242, 241, 240, 239, 238, 255, 255, 255, 254, 253,
                252, 251, 250, 249, 248, 247, 246, 245, 244, 243, 242, 241, 240, 239, 238, 255,
                255, 255, 254, 253, 252, 251, 250, 249, 248, 247, 246, 245, 244, 243, 242, 241,
                240, 239, 238,
            ];

            let dot_simd = unsafe { neon_cosine_similarity_bytes(&v1, &v2) };
            let dot = cosine_similarity_bytes(&v1, &v2);
            assert_eq!(dot_simd, dot);
        } else {
            println!("neon test skipped");
        }
    }

    #[test]
    fn test_zero_neon() {
        if is_aarch64_feature_detected!("neon") {
            let v1: Vec<u8> = vec![0, 0, 0, 0, 0, 0, 0, 0];
            let v2: Vec<u8> = vec![255, 255, 0, 254, 253, 252, 251, 250];

            let dot_simd = unsafe { neon_cosine_similarity_bytes(&v1, &v2) };
            assert_eq!(dot_simd, 0.0);

            let dot_simd = unsafe { neon_cosine_similarity_bytes(&v2, &v1) };
            assert_eq!(dot_simd, 0.0);

            let dot_simd = unsafe { neon_cosine_similarity_bytes(&v1, &v1) };
            assert_eq!(dot_simd, 0.0);
        } else {
            println!("neon test skipped");
        }
    }
}