GCC Code Coverage Report

Directory: src/
File: src/quant/rabitq_avx512.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 94 94 100.0%
Functions: 5 5 100.0%
Branches: 36 36 100.0%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * rabitq_avx512.c - AVX-512 optimized binary inner product for RaBitQ
6 *
7 * AVX-512 provides native masking support which makes this operation
8 * extremely efficient. We can process 16 floats at a time using a
9 * 16-bit mask directly.
10 */
11
12 #include "vs_config.h"
13
14 #ifdef VS_SIMD_FULL
15
16 #if defined(__x86_64__) || defined(_M_X64)
17
18 #include <immintrin.h>
19
20 #include "algo/simd_utils.h"
21 #include "core/types.h"
22 #include "quant/rabitq.h"
23
24 /*
25 * AVX-512 binary inner product implementation.
26 *
27 * Processes 16 floats at a time using native AVX-512 masking.
28 * The mask register directly controls which elements participate
29 * in the addition, making this very efficient.
30 */
31 VS_TARGET_AVX512 float
32 548427 vs_rabitq_inner_product_avx512(
33 const float *transformed, const uint8_t *bits, Dimension dim)
34 {
35 548427 __m512 sum = _mm512_setzero_ps();
36
37 /* Main loop: process 16 floats (2 bytes of bits) at a time */
38 548427 Dimension i = 0;
39
2/2
✓ Branch 0 taken 1251494 times.
✓ Branch 1 taken 548427 times.
 1799921 for (; i + 16 <= dim; i += 16)
40 {
41 /* Load 16 transformed values */
42 1251494 __m512 t = _mm512_loadu_ps(transformed + i);
43
44 /* Get 2 bytes containing bits for these 16 floats (LSB-first) */
45 /* bits[i/8] contains bits for i..i+7, bits[i/8+1] for i+8..i+15 */
46 /* With LSB-first storage, bit order matches AVX-512 mask directly */
47 1251494 uint8_t byte0 = bits[i / 8]; /* bits for i+0..i+7 */
48 1251494 uint8_t byte1 = bits[i / 8 + 1]; /* bits for i+8..i+15 */
49
50 /* Combine into 16-bit mask (byte0 = low 8 bits, byte1 = high 8 bits)
51 */
52 1251494 __mmask16 k = (__mmask16)((uint16_t)byte1 << 8 | byte0);
53
54 /* Native masked add: only adds elements where mask bit is 1 */
55 2060286 sum = _mm512_mask_add_ps(sum, k, sum, t);
56 }
57
58 /* Horizontal sum */
59 548427 float result = vs_horizontal_sum_avx512(sum);
60
61 /* Scalar tail for remaining elements */
62
63 /* Scalar tail (LSB-first) */
64
2/2
✓ Branch 0 taken 1358641 times.
✓ Branch 1 taken 548427 times.
 1907068 for (; i < dim; i++)
65 {
66 1358641 int byte_idx = i / 8;
67 1358641 int bit_idx = i % 8;
68 1358641 int bit = (bits[byte_idx] >> bit_idx) & 1;
69
2/2
✓ Branch 0 taken 694335 times.
✓ Branch 1 taken 664306 times.
 1358641 if (bit)
70 694335 result += transformed[i];
71 }
72
73 548427 return result;
74 }
75
76 /*
77 * AVX-512 sign extraction for RaBitQ encoding.
78 *
79 * Extracts sign bits from transformed floats into packed bytes.
80 * Uses AVX-512 compare to generate a 16-bit mask directly.
81 */
82 VS_TARGET_AVX512 void
83 694422 vs_rabitq_extract_signs_avx512(
84 const float *transformed, uint8_t *bits, Dimension dim)
85 {
86 694422 __m512 zero = _mm512_setzero_ps();
87
88 /* Main loop: process 16 floats -> 2 bytes at a time */
89 694422 Dimension i = 0;
90
2/2
✓ Branch 0 taken 2424701 times.
✓ Branch 1 taken 694422 times.
 3119123 for (; i + 16 <= dim; i += 16)
91 {
92 /* Load 16 floats */
93 2424701 __m512 t = _mm512_loadu_ps(transformed + i);
94
95 /* Compare > 0, returns 16-bit mask */
96 2424701 __mmask16 mask = _mm512_cmp_ps_mask(t, zero, _CMP_GT_OQ);
97
98 /* Store as 2 bytes (LSB-first: low byte = bits 0-7, high = 8-15) */
99 2424701 bits[i / 8] = (uint8_t)(mask & 0xFF);
100 2424701 bits[i / 8 + 1] = (uint8_t)(mask >> 8);
101 }
102
103 /* Handle remaining floats (8 at a time with AVX2-style or scalar) */
104
2/2
✓ Branch 0 taken 1320464 times.
✓ Branch 1 taken 694422 times.
 2014886 for (; i < dim; i++)
105 {
106 1320464 int byte_idx = i / 8;
107 1320464 int bit_idx = i % 8;
108
2/2
✓ Branch 0 taken 208740 times.
✓ Branch 1 taken 1111724 times.
 1320464 if (bit_idx == 0)
109 208740 bits[byte_idx] = 0; /* Clear byte on first bit */
110
2/2
✓ Branch 0 taken 650002 times.
✓ Branch 1 taken 670462 times.
 1320464 if (transformed[i] > 0)
111 650002 bits[byte_idx] |= (1 << bit_idx);
112 }
113 694422 }
114
115 /*
116 * AVX-512 VPOPCNTDQ Hamming distance.
117 *
118 * Processes 512 bits (64 bytes) per iteration using XOR + VPOPCNTDQ.
119 * VPOPCNTDQ computes popcount of each 64-bit element in a ZMM register.
120 * Requires Ice Lake (2019+) or later.
121 */
122 VS_TARGET_AVX512_VPOPCNTDQ uint32_t
123 126 vs_rabitq_hamming_avx512(
124 const uint8_t *a, const uint8_t *b, uint32_t packed_bytes)
125 {
126 126 __m512i total = _mm512_setzero_si512();
127
128 126 uint32_t i = 0;
129
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 126 times.
 128 for (; i + 64 <= packed_bytes; i += 64)
130 {
131 2 __m512i va = _mm512_loadu_si512(a + i);
132 4 __m512i vb = _mm512_loadu_si512(b + i);
133 2 __m512i x = _mm512_xor_si512(va, vb);
134 4 total = _mm512_add_epi64(total, _mm512_popcnt_epi64(x));
135 }
136
137 126 uint64_t result = vs_horizontal_sum_epi64_avx512(total);
138
139 /* Scalar tail */
140
2/2
✓ Branch 0 taken 1152 times.
✓ Branch 1 taken 126 times.
 1278 for (; i < packed_bytes; i++)
141 1152 result += (uint64_t)__builtin_popcount(a[i] ^ b[i]);
142
143 126 return (uint32_t)result;
144 }
145
146 /*
147 * Multi-candidate AVX-512 VPOPCNTDQ Hamming distance.
148 */
149 VS_TARGET_AVX512_VPOPCNTDQ void
150 14 vs_rabitq_hamming_multi_avx512(
151 const uint8_t *query_bits,
152 const uint8_t *data_bits,
153 uint32_t stride,
154 uint32_t packed_bytes,
155 uint32_t count,
156 uint32_t *results)
157 {
158
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 14 times.
 88 for (uint32_t c = 0; c < count; c++)
159 {
160 74 results[c] = vs_rabitq_hamming_avx512(
161 74 query_bits, data_bits + c * stride, packed_bytes);
162 }
163 14 }
164
165 /*
166 * AVX-512 multi-candidate vertical inner product.
167 *
168 * Processes 4 candidates per dimension chunk. Each iteration:
169 * 1. Load 16 floats from transformed[] (1 zmm register)
170 * 2. For each of 4 candidates: load 2 bytes of bits, form mask,
171 * masked-add into that candidate's accumulator
172 * 3. After all dimensions: horizontal sum each accumulator
173 *
174 * Tail candidates (count % 4) use the single-candidate kernel.
175 */
176 VS_TARGET_AVX512 void
177 1289194 vs_rabitq_inner_product_multi_avx512(
178 const float *transformed,
179 const uint8_t *bits,
180 uint32_t stride,
181 Dimension dim,
182 uint32_t count,
183 float *results)
184 {
185 /* Process groups of 4 candidates */
186 1289194 uint32_t groups = count / 4;
187 1289194 uint32_t tail = count % 4;
188
189
2/2
✓ Branch 0 taken 4279918 times.
✓ Branch 1 taken 1289194 times.
 5569112 for (uint32_t g = 0; g < groups; g++)
190 {
191 4279918 uint32_t base = g * 4;
192
193 4279918 const uint8_t *b0 = bits + (size_t)base * stride;
194 4279918 const uint8_t *b1 = bits + (size_t)(base + 1) * stride;
195 4279918 const uint8_t *b2 = bits + (size_t)(base + 2) * stride;
196 4279918 const uint8_t *b3 = bits + (size_t)(base + 3) * stride;
197
198 4279918 __m512 sum0 = _mm512_setzero_ps();
199 4279918 __m512 sum1 = _mm512_setzero_ps();
200 4279918 __m512 sum2 = _mm512_setzero_ps();
201 4279918 __m512 sum3 = _mm512_setzero_ps();
202
203 /* Main loop: process 16 floats at a time */
204 4279918 Dimension i = 0;
205
2/2
✓ Branch 0 taken 5488512 times.
✓ Branch 1 taken 4279918 times.
 9768430 for (; i + 16 <= dim; i += 16)
206 {
207 5488512 __m512 t = _mm512_loadu_ps(transformed + i);
208
209 5488512 uint32_t bi = i / 8;
210
211 5488512 __mmask16 k0 = (__mmask16)((uint16_t)b0[bi + 1] << 8 | b0[bi]);
212 5488512 __mmask16 k1 = (__mmask16)((uint16_t)b1[bi + 1] << 8 | b1[bi]);
213 5488512 __mmask16 k2 = (__mmask16)((uint16_t)b2[bi + 1] << 8 | b2[bi]);
214 5488512 __mmask16 k3 = (__mmask16)((uint16_t)b3[bi + 1] << 8 | b3[bi]);
215
216 5488512 sum0 = _mm512_mask_add_ps(sum0, k0, sum0, t);
217 5488512 sum1 = _mm512_mask_add_ps(sum1, k1, sum1, t);
218 5488512 sum2 = _mm512_mask_add_ps(sum2, k2, sum2, t);
219 9068542 sum3 = _mm512_mask_add_ps(sum3, k3, sum3, t);
220 }
221
222 4279918 results[base + 0] = vs_horizontal_sum_avx512(sum0);
223 4279918 results[base + 1] = vs_horizontal_sum_avx512(sum1);
224 4279918 results[base + 2] = vs_horizontal_sum_avx512(sum2);
225 4279918 results[base + 3] = vs_horizontal_sum_avx512(sum3);
226
227 /* Scalar tail for remaining dimensions */
228
2/2
✓ Branch 0 taken 24679516 times.
✓ Branch 1 taken 4279918 times.
 28959434 for (; i < dim; i++)
229 {
230 24679516 int byte_idx = i / 8;
231 24679516 int bit_idx = i % 8;
232
233
2/2
✓ Branch 0 taken 11961458 times.
✓ Branch 1 taken 12718058 times.
 24679516 if ((b0[byte_idx] >> bit_idx) & 1)
234 11961458 results[base + 0] += transformed[i];
235
2/2
✓ Branch 0 taken 12454958 times.
✓ Branch 1 taken 12224558 times.
 24679516 if ((b1[byte_idx] >> bit_idx) & 1)
236 12454958 results[base + 1] += transformed[i];
237
2/2
✓ Branch 0 taken 12783906 times.
✓ Branch 1 taken 11895610 times.
 24679516 if ((b2[byte_idx] >> bit_idx) & 1)
238 12783906 results[base + 2] += transformed[i];
239
2/2
✓ Branch 0 taken 12059096 times.
✓ Branch 1 taken 12620420 times.
 24679516 if ((b3[byte_idx] >> bit_idx) & 1)
240 12059096 results[base + 3] += transformed[i];
241 }
242 }
243
244 /* Handle remaining candidates with single-candidate kernel */
245
2/2
✓ Branch 0 taken 541481 times.
✓ Branch 1 taken 1289194 times.
 1830675 for (uint32_t i = groups * 4; i < groups * 4 + tail; i++)
246 {
247 541481 results[i] = vs_rabitq_inner_product_avx512(
248 541481 transformed, bits + (size_t)i * stride, dim);
249 }
250 1289194 }
251
252 #endif /* x86_64 */
253
254 #endif /* VS_SIMD_FULL */
255