GCC Code Coverage Report

Directory: src/
File: src/algo/distance_avx512.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 177 202 87.6%
Functions: 6 6 100.0%
Branches: 91 122 74.6%
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_avx512.c - AVX-512 SIMD distance implementations
6 *
7 * Explicit SIMD implementations using AVX-512 intrinsics for x86-64 CPUs
8 * with AVX-512F support. Processes 16 floats per iteration (512 bits).
9 *
10 * Uses per-function target attributes instead of compiler flags to enable
11 * AVX-512 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 AVX-512.
24 *
25 * Processes 16 floats per iteration using FMA (fused multiply-add).
26 * Remaining elements handled with scalar tail loop.
27 */
28 VS_TARGET_AVX512 Distance
29 729019 vs_distance_l2_avx512(Vec32Ref a, Vec32Ref b)
30 {
31
12/12
✓ Branch 0 taken 729017 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 195759 times.
✓ Branch 3 taken 533258 times.
✓ Branch 4 taken 729015 times.
✓ Branch 5 taken 4 times.
✓ Branch 6 taken 195759 times.
✓ Branch 7 taken 533256 times.
✓ Branch 8 taken 533256 times.
✓ Branch 9 taken 6 times.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 533254 times.
 729019 if (vs_unlikely(
32 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
33 b.data == NULL))
34 8 return -1.0f;
35
36 533254 const float *pa = a.data;
37 533254 const float *pb = b.data;
38 533254 Dimension dim = a.dim;
39
40 533254 __m512 sum_vec = _mm512_setzero_ps();
41
42 /* Main loop: 16 floats per iteration */
43 533254 Dimension i = 0;
44
2/2
✓ Branch 0 taken 6504650 times.
✓ Branch 1 taken 729011 times.
 7233661 for (; i + 16 <= dim; i += 16)
45 {
46 6504650 __m512 va = _mm512_loadu_ps(pa + i);
47 8001444 __m512 vb = _mm512_loadu_ps(pb + i);
48 6504650 __m512 diff = _mm512_sub_ps(va, vb);
49 6504650 sum_vec = _mm512_fmadd_ps(diff, diff, sum_vec);
50 }
51
52 /* Horizontal reduction */
53 729011 float sum = vs_horizontal_sum_avx512(sum_vec);
54
55 /* Scalar tail for remaining elements */
56
2/2
✓ Branch 0 taken 270961 times.
✓ Branch 1 taken 729011 times.
 999972 for (; i < dim; i++)
57 {
58 270961 float diff = pa[i] - pb[i];
59 270961 sum += diff * diff;
60 }
61
62 533254 return sum;
63 }
64
65 /*
66 * Negative inner product using AVX-512.
67 */
68 VS_TARGET_AVX512 Distance
69 96 vs_distance_ip_avx512(Vec32Ref a, Vec32Ref b)
70 {
71
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(
72 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
73 b.data == NULL))
74 2 return -1.0f;
75
76 90 const float *pa = a.data;
77 90 const float *pb = b.data;
78 90 Dimension dim = a.dim;
79
80 90 __m512 dot_vec = _mm512_setzero_ps();
81
82 90 Dimension i = 0;
83
2/2
✓ Branch 0 taken 736 times.
✓ Branch 1 taken 94 times.
 830 for (; i + 16 <= dim; i += 16)
84 {
85 736 __m512 va = _mm512_loadu_ps(pa + i);
86 1472 __m512 vb = _mm512_loadu_ps(pb + i);
87 736 dot_vec = _mm512_fmadd_ps(va, vb, dot_vec);
88 }
89
90 94 float dot = vs_horizontal_sum_avx512(dot_vec);
91
92
2/2
✓ Branch 0 taken 178 times.
✓ Branch 1 taken 94 times.
 272 for (; i < dim; i++)
93 178 dot += pa[i] * pb[i];
94
95 94 return -dot; /* Negate for distance metric */
96 }
97
98 /*
99 * Cosine distance using AVX-512.
100 *
101 * Computes three parallel reductions: dot product, norm_a, norm_b.
102 */
103 VS_TARGET_AVX512 Distance
104 16047 vs_distance_cosine_avx512(Vec32Ref a, Vec32Ref b)
105 {
106
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(
107 a.dim != b.dim || a.dim == 0 || a.data == NULL ||
108 b.data == NULL))
109 2 return -1.0f;
110
111 9302 const float *pa = a.data;
112 9302 const float *pb = b.data;
113 9302 Dimension dim = a.dim;
114
115 9302 __m512 dot_vec = _mm512_setzero_ps();
116 9302 __m512 norm_a_vec = _mm512_setzero_ps();
117 9302 __m512 norm_b_vec = _mm512_setzero_ps();
118
119 9302 Dimension i = 0;
120
2/2
✓ Branch 0 taken 7596 times.
✓ Branch 1 taken 16045 times.
 23641 for (; i + 16 <= dim; i += 16)
121 {
122 7596 __m512 va = _mm512_loadu_ps(pa + i);
123 14140 __m512 vb = _mm512_loadu_ps(pb + i);
124 7596 dot_vec = _mm512_fmadd_ps(va, vb, dot_vec);
125 7596 norm_a_vec = _mm512_fmadd_ps(va, va, norm_a_vec);
126 7596 norm_b_vec = _mm512_fmadd_ps(vb, vb, norm_b_vec);
127 }
128
129 16045 float dot = vs_horizontal_sum_avx512(dot_vec);
130 16045 float norm_a = vs_horizontal_sum_avx512(norm_a_vec);
131 16045 float norm_b = vs_horizontal_sum_avx512(norm_b_vec);
132
133 /* Scalar tail */
134
2/2
✓ Branch 0 taken 75241 times.
✓ Branch 1 taken 16045 times.
 91286 for (; i < dim; i++)
135 {
136 75241 float va = pa[i];
137 75241 float vb = pb[i];
138 75241 dot += va * vb;
139 75241 norm_a += va * va;
140 75241 norm_b += vb * vb;
141 }
142
143 16045 float denom = sqrtf(norm_a) * sqrtf(norm_b);
144
2/2
✓ Branch 0 taken 6749 times.
✓ Branch 1 taken 9296 times.
 16045 if (denom < 1e-8f)
145 6 return 1.0f; /* Maximum distance for zero vectors */
146
147 16039 return 1.0f - (dot / denom);
148 }
149
150 /*
151 * Batch L2 distance - simple single-accumulator loop.
152 */
153 VS_TARGET_AVX512 int
154 6 vs_distance_batch_l2_avx512(
155 Vec32Ref query,
156 const float *vectors,
157 uint32_t count,
158 Dimension dim,
159 Distance *distances)
160 {
161
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(
162 query.dim != dim || query.data == NULL || vectors == NULL ||
163 distances == NULL))
164 ✗ return -1;
165
166 6 const float *q = query.data;
167 6 const uint32_t dim32 = dim; /* Use 32-bit for efficient loop codegen */
168
169 /* Precompute loop bounds */
170 6 const size_t unrolled_bytes = (dim32 / 32) *
171 128; /* 32 floats = 128 bytes */
172 6 const uint32_t unrolled_elems = (dim32 / 32) * 32;
173
174
2/2
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 6 times.
 56 for (uint32_t v = 0; v < count; v++)
175 {
176 50 const float *vec = vectors + v * dim32;
177
178 50 __m512 sum0 = _mm512_setzero_ps();
179 50 __m512 sum1 = _mm512_setzero_ps();
180
181 /* 2x unrolled loop: 32 floats per iteration */
182 50 size_t offset = 0;
183
2/2
✓ Branch 0 taken 100 times.
✓ Branch 1 taken 50 times.
 150 while (offset < unrolled_bytes)
184 {
185 100 __m512 vq0 = _mm512_loadu_ps(
186 100 (const float *)((const char *)q + offset));
187 100 __m512 vv0 = _mm512_loadu_ps(
188 100 (const float *)((const char *)vec + offset));
189 100 __m512 diff0 = _mm512_sub_ps(vq0, vv0);
190 100 sum0 = _mm512_fmadd_ps(diff0, diff0, sum0);
191
192 100 __m512 vq1 = _mm512_loadu_ps(
193 100 (const float *)((const char *)q + offset + 64));
194 100 __m512 vv1 = _mm512_loadu_ps(
195 100 (const float *)((const char *)vec + offset + 64));
196 100 __m512 diff1 = _mm512_sub_ps(vq1, vv1);
197 100 sum1 = _mm512_fmadd_ps(diff1, diff1, sum1);
198
199 100 offset += 128;
200 }
201
202 /* Combine accumulators */
203
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 50 __m512 sum_vec = _mm512_add_ps(sum0, sum1);
204
205 /* Handle remaining 16-float chunk if dim not divisible by 32 */
206
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 40 times.
 50 if (unrolled_elems + 16 <= dim32)
207 {
208 10 __m512 vq = _mm512_loadu_ps(
209 10 (const float *)((const char *)q + offset));
210 10 __m512 vv = _mm512_loadu_ps(
211 10 (const float *)((const char *)vec + offset));
212 10 __m512 diff = _mm512_sub_ps(vq, vv);
213 10 sum_vec = _mm512_fmadd_ps(diff, diff, sum_vec);
214 }
215
216 50 float sum = vs_horizontal_sum_avx512(sum_vec);
217
218 /* Scalar tail */
219 50 uint32_t tail_start = (dim32 / 16) * 16;
220
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 50 times.
 60 for (uint32_t i = tail_start; i < dim32; i++)
221 {
222 10 float diff = q[i] - vec[i];
223 10 sum += diff * diff;
224 }
225
226 50 distances[v] = sum;
227 }
228
229 6 return 0;
230 }
231
232 /*
233 * Batch inner product - 2x unrolled loop.
234 */
235 VS_TARGET_AVX512 int
236 6 vs_distance_batch_ip_avx512(
237 Vec32Ref query,
238 const float *vectors,
239 uint32_t count,
240 Dimension dim,
241 Distance *distances)
242 {
243
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(
244 query.dim != dim || query.data == NULL || vectors == NULL ||
245 distances == NULL))
246 ✗ return -1;
247
248 6 const float *q = query.data;
249 6 const uint32_t dim32 = dim; /* Use 32-bit for efficient loop codegen */
250
251 /* Precompute loop bounds */
252 6 const size_t unrolled_bytes = (dim32 / 32) *
253 128; /* 32 floats = 128 bytes */
254 6 const uint32_t unrolled_elems = (dim32 / 32) * 32;
255
256
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 6 times.
 42 for (uint32_t v = 0; v < count; v++)
257 {
258 36 const float *vec = vectors + v * dim32;
259
260 36 __m512 dot0 = _mm512_setzero_ps();
261 36 __m512 dot1 = _mm512_setzero_ps();
262
263 /* 2x unrolled loop: 32 floats per iteration */
264 36 size_t offset = 0;
265
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 36 times.
 76 while (offset < unrolled_bytes)
266 {
267 40 __m512 vq0 = _mm512_loadu_ps(
268 40 (const float *)((const char *)q + offset));
269 40 __m512 vv0 = _mm512_loadu_ps(
270 40 (const float *)((const char *)vec + offset));
271 40 dot0 = _mm512_fmadd_ps(vq0, vv0, dot0);
272
273 40 __m512 vq1 = _mm512_loadu_ps(
274 40 (const float *)((const char *)q + offset + 64));
275 40 __m512 vv1 = _mm512_loadu_ps(
276 40 (const float *)((const char *)vec + offset + 64));
277 40 dot1 = _mm512_fmadd_ps(vq1, vv1, dot1);
278
279 40 offset += 128;
280 }
281
282 /* Combine accumulators */
283
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 36 __m512 dot_vec = _mm512_add_ps(dot0, dot1);
284
285 /* Handle remaining 16-float chunk if dim not divisible by 32 */
286
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 36 times.
 36 if (unrolled_elems + 16 <= dim32)
287 {
288 ✗ __m512 vq = _mm512_loadu_ps(
289 ✗ (const float *)((const char *)q + offset));
290 ✗ __m512 vv = _mm512_loadu_ps(
291 ✗ (const float *)((const char *)vec + offset));
292 ✗ dot_vec = _mm512_fmadd_ps(vq, vv, dot_vec);
293 }
294
295 36 float dot = vs_horizontal_sum_avx512(dot_vec);
296
297 /* Scalar tail */
298 36 uint32_t tail_start = (dim32 / 16) * 16;
299
2/2
✓ Branch 0 taken 78 times.
✓ Branch 1 taken 36 times.
 114 for (uint32_t i = tail_start; i < dim32; i++)
300 78 dot += q[i] * vec[i];
301
302 36 distances[v] = -dot;
303 }
304
305 6 return 0;
306 }
307
308 /*
309 * Batch cosine distance with 4 accumulators.
310 */
311 VS_TARGET_AVX512 int
312 6 vs_distance_batch_cosine_avx512(
313 Vec32Ref query,
314 const float *vectors,
315 uint32_t count,
316 Dimension dim,
317 Distance *distances)
318 {
319
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(
320 query.dim != dim || query.data == NULL || vectors == NULL ||
321 distances == NULL))
322 ✗ return -1;
323
324 6 const float *q = query.data;
325 6 const uint32_t dim32 = dim; /* Use 32-bit for efficient loop codegen */
326
327 /* Pre-compute query norm (reused for all vectors) */
328 6 __m512 norm_q_vec = _mm512_setzero_ps();
329 6 uint32_t i = 0;
330
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 6 times.
 16 for (; i + 16 <= dim32; i += 16)
331 {
332 20 __m512 vq = _mm512_loadu_ps(q + i);
333 10 norm_q_vec = _mm512_fmadd_ps(vq, vq, norm_q_vec);
334 }
335 6 float norm_q = vs_horizontal_sum_avx512(norm_q_vec);
336
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 6 times.
 24 for (; i < dim32; i++)
337 18 norm_q += q[i] * q[i];
338
339 6 float sqrt_norm_q = sqrtf(norm_q);
340
341
2/2
✓ Branch 0 taken 44 times.
✓ Branch 1 taken 6 times.
 50 for (uint32_t v = 0; v < count; v++)
342 {
343 44 const float *vec = vectors + v * dim32;
344
345
2/2
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 12 times.
 44 if (v + VS_PREFETCH_DISTANCE < count)
346 32 vs_prefetch_read(vectors + (v + VS_PREFETCH_DISTANCE) * dim32);
347
348 /* Use 4 accumulators to hide FMA latency */
349 44 __m512 dot0 = _mm512_setzero_ps();
350 44 __m512 dot1 = _mm512_setzero_ps();
351 44 __m512 dot2 = _mm512_setzero_ps();
352 44 __m512 dot3 = _mm512_setzero_ps();
353 44 __m512 norm_v0 = _mm512_setzero_ps();
354 44 __m512 norm_v1 = _mm512_setzero_ps();
355 44 __m512 norm_v2 = _mm512_setzero_ps();
356 44 __m512 norm_v3 = _mm512_setzero_ps();
357
358 44 i = 0;
359
360 /* Process 64 floats per iteration with 4 accumulators */
361
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 44 times.
 44 for (; i + 64 <= dim32; i += 64)
362 {
363 ✗ __m512 vq0 = _mm512_loadu_ps(q + i);
364 ✗ __m512 vv0 = _mm512_loadu_ps(vec + i);
365 ✗ dot0 = _mm512_fmadd_ps(vq0, vv0, dot0);
366 ✗ norm_v0 = _mm512_fmadd_ps(vv0, vv0, norm_v0);
367
368 ✗ __m512 vq1 = _mm512_loadu_ps(q + i + 16);
369 ✗ __m512 vv1 = _mm512_loadu_ps(vec + i + 16);
370 ✗ dot1 = _mm512_fmadd_ps(vq1, vv1, dot1);
371 ✗ norm_v1 = _mm512_fmadd_ps(vv1, vv1, norm_v1);
372
373 ✗ __m512 vq2 = _mm512_loadu_ps(q + i + 32);
374 ✗ __m512 vv2 = _mm512_loadu_ps(vec + i + 32);
375 ✗ dot2 = _mm512_fmadd_ps(vq2, vv2, dot2);
376 ✗ norm_v2 = _mm512_fmadd_ps(vv2, vv2, norm_v2);
377
378 ✗ __m512 vq3 = _mm512_loadu_ps(q + i + 48);
379 ✗ __m512 vv3 = _mm512_loadu_ps(vec + i + 48);
380 ✗ dot3 = _mm512_fmadd_ps(vq3, vv3, dot3);
381 ✗ norm_v3 = _mm512_fmadd_ps(vv3, vv3, norm_v3);
382 }
383
384 /* Combine accumulators */
385 132 __m512 dot_vec = _mm512_add_ps(
386 _mm512_add_ps(dot0, dot1), _mm512_add_ps(dot2, dot3));
387 88 __m512 norm_v_vec = _mm512_add_ps(
388 _mm512_add_ps(norm_v0, norm_v1),
389 _mm512_add_ps(norm_v2, norm_v3));
390
391 /* Handle remaining 16-float chunks */
392
2/2
✓ Branch 0 taken 80 times.
✓ Branch 1 taken 44 times.
 124 for (; i + 16 <= dim32; i += 16)
393 {
394 80 __m512 vq = _mm512_loadu_ps(q + i);
395 160 __m512 vv = _mm512_loadu_ps(vec + i);
396 80 dot_vec = _mm512_fmadd_ps(vq, vv, dot_vec);
397 80 norm_v_vec = _mm512_fmadd_ps(vv, vv, norm_v_vec);
398 }
399
400 44 float dot = vs_horizontal_sum_avx512(dot_vec);
401 44 float norm_v = vs_horizontal_sum_avx512(norm_v_vec);
402
403 /* Scalar tail */
404
2/2
✓ Branch 0 taken 72 times.
✓ Branch 1 taken 44 times.
 116 for (; i < dim32; i++)
405 {
406 72 dot += q[i] * vec[i];
407 72 norm_v += vec[i] * vec[i];
408 }
409
410 44 float denom = sqrt_norm_q * sqrtf(norm_v);
411
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 44 times.
 44 if (denom < 1e-8f)
412 ✗ distances[v] = 1.0f;
413 else
414 44 distances[v] = 1.0f - (dot / denom);
415 }
416
417 6 return 0;
418 }
419