GCC Code Coverage Report

Directory: src/
File: src/algo/distance_avx2.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 124 138 89.9%
Functions: 6 6 100.0%
Branches: 74 116 63.8%
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 102 vs_distance_l2_avx2(Vec32Ref a, Vec32Ref b)
27 {
28
6/12
✓ Branch 0 taken 102 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 102 times.
✓ Branch 4 taken 102 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 102 times.
✓ Branch 8 taken 102 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 102 times.
 102 if (vs_unlikely(
29 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
30 b.data == NULL))
31 ✗ return -1.0f;
32
33 102 const float *pa = a.data;
34 102 const float *pb = b.data;
35 102 Dimension dim = a.dim;
36
37 102 __m256 sum_vec = _mm256_setzero_ps();
38
39 /* Main loop: 8 floats per iteration */
40 102 Dimension i = 0;
41
2/2
✓ Branch 0 taken 2450 times.
✓ Branch 1 taken 102 times.
 2552 for (; i + 8 <= dim; i += 8)
42 {
43 2450 __m256 va = _mm256_loadu_ps(pa + i);
44 4900 __m256 vb = _mm256_loadu_ps(pb + i);
45 2450 __m256 diff = _mm256_sub_ps(va, vb);
46 2450 sum_vec = _mm256_fmadd_ps(diff, diff, sum_vec);
47 }
48
49 /* Horizontal reduction */
50 102 float sum = vs_horizontal_sum_avx2(sum_vec);
51
52 /* Scalar tail */
53
2/2
✓ Branch 0 taken 268 times.
✓ Branch 1 taken 102 times.
 370 for (; i < dim; i++)
54 {
55 268 float diff = pa[i] - pb[i];
56 268 sum += diff * diff;
57 }
58
59 102 return sum;
60 }
61
62 /*
63 * Negative inner product using AVX2.
64 */
65 VS_TARGET_AVX2 Distance
66 52 vs_distance_ip_avx2(Vec32Ref a, Vec32Ref b)
67 {
68
6/12
✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 52 times.
✓ Branch 4 taken 52 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 52 times.
✓ Branch 8 taken 52 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 52 times.
 52 if (vs_unlikely(
69 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
70 b.data == NULL))
71 ✗ return -1.0f;
72
73 52 const float *pa = a.data;
74 52 const float *pb = b.data;
75 52 Dimension dim = a.dim;
76
77 52 __m256 dot_vec = _mm256_setzero_ps();
78
79 52 Dimension i = 0;
80
2/2
✓ Branch 0 taken 1268 times.
✓ Branch 1 taken 52 times.
 1320 for (; i + 8 <= dim; i += 8)
81 {
82 1268 __m256 va = _mm256_loadu_ps(pa + i);
83 2536 __m256 vb = _mm256_loadu_ps(pb + i);
84 1268 dot_vec = _mm256_fmadd_ps(va, vb, dot_vec);
85 }
86
87 52 float dot = vs_horizontal_sum_avx2(dot_vec);
88
89
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 52 times.
 76 for (; i < dim; i++)
90 24 dot += pa[i] * pb[i];
91
92 52 return -dot;
93 }
94
95 /*
96 * Cosine distance using AVX2.
97 */
98 VS_TARGET_AVX2 Distance
99 48 vs_distance_cosine_avx2(Vec32Ref a, Vec32Ref b)
100 {
101
6/12
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 48 times.
✓ Branch 4 taken 48 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 48 times.
✓ Branch 8 taken 48 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 48 times.
 48 if (vs_unlikely(
102 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
103 b.data == NULL))
104 ✗ return -1.0f;
105
106 48 const float *pa = a.data;
107 48 const float *pb = b.data;
108 48 Dimension dim = a.dim;
109
110 48 __m256 dot_vec = _mm256_setzero_ps();
111 48 __m256 norm_a_vec = _mm256_setzero_ps();
112 48 __m256 norm_b_vec = _mm256_setzero_ps();
113
114 48 Dimension i = 0;
115
2/2
✓ Branch 0 taken 884 times.
✓ Branch 1 taken 48 times.
 932 for (; i + 8 <= dim; i += 8)
116 {
117 884 __m256 va = _mm256_loadu_ps(pa + i);
118 1768 __m256 vb = _mm256_loadu_ps(pb + i);
119 884 dot_vec = _mm256_fmadd_ps(va, vb, dot_vec);
120 884 norm_a_vec = _mm256_fmadd_ps(va, va, norm_a_vec);
121 884 norm_b_vec = _mm256_fmadd_ps(vb, vb, norm_b_vec);
122 }
123
124 48 float dot = vs_horizontal_sum_avx2(dot_vec);
125 48 float norm_a = vs_horizontal_sum_avx2(norm_a_vec);
126 48 float norm_b = vs_horizontal_sum_avx2(norm_b_vec);
127
128
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 48 times.
 68 for (; i < dim; i++)
129 {
130 20 float va = pa[i];
131 20 float vb = pb[i];
132 20 dot += va * vb;
133 20 norm_a += va * va;
134 20 norm_b += vb * vb;
135 }
136
137 48 float denom = sqrtf(norm_a) * sqrtf(norm_b);
138
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 48 times.
 48 if (denom < 1e-8f)
139 ✗ return 1.0f;
140
141 48 return 1.0f - (dot / denom);
142 }
143
144 /*
145 * Batch L2 distance with prefetching.
146 */
147 VS_TARGET_AVX2 int
148 2 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 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
 2 if (vs_unlikely(
156 query.dim != dim || query.data == NULL || vectors == NULL ||
157 distances == NULL))
158 ✗ return -1;
159
160 2 const float *q = query.data;
161
162
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 2 times.
 22 for (uint32_t v = 0; v < count; v++)
163 {
164 20 const float *vec = vectors + v * dim;
165
166
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 4 times.
 20 if (v + VS_PREFETCH_DISTANCE < count)
167 16 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
168
169 20 __m256 sum_vec = _mm256_setzero_ps();
170
171 20 Dimension i = 0;
172
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
 60 for (; i + 8 <= dim; i += 8)
173 {
174 40 __m256 vq = _mm256_loadu_ps(q + i);
175 80 __m256 vv = _mm256_loadu_ps(vec + i);
176 40 __m256 diff = _mm256_sub_ps(vq, vv);
177 40 sum_vec = _mm256_fmadd_ps(diff, diff, sum_vec);
178 }
179
180 20 float sum = vs_horizontal_sum_avx2(sum_vec);
181
182
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 for (; i < dim; i++)
183 {
184 ✗ float diff = q[i] - vec[i];
185 ✗ sum += diff * diff;
186 }
187
188 20 distances[v] = sum;
189 }
190
191 2 return 0;
192 }
193
194 /*
195 * Batch inner product with prefetching.
196 */
197 VS_TARGET_AVX2 int
198 2 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 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
 2 if (vs_unlikely(
206 query.dim != dim || query.data == NULL || vectors == NULL ||
207 distances == NULL))
208 ✗ return -1;
209
210 2 const float *q = query.data;
211
212
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 2 times.
 22 for (uint32_t v = 0; v < count; v++)
213 {
214 20 const float *vec = vectors + v * dim;
215
216
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 4 times.
 20 if (v + VS_PREFETCH_DISTANCE < count)
217 16 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
218
219 20 __m256 dot_vec = _mm256_setzero_ps();
220
221 20 Dimension i = 0;
222
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
 60 for (; i + 8 <= dim; i += 8)
223 {
224 40 __m256 vq = _mm256_loadu_ps(q + i);
225 80 __m256 vv = _mm256_loadu_ps(vec + i);
226 40 dot_vec = _mm256_fmadd_ps(vq, vv, dot_vec);
227 }
228
229 20 float dot = vs_horizontal_sum_avx2(dot_vec);
230
231
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 for (; i < dim; i++)
232 ✗ dot += q[i] * vec[i];
233
234 20 distances[v] = -dot;
235 }
236
237 2 return 0;
238 }
239
240 /*
241 * Batch cosine distance with prefetching.
242 */
243 VS_TARGET_AVX2 int
244 2 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 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
 2 if (vs_unlikely(
252 query.dim != dim || query.data == NULL || vectors == NULL ||
253 distances == NULL))
254 ✗ return -1;
255
256 2 const float *q = query.data;
257
258 /* Pre-compute query norm */
259 2 __m256 norm_q_vec = _mm256_setzero_ps();
260 2 Dimension i = 0;
261
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 for (; i + 8 <= dim; i += 8)
262 {
263 8 __m256 vq = _mm256_loadu_ps(q + i);
264 4 norm_q_vec = _mm256_fmadd_ps(vq, vq, norm_q_vec);
265 }
266 2 float norm_q = vs_horizontal_sum_avx2(norm_q_vec);
267
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 for (; i < dim; i++)
268 ✗ norm_q += q[i] * q[i];
269
270 2 float sqrt_norm_q = sqrtf(norm_q);
271
272
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 2 times.
 22 for (uint32_t v = 0; v < count; v++)
273 {
274 20 const float *vec = vectors + v * dim;
275
276
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 4 times.
 20 if (v + VS_PREFETCH_DISTANCE < count)
277 16 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim);
278
279 20 __m256 dot_vec = _mm256_setzero_ps();
280 20 __m256 norm_v_vec = _mm256_setzero_ps();
281
282 20 i = 0;
283
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
 60 for (; i + 8 <= dim; i += 8)
284 {
285 40 __m256 vq = _mm256_loadu_ps(q + i);
286 80 __m256 vv = _mm256_loadu_ps(vec + i);
287 40 dot_vec = _mm256_fmadd_ps(vq, vv, dot_vec);
288 40 norm_v_vec = _mm256_fmadd_ps(vv, vv, norm_v_vec);
289 }
290
291 20 float dot = vs_horizontal_sum_avx2(dot_vec);
292 20 float norm_v = vs_horizontal_sum_avx2(norm_v_vec);
293
294
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 for (; i < dim; i++)
295 {
296 ✗ dot += q[i] * vec[i];
297 ✗ norm_v += vec[i] * vec[i];
298 }
299
300 20 float denom = sqrt_norm_q * sqrtf(norm_v);
301
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 if (denom < 1e-8f)
302 ✗ distances[v] = 1.0f;
303 else
304 20 distances[v] = 1.0f - (dot / denom);
305 }
306
307 2 return 0;
308 }
309