GCC Code Coverage Report

Directory: src/
File: src/algo/distance_avx2.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 134 138 97.1%
Functions: 6 6 100.0%
Branches: 91 116 78.4%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * distance_avx2.c - AVX2 SIMD distance implementations
6 *
7 * Explicit SIMD implementations using AVX2 intrinsics for x86-64 CPUs.
8 * Processes 8 floats per iteration (256 bits).
9 *
10 * Uses per-function target attributes instead of compiler flags to enable
11 * AVX2 code generation. This allows the file to be compiled as part of
12 * the main build without requiring separate library compilation.
13 */
14
15 #include <immintrin.h>
16 #include <math.h>
17
18 #include "algo/distance.h"
19 #include "algo/simd_utils.h"
20 #include "core/platform.h"
21
22 /*
23 * L2 squared distance using AVX2.
24 */
25 VS_TARGET_AVX2 Distance
26 739877 vs_distance_l2_avx2(Vec32Ref a, Vec32Ref b)
27 {
28
12/12
✓ Branch 0 taken 739875 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 195757 times.
✓ Branch 3 taken 544118 times.
✓ Branch 4 taken 739873 times.
✓ Branch 5 taken 4 times.
✓ Branch 6 taken 195757 times.
✓ Branch 7 taken 544116 times.
✓ Branch 8 taken 544116 times.
✓ Branch 9 taken 6 times.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 544114 times.
 739877 if (vs_unlikely(
29 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
30 b.data == NULL))
31 8 return -1.0f;
32
33 544114 const float *pa = a.data;
34 544114 const float *pb = b.data;
35 544114 Dimension dim = a.dim;
36
37 544114 __m256 sum_vec = _mm256_setzero_ps();
38
39 /* Main loop: 8 floats per iteration */
40 544114 Dimension i = 0;
41
2/2
✓ Branch 0 taken 13161819 times.
✓ Branch 1 taken 739869 times.
 13901688 for (; i + 8 <= dim; i += 8)
42 {
43 13161819 __m256 va = _mm256_loadu_ps(pa + i);
44 16307101 __m256 vb = _mm256_loadu_ps(pb + i);
45 13161819 __m256 diff = _mm256_sub_ps(va, vb);
46 13161819 sum_vec = _mm256_fmadd_ps(diff, diff, sum_vec);
47 }
48
49 /* Horizontal reduction */
50 739869 float sum = vs_horizontal_sum_avx2(sum_vec);
51
52 /* Scalar tail */
53
2/2
✓ Branch 0 taken 262843 times.
✓ Branch 1 taken 739869 times.
 1002712 for (; i < dim; i++)
54 {
55 262843 float diff = pa[i] - pb[i];
56 262843 sum += diff * diff;
57 }
58
59 544114 return sum;
60 }
61
62 /*
63 * Negative inner product using AVX2.
64 */
65 VS_TARGET_AVX2 Distance
66 96 vs_distance_ip_avx2(Vec32Ref a, Vec32Ref b)
67 {
68
9/12
✓ Branch 0 taken 96 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 92 times.
✓ Branch 4 taken 96 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✓ Branch 7 taken 92 times.
✓ Branch 8 taken 92 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 90 times.
 96 if (vs_unlikely(
69 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
70 b.data == NULL))
71 2 return -1.0f;
72
73 90 const float *pa = a.data;
74 90 const float *pb = b.data;
75 90 Dimension dim = a.dim;
76
77 90 __m256 dot_vec = _mm256_setzero_ps();
78
79 90 Dimension i = 0;
80
2/2
✓ Branch 0 taken 1480 times.
✓ Branch 1 taken 94 times.
 1574 for (; i + 8 <= dim; i += 8)
81 {
82 1480 __m256 va = _mm256_loadu_ps(pa + i);
83 2960 __m256 vb = _mm256_loadu_ps(pb + i);
84 1480 dot_vec = _mm256_fmadd_ps(va, vb, dot_vec);
85 }
86
87 94 float dot = vs_horizontal_sum_avx2(dot_vec);
88
89
2/2
✓ Branch 0 taken 98 times.
✓ Branch 1 taken 94 times.
 192 for (; i < dim; i++)
90 98 dot += pa[i] * pb[i];
91
92 94 return -dot;
93 }
94
95 /*
96 * Cosine distance using AVX2.
97 */
98 VS_TARGET_AVX2 Distance
99 16047 vs_distance_cosine_avx2(Vec32Ref a, Vec32Ref b)
100 {
101
9/12
✓ Branch 0 taken 16047 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6743 times.
✓ Branch 3 taken 9304 times.
✓ Branch 4 taken 16047 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6745 times.
✓ Branch 7 taken 9302 times.
✓ Branch 8 taken 9302 times.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 9302 times.
 16047 if (vs_unlikely(
102 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
103 b.data == NULL))
104 2 return -1.0f;
105
106 9302 const float *pa = a.data;
107 9302 const float *pb = b.data;
108 9302 Dimension dim = a.dim;
109
110 9302 __m256 dot_vec = _mm256_setzero_ps();
111 9302 __m256 norm_a_vec = _mm256_setzero_ps();
112 9302 __m256 norm_b_vec = _mm256_setzero_ps();
113
114 9302 Dimension i = 0;
115
2/2
✓ Branch 0 taken 24565 times.
✓ Branch 1 taken 16045 times.
 40610 for (; i + 8 <= dim; i += 8)
116 {
117 24565 __m256 va = _mm256_loadu_ps(pa + i);
118 40863 __m256 vb = _mm256_loadu_ps(pb + i);
119 24565 dot_vec = _mm256_fmadd_ps(va, vb, dot_vec);
120 24565 norm_a_vec = _mm256_fmadd_ps(va, va, norm_a_vec);
121 24565 norm_b_vec = _mm256_fmadd_ps(vb, vb, norm_b_vec);
122 }
123
124 16045 float dot = vs_horizontal_sum_avx2(dot_vec);
125 16045 float norm_a = vs_horizontal_sum_avx2(norm_a_vec);
126 16045 float norm_b = vs_horizontal_sum_avx2(norm_b_vec);
127
128
2/2
✓ Branch 0 taken 241 times.
✓ Branch 1 taken 16045 times.
 16286 for (; i < dim; i++)
129 {
130 241 float va = pa[i];
131 241 float vb = pb[i];
132 241 dot += va * vb;
133 241 norm_a += va * va;
134 241 norm_b += vb * vb;
135 }
136
137 16045 float denom = sqrtf(norm_a) * sqrtf(norm_b);
138
2/2
✓ Branch 0 taken 6749 times.
✓ Branch 1 taken 9296 times.
 16045 if (denom < 1e-8f)
139 6 return 1.0f;
140
141 16039 return 1.0f - (dot / denom);
142 }
143
144 /*
145 * Batch L2 distance with prefetching.
146 */
147 VS_TARGET_AVX2 int
148 6 vs_distance_batch_l2_avx2(
149 Vec32Ref query,
150 const float *vectors,
151 uint32_t count,
152 Dimension dim,
153 Distance *distances)
154 {
155
6/12
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6 times.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 6 times.
 6 if (vs_unlikely(
156 query.dim != dim || query.data == NULL || vectors == NULL ||
157 distances == NULL))
158 ✗ return -1;
159
160 6 const float *q = query.data;
161
162
2/2
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 6 times.
 56 for (uint32_t v = 0; v < count; v++)
163 {
164 50 const float *vec = vectors + v * dim;
165
166
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 12 times.
 50 if (v + VS_PREFETCH_DISTANCE < count)
167 38 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
168
169 50 __m256 sum_vec = _mm256_setzero_ps();
170
171 50 Dimension i = 0;
172
2/2
✓ Branch 0 taken 380 times.
✓ Branch 1 taken 50 times.
 430 for (; i + 8 <= dim; i += 8)
173 {
174 380 __m256 vq = _mm256_loadu_ps(q + i);
175 760 __m256 vv = _mm256_loadu_ps(vec + i);
176 380 __m256 diff = _mm256_sub_ps(vq, vv);
177 380 sum_vec = _mm256_fmadd_ps(diff, diff, sum_vec);
178 }
179
180 50 float sum = vs_horizontal_sum_avx2(sum_vec);
181
182
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 50 times.
 60 for (; i < dim; i++)
183 {
184 10 float diff = q[i] - vec[i];
185 10 sum += diff * diff;
186 }
187
188 50 distances[v] = sum;
189 }
190
191 6 return 0;
192 }
193
194 /*
195 * Batch inner product with prefetching.
196 */
197 VS_TARGET_AVX2 int
198 6 vs_distance_batch_ip_avx2(
199 Vec32Ref query,
200 const float *vectors,
201 uint32_t count,
202 Dimension dim,
203 Distance *distances)
204 {
205
6/12
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6 times.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 6 times.
 6 if (vs_unlikely(
206 query.dim != dim || query.data == NULL || vectors == NULL ||
207 distances == NULL))
208 ✗ return -1;
209
210 6 const float *q = query.data;
211
212
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 6 times.
 42 for (uint32_t v = 0; v < count; v++)
213 {
214 36 const float *vec = vectors + v * dim;
215
216
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 12 times.
 36 if (v + VS_PREFETCH_DISTANCE < count)
217 24 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
218
219 36 __m256 dot_vec = _mm256_setzero_ps();
220
221 36 Dimension i = 0;
222
2/2
✓ Branch 0 taken 126 times.
✓ Branch 1 taken 36 times.
 162 for (; i + 8 <= dim; i += 8)
223 {
224 126 __m256 vq = _mm256_loadu_ps(q + i);
225 252 __m256 vv = _mm256_loadu_ps(vec + i);
226 126 dot_vec = _mm256_fmadd_ps(vq, vv, dot_vec);
227 }
228
229 36 float dot = vs_horizontal_sum_avx2(dot_vec);
230
231
2/2
✓ Branch 0 taken 30 times.
✓ Branch 1 taken 36 times.
 66 for (; i < dim; i++)
232 30 dot += q[i] * vec[i];
233
234 36 distances[v] = -dot;
235 }
236
237 6 return 0;
238 }
239
240 /*
241 * Batch cosine distance with prefetching.
242 */
243 VS_TARGET_AVX2 int
244 6 vs_distance_batch_cosine_avx2(
245 Vec32Ref query,
246 const float *vectors,
247 uint32_t count,
248 Dimension dim,
249 Distance *distances)
250 {
251
6/12
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6 times.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 6 times.
 6 if (vs_unlikely(
252 query.dim != dim || query.data == NULL || vectors == NULL ||
253 distances == NULL))
254 ✗ return -1;
255
256 6 const float *q = query.data;
257
258 /* Pre-compute query norm */
259 6 __m256 norm_q_vec = _mm256_setzero_ps();
260 6 Dimension i = 0;
261
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 6 times.
 24 for (; i + 8 <= dim; i += 8)
262 {
263 36 __m256 vq = _mm256_loadu_ps(q + i);
264 18 norm_q_vec = _mm256_fmadd_ps(vq, vq, norm_q_vec);
265 }
266 6 float norm_q = vs_horizontal_sum_avx2(norm_q_vec);
267
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 for (; i < dim; i++)
268 2 norm_q += q[i] * q[i];
269
270 6 float sqrt_norm_q = sqrtf(norm_q);
271
272
2/2
✓ Branch 0 taken 44 times.
✓ Branch 1 taken 6 times.
 50 for (uint32_t v = 0; v < count; v++)
273 {
274 44 const float *vec = vectors + v * dim;
275
276
2/2
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 12 times.
 44 if (v + VS_PREFETCH_DISTANCE < count)
277 32 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
278
279 44 __m256 dot_vec = _mm256_setzero_ps();
280 44 __m256 norm_v_vec = _mm256_setzero_ps();
281
282 44 i = 0;
283
2/2
✓ Branch 0 taken 128 times.
✓ Branch 1 taken 44 times.
 172 for (; i + 8 <= dim; i += 8)
284 {
285 128 __m256 vq = _mm256_loadu_ps(q + i);
286 256 __m256 vv = _mm256_loadu_ps(vec + i);
287 128 dot_vec = _mm256_fmadd_ps(vq, vv, dot_vec);
288 128 norm_v_vec = _mm256_fmadd_ps(vv, vv, norm_v_vec);
289 }
290
291 44 float dot = vs_horizontal_sum_avx2(dot_vec);
292 44 float norm_v = vs_horizontal_sum_avx2(norm_v_vec);
293
294
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 44 times.
 52 for (; i < dim; i++)
295 {
296 8 dot += q[i] * vec[i];
297 8 norm_v += vec[i] * vec[i];
298 }
299
300 44 float denom = sqrt_norm_q * sqrtf(norm_v);
301
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 44 times.
 44 if (denom < 1e-8f)
302 ✗ distances[v] = 1.0f;
303 else
304 44 distances[v] = 1.0f - (dot / denom);
305 }
306
307 6 return 0;
308 }
309