GCC Code Coverage Report

Directory: src/
File: src/quant/fast_rotate.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 146 254 57.5%
Functions: 7 11 63.6%
Branches: 99 176 56.2%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * fast_rotate.c - O(d log d) orthonormal rotation
6 *
7 * Generalised Walsh-Hadamard transform with random sign-flip prefix
8 * and a K×K mixer across K sub-blocks, scaled to be orthonormal. Used
9 * as a drop-in for the dense `P^T * x` rotation in RaBitQ.
10 *
11 * For dim = N (power of two) we set K=1 and apply the classic
12 * randomised Hadamard transform: D · H_N · scale.
13 *
14 * For dim = N · K with N power-of-two and K ≥ 2 we view x as a K×N
15 * matrix (K rows of length N), apply H_N to each row in place, then
16 * apply a K×K orthonormal matrix M across the rows at each column.
17 * The combined map M · (I_K ⊗ H_N) · D · scale is orthonormal because
18 * each factor is, and the Kronecker structure makes it O(d log N + d·K)
19 * to evaluate — for dim=768 = 256·3 that's ~9 k ops vs ~1.18 M ops
20 * for the dense sgemv it replaces.
21 */
22
23 #include <math.h>
24 #include <string.h>
25
26 #include "core/memory.h"
27 #include "core/platform.h"
28 #include "quant/fast_rotate.h"
29
30 #if defined(__x86_64__) || defined(_M_X64)
31 #include <immintrin.h>
32 #elif defined(__aarch64__) || defined(_M_ARM64)
33 #include <arm_neon.h>
34 #endif
35
36 bool
37 1888 vs_fast_rotate_supported(Dimension dim)
38 {
39
2/2
✓ Branch 0 taken 366 times.
✓ Branch 1 taken 1522 times.
 1888 if (dim < 4)
40 6 return false;
41
42 /* Factor dim = N * K with N the largest power-of-two divisor.
43 * K must be ≤ VS_FAST_ROTATE_K_MAX, and N ≥ 4 so the inner
44 * FWHT has at least two butterfly stages. */
45 1494 uint32_t n = dim;
46 1494 uint32_t k = 1;
47
2/2
✓ Branch 0 taken 7080 times.
✓ Branch 1 taken 1494 times.
 8574 while ((n & 1) == 0)
48 7080 n >>= 1;
49 /* Now n is the odd part, k = original/odd_part is the power-of-two
50 * part of dim. The decomposition we want is the *opposite*:
51 * fwht_n = power-of-two part, k = odd cofactor. */
52 1494 uint32_t fwht_n = dim / n;
53 1494 k = n;
54
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 1094 times.
 1494 if (fwht_n < 4)
55 40 return false;
56
2/2
✓ Branch 0 taken 26 times.
✓ Branch 1 taken 1428 times.
 1454 if (k > VS_FAST_ROTATE_K_MAX)
57 26 return false;
58 1092 return true;
59 }
60
61 /* SplitMix64: deterministic PRNG used only for sign-vector / mixer
62 * init. Self-contained so init doesn't depend on the rest of the
63 * codebase. */
64 static inline uint64_t
65 139134 splitmix64(uint64_t *state)
66 {
67 139134 uint64_t z = (*state += 0x9E3779B97F4A7C15ULL);
68 139134 z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
69 139134 z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
70 139134 return z ^ (z >> 31);
71 }
72
73 /* Build a K×K orthonormal matrix from PRNG state via Gram-Schmidt on
74 * K random vectors. Cheap because K is small (≤ K_MAX). */
75 static void
76 86 init_mixer(float *m, uint32_t k, uint64_t *prng_state)
77 {
78 22 float v[VS_FAST_ROTATE_K_MAX][VS_FAST_ROTATE_K_MAX];
79
80 /* Random vectors */
81
2/2
✓ Branch 0 taken 274 times.
✓ Branch 1 taken 86 times.
 360 for (uint32_t i = 0; i < k; i++)
82
2/2
✓ Branch 0 taken 918 times.
✓ Branch 1 taken 274 times.
 1192 for (uint32_t j = 0; j < k; j++)
83 {
84 918 uint64_t r = splitmix64(prng_state);
85 /* Standard normal via Box-Muller would be cleaner but the
86 * mixer just needs to be a non-degenerate basis to
87 * orthogonalise — uniform in [-1, 1] is fine. */
88 918 v[i][j] = (float)((double)r / (double)UINT64_MAX) * 2.0f - 1.0f;
89 }
90
91 /* Modified Gram-Schmidt */
92
2/2
✓ Branch 0 taken 274 times.
✓ Branch 1 taken 86 times.
 360 for (uint32_t i = 0; i < k; i++)
93 {
94
2/2
✓ Branch 0 taken 322 times.
✓ Branch 1 taken 274 times.
 596 for (uint32_t j = 0; j < i; j++)
95 {
96 192 float dot = 0.0f;
97
2/2
✓ Branch 0 taken 1214 times.
✓ Branch 1 taken 322 times.
 1536 for (uint32_t c = 0; c < k; c++)
98 1214 dot += v[i][c] * v[j][c];
99
2/2
✓ Branch 0 taken 1214 times.
✓ Branch 1 taken 322 times.
 1536 for (uint32_t c = 0; c < k; c++)
100 1214 v[i][c] -= dot * v[j][c];
101 }
102 192 float norm = 0.0f;
103
2/2
✓ Branch 0 taken 918 times.
✓ Branch 1 taken 274 times.
 1192 for (uint32_t c = 0; c < k; c++)
104 918 norm += v[i][c] * v[i][c];
105 274 norm = sqrtf(norm);
106 /* Degenerate rows are vanishingly unlikely for random uniform
107 * inputs at K ≤ 8; if it happens, fall back to a canonical
108 * basis vector for that slot. */
109
2/2
✓ Branch 0 taken 82 times.
✓ Branch 1 taken 192 times.
 274 if (norm < 1e-6f)
110 {
111 ✗ for (uint32_t c = 0; c < k; c++)
112 ✗ v[i][c] = (c == i) ? 1.0f : 0.0f;
113 }
114 else
115 {
116
2/2
✓ Branch 0 taken 918 times.
✓ Branch 1 taken 274 times.
 1192 for (uint32_t c = 0; c < k; c++)
117 918 v[i][c] /= norm;
118 }
119 }
120
121
2/2
✓ Branch 0 taken 274 times.
✓ Branch 1 taken 86 times.
 360 for (uint32_t i = 0; i < k; i++)
122
2/2
✓ Branch 0 taken 918 times.
✓ Branch 1 taken 274 times.
 1192 for (uint32_t j = 0; j < k; j++)
123 918 m[i * k + j] = v[i][j];
124 86 }
125
126 void
127 732 vs_fast_rotate_init(VsFastRotateParams *p, Dimension dim, uint64_t seed)
128 {
129 732 p->dim = dim;
130 732 p->seed = seed;
131
132 /* Factor dim = fwht_n * k with fwht_n the power-of-two part. */
133 732 uint32_t odd = dim;
134
2/2
✓ Branch 0 taken 3634 times.
✓ Branch 1 taken 732 times.
 4366 while ((odd & 1) == 0)
135 3634 odd >>= 1;
136 732 p->fwht_n = dim / odd;
137 732 p->k = odd;
138
139
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 168 times.
 732 memset(p->signs1, 0, sizeof(p->signs1));
140 732 memset(p->signs2, 0, sizeof(p->signs2));
141
142
2/2
✓ Branch 0 taken 564 times.
✓ Branch 1 taken 168 times.
 732 uint64_t s = seed ? seed : 0xDEADBEEFCAFEBABEULL;
143
2/2
✓ Branch 0 taken 69108 times.
✓ Branch 1 taken 732 times.
 69840 for (Dimension i = 0; i < dim; i++)
144 {
145 69108 uint64_t r = splitmix64(&s);
146
2/2
✓ Branch 0 taken 33893 times.
✓ Branch 1 taken 35215 times.
 69108 if (r & 1)
147 33893 p->signs1[i >> 3] |= (uint8_t)(1u << (i & 7));
148 }
149
2/2
✓ Branch 0 taken 69108 times.
✓ Branch 1 taken 732 times.
 69840 for (Dimension i = 0; i < dim; i++)
150 {
151 69108 uint64_t r = splitmix64(&s);
152
2/2
✓ Branch 0 taken 35305 times.
✓ Branch 1 taken 33803 times.
 69108 if (r & 1)
153 35305 p->signs2[i >> 3] |= (uint8_t)(1u << (i & 7));
154 }
155
156
2/2
✓ Branch 0 taken 146 times.
✓ Branch 1 taken 22 times.
 732 memset(p->mixer, 0, sizeof(p->mixer));
157
2/2
✓ Branch 0 taken 646 times.
✓ Branch 1 taken 86 times.
 732 if (p->k == 1)
158 646 p->mixer[0] = 1.0f;
159 else
160 86 init_mixer(p->mixer, p->k, &s);
161 732 }
162
163 /* In-place Walsh-Hadamard on a length-N block (N power of two). */
164 static inline void
165 ✗ fwht_inplace(float *a, Dimension n)
166 {
167 ✗ for (uint32_t h = 1; h < n; h <<= 1)
168 {
169 ✗ for (uint32_t i = 0; i < n; i += (h << 1))
170 {
171 ✗ for (uint32_t j = i; j < i + h; j++)
172 {
173 ✗ float x = a[j];
174 ✗ float y = a[j + h];
175 ✗ a[j] = x + y;
176 ✗ a[j + h] = x - y;
177 }
178 }
179 }
180 ✗ }
181
182 /* Apply the K×K mixer matrix M to the K-vector formed by taking
183 * `col` from each of K rows. Mixer rows = K, columns = K. */
184 static inline void
185 ✗ apply_mixer_column(
186 float *out, const float *m, uint32_t k, Dimension n, Dimension col)
187 {
188 ✗ float tmp[VS_FAST_ROTATE_K_MAX];
189 ✗ for (uint32_t i = 0; i < k; i++)
190 ✗ tmp[i] = out[i * n + col];
191
192 ✗ for (uint32_t i = 0; i < k; i++)
193 {
194 ✗ float s = 0.0f;
195 ✗ for (uint32_t j = 0; j < k; j++)
196 ✗ s += m[i * k + j] * tmp[j];
197 ✗ out[i * n + col] = s;
198 }
199 ✗ }
200
201 /* Reference scalar implementation. */
202 static void
203 ✗ fast_rotate_apply_scalar(
204 const VsFastRotateParams *p, const float *in, float *out)
205 {
206 ✗ Dimension dim = p->dim;
207 ✗ Dimension n = p->fwht_n;
208 ✗ uint32_t k = p->k;
209 /* Unscaled FWHT on a length-N block has Parseval factor N
210 * (||y||² = N · ||x||²). The K×K mixer is unit-norm. So pre-scale
211 * by 1/sqrt(N), NOT 1/sqrt(dim) — that gets the overall map to
212 * unit-norm regardless of K. */
213 ✗ float invsq = 1.0f / sqrtf((float)n);
214
215 /* Pre-FWHT sign flip (D1) + scale. Folding the 1/sqrt(N) scale
216 * into this pass means the butterflies and mixer can run on plain
217 * sums without a final scaling sweep. */
218 ✗ for (Dimension i = 0; i < dim; i++)
219 {
220 ✗ float sign = ((p->signs1[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
221 ✗ out[i] = in[i] * sign * invsq;
222 }
223
224 /* FWHT on each of K sub-blocks of length N. When K==1 this is
225 * just the classic randomised Hadamard rotation. */
226 ✗ for (uint32_t b = 0; b < k; b++)
227 ✗ fwht_inplace(out + (size_t)b * n, n);
228
229 /* K×K mixer across the blocks. Skip when K==1: mixer is identity. */
230 ✗ if (k > 1)
231 {
232 ✗ for (Dimension col = 0; col < n; col++)
233 ✗ apply_mixer_column(out, p->mixer, k, n, col);
234 }
235
236 /* Post-FWHT sign flip (D2). This second round of randomisation
237 * tightens the concentration of the transform: with one round,
238 * worst-case Lipschitz can break RaBitQ's strict lower bound at
239 * small dim; two rounds match the FJLT analysis (Ailon-Chazelle)
240 * and restore the bound. Cost is one extra O(d) pass. */
241 ✗ for (Dimension i = 0; i < dim; i++)
242 {
243 ✗ float sign = ((p->signs2[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
244 ✗ out[i] *= sign;
245 }
246 ✗ }
247
248 #if defined(__x86_64__) || defined(_M_X64)
249 /*
250 * AVX2 kernel. Vectorizes the parts that map cleanly to 8-wide SIMD:
251 * - D1/D2 sign flips: expand 8 packed sign bits to a ±1 vector, multiply.
252 * - FWHT stages with stride h >= 8: contiguous wide add/sub (bit-identical
253 * to the scalar butterfly). Small stages (h < 8) stay scalar -- they'd
254 * need in-register shuffle networks for little of the total work.
255 * - K x K mixer: vectorized across 8 columns (contiguous within a block).
256 */
257 __attribute__((target("avx2,fma"))) static void
258 ✗ fast_rotate_apply_avx2(
259 const VsFastRotateParams *p, const float *in, float *out)
260 {
261 ✗ Dimension dim = p->dim;
262 ✗ Dimension n = p->fwht_n;
263 ✗ uint32_t k = p->k;
264 ✗ float invsq = 1.0f / sqrtf((float)n);
265 ✗ const __m256i lanebits = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
266 ✗ const __m256 vpos = _mm256_set1_ps(1.0f);
267 ✗ const __m256 vneg = _mm256_set1_ps(-1.0f);
268
269 /* D1: sign flip + 1/sqrt(N) scale. */
270 ✗ __m256 vscale = _mm256_set1_ps(invsq);
271 ✗ Dimension i = 0;
272 ✗ for (; i + 8 <= dim; i += 8)
273 {
274 ✗ __m256i bb = _mm256_and_si256(
275 ✗ _mm256_set1_epi32(p->signs1[i >> 3]), lanebits);
276 ✗ __m256 sign = _mm256_blendv_ps(
277 vpos,
278 vneg,
279 _mm256_castsi256_ps(_mm256_cmpeq_epi32(bb, lanebits)));
280 ✗ __m256 v = _mm256_mul_ps(
281 ✗ _mm256_mul_ps(_mm256_loadu_ps(in + i), sign), vscale);
282 ✗ _mm256_storeu_ps(out + i, v);
283 }
284 ✗ for (; i < dim; i++)
285 {
286 ✗ float s = ((p->signs1[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
287 ✗ out[i] = in[i] * s * invsq;
288 }
289
290 /* FWHT per block: wide add/sub for h >= 8, scalar for the small tail. */
291 ✗ for (uint32_t blk = 0; blk < k; blk++)
292 {
293 ✗ float *a = out + (size_t)blk * n;
294 ✗ for (uint32_t h = 1; h < n; h <<= 1)
295 {
296 ✗ if (h >= 8)
297 {
298 ✗ for (uint32_t s = 0; s < n; s += (h << 1))
299 ✗ for (uint32_t j = s; j < s + h; j += 8)
300 {
301 ✗ __m256 x = _mm256_loadu_ps(a + j);
302 ✗ __m256 y = _mm256_loadu_ps(a + j + h);
303 ✗ _mm256_storeu_ps(a + j, _mm256_add_ps(x, y));
304 ✗ _mm256_storeu_ps(a + j + h, _mm256_sub_ps(x, y));
305 }
306 }
307 else
308 {
309 ✗ for (uint32_t s = 0; s < n; s += (h << 1))
310 ✗ for (uint32_t j = s; j < s + h; j++)
311 {
312 ✗ float x = a[j];
313 ✗ float y = a[j + h];
314 ✗ a[j] = x + y;
315 ✗ a[j + h] = x - y;
316 }
317 }
318 }
319 }
320
321 /* K x K mixer across blocks, vectorized across columns. */
322 ✗ if (k > 1)
323 {
324 ✗ Dimension col = 0;
325 ✗ for (; col + 8 <= n; col += 8)
326 {
327 __m256 blkv[VS_FAST_ROTATE_K_MAX];
328 ✗ for (uint32_t j = 0; j < k; j++)
329 ✗ blkv[j] = _mm256_loadu_ps(out + (size_t)j * n + col);
330 ✗ for (uint32_t r = 0; r < k; r++)
331 {
332 ✗ __m256 acc = _mm256_mul_ps(
333 ✗ _mm256_set1_ps(p->mixer[r * k]), blkv[0]);
334 ✗ for (uint32_t j = 1; j < k; j++)
335 ✗ acc = _mm256_fmadd_ps(
336 ✗ _mm256_set1_ps(p->mixer[r * k + j]), blkv[j], acc);
337 ✗ _mm256_storeu_ps(out + (size_t)r * n + col, acc);
338 }
339 }
340 ✗ for (; col < n; col++)
341 {
342 float tmp[VS_FAST_ROTATE_K_MAX];
343 ✗ for (uint32_t j = 0; j < k; j++)
344 ✗ tmp[j] = out[(size_t)j * n + col];
345 ✗ for (uint32_t r = 0; r < k; r++)
346 {
347 ✗ float s = 0.0f;
348 ✗ for (uint32_t j = 0; j < k; j++)
349 ✗ s += p->mixer[r * k + j] * tmp[j];
350 ✗ out[(size_t)r * n + col] = s;
351 }
352 }
353 }
354
355 /* D2: post-FWHT sign flip. */
356 ✗ i = 0;
357 ✗ for (; i + 8 <= dim; i += 8)
358 {
359 ✗ __m256i bb = _mm256_and_si256(
360 ✗ _mm256_set1_epi32(p->signs2[i >> 3]), lanebits);
361 ✗ __m256 sign = _mm256_blendv_ps(
362 vpos,
363 vneg,
364 _mm256_castsi256_ps(_mm256_cmpeq_epi32(bb, lanebits)));
365 ✗ _mm256_storeu_ps(
366 ✗ out + i, _mm256_mul_ps(_mm256_loadu_ps(out + i), sign));
367 }
368 ✗ for (; i < dim; i++)
369 {
370 ✗ float s = ((p->signs2[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
371 ✗ out[i] *= s;
372 }
373 ✗ }
374
375 /*
376 * AVX-512 kernel (16-wide). Same structure as the AVX2 kernel; the packed
377 * sign bits map directly onto a __mmask16 so D1/D2 need no compare trick.
378 * Uses only avx512f (float arithmetic + masked blend), so it is gated on
379 * SIMD_AVX512F alone. FWHT: 512-wide for h >= 16, 256-wide for h == 8,
380 * scalar below that.
381 */
382 __attribute__((target("avx512f"))) static void
383 711804 fast_rotate_apply_avx512(
384 const VsFastRotateParams *p, const float *in, float *out)
385 {
386 711804 Dimension dim = p->dim;
387 711804 Dimension n = p->fwht_n;
388 711804 uint32_t k = p->k;
389 711804 float invsq = 1.0f / sqrtf((float)n);
390 711804 const __m512 vpos = _mm512_set1_ps(1.0f);
391 711804 const __m512 vneg = _mm512_set1_ps(-1.0f);
392 711804 const __m512 vscale = _mm512_set1_ps(invsq);
393
394 /* D1: sign flip + scale. */
395 711804 Dimension i = 0;
396
2/2
✓ Branch 0 taken 2360156 times.
✓ Branch 1 taken 711804 times.
 3071960 for (; i + 16 <= dim; i += 16)
397 {
398 2360156 __mmask16 m = (__mmask16)(*(const uint16_t *)(p->signs1 + (i >> 3)));
399 2360156 __m512 sign = _mm512_mask_blend_ps(m, vpos, vneg);
400 2870018 __m512 v = _mm512_mul_ps(
401 2360156 _mm512_mul_ps(_mm512_loadu_ps(in + i), sign), vscale);
402 2360156 _mm512_storeu_ps(out + i, v);
403 }
404
2/2
✓ Branch 0 taken 1294668 times.
✓ Branch 1 taken 711804 times.
 2006472 for (; i < dim; i++)
405 {
406
2/2
✓ Branch 0 taken 888435 times.
✓ Branch 1 taken 406233 times.
 1294668 float s = ((p->signs1[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
407 1294668 out[i] = in[i] * s * invsq;
408 }
409
410 /* FWHT per block. */
411
2/2
✓ Branch 0 taken 926150 times.
✓ Branch 1 taken 711804 times.
 1637954 for (uint32_t blk = 0; blk < k; blk++)
412 {
413 926150 float *a = out + (size_t)blk * n;
414
2/2
✓ Branch 0 taken 4162822 times.
✓ Branch 1 taken 926150 times.
 5088972 for (uint32_t h = 1; h < n; h <<= 1)
415 {
416
2/2
✓ Branch 0 taken 826288 times.
✓ Branch 1 taken 3336534 times.
 4162822 if (h >= 16)
417 {
418
2/2
✓ Branch 0 taken 1695607 times.
✓ Branch 1 taken 826288 times.
 2521895 for (uint32_t s = 0; s < n; s += (h << 1))
419
2/2
✓ Branch 0 taken 3020714 times.
✓ Branch 1 taken 1695607 times.
 4716321 for (uint32_t j = s; j < s + h; j += 16)
420 {
421 3020714 __m512 x = _mm512_loadu_ps(a + j);
422 3339320 __m512 y = _mm512_loadu_ps(a + j + h);
423 3020714 _mm512_storeu_ps(a + j, _mm512_add_ps(x, y));
424 3020714 _mm512_storeu_ps(a + j + h, _mm512_sub_ps(x, y));
425 }
426 }
427
2/2
✓ Branch 0 taken 571479 times.
✓ Branch 1 taken 2765055 times.
 3336534 else if (h == 8)
428 {
429
2/2
✓ Branch 0 taken 2267086 times.
✓ Branch 1 taken 571479 times.
 2838565 for (uint32_t s = 0; s < n; s += 16)
430 {
431 2267086 __m256 x = _mm256_loadu_ps(a + s);
432 2685508 __m256 y = _mm256_loadu_ps(a + s + 8);
433 2267086 _mm256_storeu_ps(a + s, _mm256_add_ps(x, y));
434 2267086 _mm256_storeu_ps(a + s + 8, _mm256_sub_ps(x, y));
435 }
436 }
437 else
438 {
439
2/2
✓ Branch 0 taken 34168321 times.
✓ Branch 1 taken 2765055 times.
 36933376 for (uint32_t s = 0; s < n; s += (h << 1))
440
2/2
✓ Branch 0 taken 58558956 times.
✓ Branch 1 taken 34168321 times.
 92727277 for (uint32_t j = s; j < s + h; j++)
441 {
442 58558956 float x = a[j];
443 58558956 float y = a[j + h];
444 58558956 a[j] = x + y;
445 58558956 a[j + h] = x - y;
446 }
447 }
448 }
449 }
450
451 /* K x K mixer across blocks, vectorized across columns. */
452
2/2
✓ Branch 0 taken 106087 times.
✓ Branch 1 taken 605717 times.
 711804 if (k > 1)
453 {
454 91834 Dimension col = 0;
455
2/2
✓ Branch 0 taken 213261 times.
✓ Branch 1 taken 106087 times.
 319348 for (; col + 16 <= n; col += 16)
456 {
457 __m512 blkv[VS_FAST_ROTATE_K_MAX];
458
2/2
✓ Branch 0 taken 640321 times.
✓ Branch 1 taken 213261 times.
 853582 for (uint32_t j = 0; j < k; j++)
459 659329 blkv[j] = _mm512_loadu_ps(out + (size_t)j * n + col);
460
2/2
✓ Branch 0 taken 640321 times.
✓ Branch 1 taken 213261 times.
 853582 for (uint32_t r = 0; r < k; r++)
461 {
462 640321 __m512 acc = _mm512_mul_ps(
463 640321 _mm512_set1_ps(p->mixer[r * k]), blkv[0]);
464
2/2
✓ Branch 0 taken 1283332 times.
✓ Branch 1 taken 640321 times.
 1923653 for (uint32_t j = 1; j < k; j++)
465 1321348 acc = _mm512_fmadd_ps(
466 1283332 _mm512_set1_ps(p->mixer[r * k + j]), blkv[j], acc);
467 640321 _mm512_storeu_ps(out + (size_t)r * n + col, acc);
468 }
469 }
470
2/2
✓ Branch 0 taken 738024 times.
✓ Branch 1 taken 106087 times.
 844111 for (; col < n; col++)
471 {
472 float tmp[VS_FAST_ROTATE_K_MAX];
473
2/2
✓ Branch 0 taken 2227144 times.
✓ Branch 1 taken 738024 times.
 2965168 for (uint32_t j = 0; j < k; j++)
474 2227144 tmp[j] = out[(size_t)j * n + col];
475
2/2
✓ Branch 0 taken 2227144 times.
✓ Branch 1 taken 738024 times.
 2965168 for (uint32_t r = 0; r < k; r++)
476 {
477 2194560 float s = 0.0f;
478
2/2
✓ Branch 0 taken 6764264 times.
✓ Branch 1 taken 2227144 times.
 8991408 for (uint32_t j = 0; j < k; j++)
479 6764264 s += p->mixer[r * k + j] * tmp[j];
480 2227144 out[(size_t)r * n + col] = s;
481 }
482 }
483 }
484
485 /* D2: post-FWHT sign flip. */
486 284514 i = 0;
487
2/2
✓ Branch 0 taken 2360156 times.
✓ Branch 1 taken 711804 times.
 3071960 for (; i + 16 <= dim; i += 16)
488 {
489 2360156 __mmask16 m = (__mmask16)(*(const uint16_t *)(p->signs2 + (i >> 3)));
490 2360156 __m512 sign = _mm512_mask_blend_ps(m, vpos, vneg);
491 2360156 _mm512_storeu_ps(
492 2870018 out + i, _mm512_mul_ps(_mm512_loadu_ps(out + i), sign));
493 }
494
2/2
✓ Branch 0 taken 1294668 times.
✓ Branch 1 taken 711804 times.
 2006472 for (; i < dim; i++)
495 {
496
2/2
✓ Branch 0 taken 720512 times.
✓ Branch 1 taken 574156 times.
 1294668 float s = ((p->signs2[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
497 1294668 out[i] *= s;
498 }
499 711804 }
500 #endif /* x86_64 */
501
502 #if defined(__aarch64__) || defined(_M_ARM64)
503 /*
504 * NEON kernel (128-bit / 4-wide). Same structure as the AVX2 kernel: D1/D2
505 * expand packed sign bits to a +/-1 vector via vtstq/vbslq; FWHT stages with
506 * stride h >= 4 use contiguous wide add/sub (smaller stages stay scalar);
507 * the K x K mixer is vectorized across 4 columns.
508 */
509 static void
510 fast_rotate_apply_neon(
511 const VsFastRotateParams *p, const float *in, float *out)
512 {
513 Dimension dim = p->dim;
514 Dimension n = p->fwht_n;
515 uint32_t k = p->k;
516 float invsq = 1.0f / sqrtf((float)n);
517 const uint32_t lanebits_arr[4] = {1, 2, 4, 8};
518 const uint32x4_t lanebits = vld1q_u32(lanebits_arr);
519 const float32x4_t vpos = vdupq_n_f32(1.0f);
520 const float32x4_t vneg = vdupq_n_f32(-1.0f);
521
522 /* D1: sign flip + scale. */
523 float32x4_t vscale = vdupq_n_f32(invsq);
524 Dimension i = 0;
525 for (; i + 4 <= dim; i += 4)
526 {
527 uint32_t nib = (uint32_t)((p->signs1[i >> 3] >> (i & 7)) & 0xF);
528 uint32x4_t iss = vtstq_u32(vdupq_n_u32(nib), lanebits);
529 float32x4_t sign = vbslq_f32(iss, vneg, vpos);
530 float32x4_t v = vmulq_f32(vmulq_f32(vld1q_f32(in + i), sign), vscale);
531 vst1q_f32(out + i, v);
532 }
533 for (; i < dim; i++)
534 {
535 float s = ((p->signs1[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
536 out[i] = in[i] * s * invsq;
537 }
538
539 /* FWHT per block: wide add/sub for h >= 4, scalar below. */
540 for (uint32_t blk = 0; blk < k; blk++)
541 {
542 float *a = out + (size_t)blk * n;
543 for (uint32_t h = 1; h < n; h <<= 1)
544 {
545 if (h >= 4)
546 {
547 for (uint32_t s = 0; s < n; s += (h << 1))
548 for (Dimension j = s; j < s + h; j += 4)
549 {
550 float32x4_t x = vld1q_f32(a + j);
551 float32x4_t y = vld1q_f32(a + j + h);
552 vst1q_f32(a + j, vaddq_f32(x, y));
553 vst1q_f32(a + j + h, vsubq_f32(x, y));
554 }
555 }
556 else
557 {
558 for (uint32_t s = 0; s < n; s += (h << 1))
559 for (uint32_t j = s; j < s + h; j++)
560 {
561 float x = a[j];
562 float y = a[j + h];
563 a[j] = x + y;
564 a[j + h] = x - y;
565 }
566 }
567 }
568 }
569
570 /* K x K mixer across blocks, vectorized across columns. */
571 if (k > 1)
572 {
573 Dimension col = 0;
574 for (; col + 4 <= n; col += 4)
575 {
576 float32x4_t blkv[VS_FAST_ROTATE_K_MAX];
577 for (uint32_t j = 0; j < k; j++)
578 blkv[j] = vld1q_f32(out + (size_t)j * n + col);
579 for (uint32_t r = 0; r < k; r++)
580 {
581 float32x4_t acc =
582 vmulq_f32(vdupq_n_f32(p->mixer[r * k]), blkv[0]);
583 for (uint32_t j = 1; j < k; j++)
584 acc = vfmaq_f32(
585 acc, vdupq_n_f32(p->mixer[r * k + j]), blkv[j]);
586 vst1q_f32(out + (size_t)r * n + col, acc);
587 }
588 }
589 for (; col < n; col++)
590 {
591 float tmp[VS_FAST_ROTATE_K_MAX];
592 for (uint32_t j = 0; j < k; j++)
593 tmp[j] = out[(size_t)j * n + col];
594 for (uint32_t r = 0; r < k; r++)
595 {
596 float s = 0.0f;
597 for (uint32_t j = 0; j < k; j++)
598 s += p->mixer[r * k + j] * tmp[j];
599 out[(size_t)r * n + col] = s;
600 }
601 }
602 }
603
604 /* D2: post-FWHT sign flip. */
605 i = 0;
606 for (; i + 4 <= dim; i += 4)
607 {
608 uint32_t nib = (uint32_t)((p->signs2[i >> 3] >> (i & 7)) & 0xF);
609 uint32x4_t iss = vtstq_u32(vdupq_n_u32(nib), lanebits);
610 float32x4_t sign = vbslq_f32(iss, vneg, vpos);
611 vst1q_f32(out + i, vmulq_f32(vld1q_f32(out + i), sign));
612 }
613 for (; i < dim; i++)
614 {
615 float s = ((p->signs2[i >> 3] >> (i & 7)) & 1u) ? -1.0f : 1.0f;
616 out[i] *= s;
617 }
618 }
619 #endif /* aarch64 */
620
621 /*
622 * Pick the best kernel for the detected CPU. Resolved once and cached in
623 * the function pointer below (mirroring the distance/fastscan dispatchers)
624 * so the hot path -- which runs per encoded vector at build and once per
625 * query -- pays no per-call capability check.
626 */
627 typedef void (*vs_fast_rotate_fn)(
628 const VsFastRotateParams *, const float *, float *);
629
630 static vs_fast_rotate_fn
631 96 resolve_fast_rotate(void)
632 {
633 #if defined(__x86_64__) || defined(_M_X64)
634
2/2
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 86 times.
 96 if (vs_has_simd(SIMD_AVX512F))
635 10 return fast_rotate_apply_avx512;
636 ✗ if (vs_has_simd(SIMD_AVX2))
637 ✗ return fast_rotate_apply_avx2;
638 #elif defined(__aarch64__) || defined(_M_ARM64)
639 if (vs_has_simd(SIMD_NEON))
640 return fast_rotate_apply_neon;
641 #endif
642 ✗ return fast_rotate_apply_scalar;
643 }
644
645 void
646 711804 vs_fast_rotate_apply(const VsFastRotateParams *p, const float *in, float *out)
647 {
648 /* Benign race: concurrent resolvers all compute the same pointer, and a
649 * pointer store is atomic on supported platforms. */
650 427290 static vs_fast_rotate_fn fn = NULL;
651
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 711708 times.
 711804 if (vs_unlikely(fn == NULL))
652 96 fn = resolve_fast_rotate();
653 711804 fn(p, in, out);
654 711804 }
655