kannolo 0.1.6

kANNolo is designed for easy prototyping of ANN Search algorithms while ensuring high effectiveness and efficiency over both dense and sparse vectors.
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236

use crate::distances::dot_product_unrolled_avx;
use crate::simd_utils::horizontal_sum_256;
use crate::{AsRefItem, DArray1, DenseDArray1};
use half::f16;
use itertools::izip;
use std::arch::x86_64::*;
use std::iter::zip;

use crate::Float;

#[inline]
pub fn vectors_norm(vectors: &[f32], d: usize) -> Vec<f32> {
    vectors
        .chunks_exact(d)
        .map(|v| dot_product_unrolled_avx(v, v))
        .collect()
}

#[inline]
pub fn euclidean_distance_batch_4<T>(query: &[T], values: [&[T]; 4]) -> [f32; 4]
where
    T: Float + EuclideanDistance<T>,
{
    unsafe { T::euclidean_distance_batch_4(query, values) }
}

#[inline]
pub fn euclidean_distance_unrolled<T>(query: &[T], values: &[T]) -> f32
where
    T: Float + EuclideanDistance<T>,
{
    unsafe { T::euclidean_distance_unrolled(query, values) }
}

pub trait EuclideanDistance<U> {
    unsafe fn euclidean_distance_unrolled(query: &[U], values: &[U]) -> f32;
    fn euclidean_distance_general(query: &[U], values: &[U]) -> f32;
    unsafe fn euclidean_distance_batch_4(query: &[U], vectors: [&[U]; 4]) -> [f32; 4];
}

impl EuclideanDistance<f16> for f16 {
    #[inline]
    #[target_feature(enable = "avx2,fma,f16c")]
    unsafe fn euclidean_distance_batch_4(query: &[f16], vectors: [&[f16]; 4]) -> [f32; 4] {
        const N_LANES: usize = 8; // 8 half values per 128-bit register (8*16 = 128)

        // Zero out accumulators (each is a __m256, holding 8 f32s).
        let mut sum_0 = _mm256_setzero_ps();
        let mut sum_1 = _mm256_setzero_ps();
        let mut sum_2 = _mm256_setzero_ps();
        let mut sum_3 = _mm256_setzero_ps();

        for (q_chunk, v0_chunk, v1_chunk, v2_chunk, v3_chunk) in izip!(
            query.chunks_exact(N_LANES),
            vectors[0].chunks_exact(N_LANES),
            vectors[1].chunks_exact(N_LANES),
            vectors[2].chunks_exact(N_LANES),
            vectors[3].chunks_exact(N_LANES)
        ) {
            // Load 8 f16 values (16 bytes) into a __m128i and convert to f32.
            let q_values = _mm256_cvtph_ps(_mm_loadu_si128(q_chunk.as_ptr() as *const __m128i));
            let v0_values = _mm256_cvtph_ps(_mm_loadu_si128(v0_chunk.as_ptr() as *const __m128i));
            let v1_values = _mm256_cvtph_ps(_mm_loadu_si128(v1_chunk.as_ptr() as *const __m128i));
            let v2_values = _mm256_cvtph_ps(_mm_loadu_si128(v2_chunk.as_ptr() as *const __m128i));
            let v3_values = _mm256_cvtph_ps(_mm_loadu_si128(v3_chunk.as_ptr() as *const __m128i));

            // Compute differences.
            let diff_0 = _mm256_sub_ps(q_values, v0_values);
            let diff_1 = _mm256_sub_ps(q_values, v1_values);
            let diff_2 = _mm256_sub_ps(q_values, v2_values);
            let diff_3 = _mm256_sub_ps(q_values, v3_values);

            // Accumulate squared differences.
            sum_0 = _mm256_fmadd_ps(diff_0, diff_0, sum_0);
            sum_1 = _mm256_fmadd_ps(diff_1, diff_1, sum_1);
            sum_2 = _mm256_fmadd_ps(diff_2, diff_2, sum_2);
            sum_3 = _mm256_fmadd_ps(diff_3, diff_3, sum_3);
        }

        // Horizontally sum the 8 f32 lanes in each __m256 accumulator.
        let distance_0 = horizontal_sum_256(sum_0);
        let distance_1 = horizontal_sum_256(sum_1);
        let distance_2 = horizontal_sum_256(sum_2);
        let distance_3 = horizontal_sum_256(sum_3);

        [distance_0, distance_1, distance_2, distance_3]
    }

    #[inline]
    unsafe fn euclidean_distance_unrolled(query: &[f16], values: &[f16]) -> f32 {
        const N_LANES: usize = 16;
        let mut r = [f16::ZERO; N_LANES];

        let chunks = query.len() / N_LANES;
        for (q_chunk, v_chunk) in zip(query.chunks_exact(N_LANES), values.chunks_exact(N_LANES)) {
            for i in 0..N_LANES {
                let diff = q_chunk[i] - v_chunk[i];
                r[i] = r[i] + diff * diff;
            }
        }

        let result: f16 = r.iter().cloned().fold(f16::default(), |acc, val| acc + val);
        result.to_f32()
            + Self::euclidean_distance_general(
                &query[N_LANES * chunks..],
                &values[N_LANES * chunks..],
            )
    }

    #[inline]
    fn euclidean_distance_general(query: &[f16], values: &[f16]) -> f32 {
        query
            .iter()
            .zip(values)
            .fold(f16::default(), |acc: f16, (a, b)| {
                let diff = *a - *b;
                acc + diff * diff
            })
            .to_f32()
    }
}

impl EuclideanDistance<f32> for f32 {
    #[inline]
    unsafe fn euclidean_distance_unrolled(query: &[f32], values: &[f32]) -> f32 {
        const N_LANES: usize = 8;

        let mut sum = _mm256_setzero_ps();
        for (q_chunk, v_chunk) in zip(query.chunks_exact(N_LANES), values.chunks_exact(N_LANES)) {
            let q_values = _mm256_loadu_ps(q_chunk.as_ptr());
            let v_values = _mm256_loadu_ps(v_chunk.as_ptr());

            let diff = _mm256_sub_ps(q_values, v_values);

            sum = _mm256_fmadd_ps(diff, diff, sum);
        }

        horizontal_sum_256(sum)
    }

    #[inline]
    fn euclidean_distance_general(query: &[f32], values: &[f32]) -> f32 {
        query.iter().zip(values).fold(0.0, |acc, (a, b)| {
            let diff = a - b;
            acc + diff * diff
        })
    }

    #[inline]
    unsafe fn euclidean_distance_batch_4(query: &[f32], vectors: [&[f32]; 4]) -> [f32; 4] {
        const N_LANES: usize = 8;

        let mut sum_0 = _mm256_setzero_ps();
        let mut sum_1 = _mm256_setzero_ps();
        let mut sum_2 = _mm256_setzero_ps();
        let mut sum_3 = _mm256_setzero_ps();

        for (q_chunk, v0_chunk, v1_chunk, v2_chunk, v3_chunk) in izip!(
            query.chunks_exact(N_LANES),
            vectors[0].chunks_exact(N_LANES),
            vectors[1].chunks_exact(N_LANES),
            vectors[2].chunks_exact(N_LANES),
            vectors[3].chunks_exact(N_LANES)
        ) {
            let q_values = _mm256_loadu_ps(q_chunk.as_ptr());

            let v0_values = _mm256_loadu_ps(v0_chunk.as_ptr());
            let v1_values = _mm256_loadu_ps(v1_chunk.as_ptr());
            let v2_values = _mm256_loadu_ps(v2_chunk.as_ptr());
            let v3_values = _mm256_loadu_ps(v3_chunk.as_ptr());

            let diff_0 = _mm256_sub_ps(q_values, v0_values);
            let diff_1 = _mm256_sub_ps(q_values, v1_values);
            let diff_2 = _mm256_sub_ps(q_values, v2_values);
            let diff_3 = _mm256_sub_ps(q_values, v3_values);

            sum_0 = _mm256_fmadd_ps(diff_0, diff_0, sum_0);
            sum_1 = _mm256_fmadd_ps(diff_1, diff_1, sum_1);
            sum_2 = _mm256_fmadd_ps(diff_2, diff_2, sum_2);
            sum_3 = _mm256_fmadd_ps(diff_3, diff_3, sum_3);
        }

        let distance_0 = horizontal_sum_256(sum_0);
        let distance_1 = horizontal_sum_256(sum_1);
        let distance_2 = horizontal_sum_256(sum_2);
        let distance_3 = horizontal_sum_256(sum_3);

        [distance_0, distance_1, distance_2, distance_3]
    }
}

#[inline]
fn dense_euclidean_distance_general<T, U>(query: &DenseDArray1<T>, values: &DenseDArray1<T>) -> f32
where
    T: AsRefItem<Item = U>,
    U: Float,
{
    query
        .values_as_slice()
        .iter()
        .zip(values.values_as_slice())
        .fold(0.0, |acc, (a, b)| {
            let diff = a.to_f32().unwrap() - b.to_f32().unwrap();
            acc + diff * diff
        })
}

#[inline]
pub fn dense_euclidean_distance_unrolled<T, U>(
    query: &DenseDArray1<T>,
    values: &DenseDArray1<T>,
) -> f32
where
    T: AsRefItem<Item = U>,
    U: Float,
{
    const N_LANES: usize = 16;
    let mut r = [0.0; N_LANES];

    let chunks = query.len() / N_LANES;
    for (q_chunk, v_chunk) in zip(
        query.values_as_slice().chunks_exact(N_LANES),
        values.values_as_slice().chunks_exact(N_LANES),
    ) {
        for i in 0..N_LANES {
            let d = q_chunk[i].to_f32().unwrap() - v_chunk[i].to_f32().unwrap();
            r[i] += d * d;
        }
    }

    r.iter().fold(0.0, |acc, &val| acc + val)
        + dense_euclidean_distance_general(
            &DenseDArray1::new(&query.values_as_slice()[N_LANES * chunks..]),
            &DenseDArray1::new(&values.values_as_slice()[N_LANES * chunks..]),
        )
}