GCC Code Coverage Report

Directory: src/
File: src/quant/matrix.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 90 122 73.8%
Functions: 11 17 64.7%
Branches: 35 50 70.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 * matrix.c - Random orthogonal matrix generation for RaBitQ
6 *
7 * Generates random orthogonal matrices via QR decomposition of Gaussian
8 * random matrices. Uses the Householder algorithm for numerical stability.
9 *
10 * The orthogonal matrix P ensures that residual vectors have isotropic
11 * distribution, which is essential for RaBitQ's theoretical error bounds.
12 */
13
14 #include "vs_config.h"
15
16 #include <math.h>
17 #include <stdint.h>
18 #include <string.h>
19
20 #ifdef VS_HAVE_CBLAS
21 /*
22 * macOS: cblas lives inside Accelerate's vecLib sub-framework. We
23 * include it directly rather than via the `Accelerate/Accelerate.h`
24 * umbrella because the umbrella pulls in MacTypes.h, which
25 * `typedef long Size` clashes with PostgreSQL's `typedef size_t
26 * Size` in the PG-extension build.
27 *
28 * The legacy `vecLib/cblas.h` deprecates every prototype past macOS
29 * 13.3; under -DACCELERATE_NEW_LAPACK (set in meson.build for the
30 * Apple branch) we use `vecLib/cblas_new.h` instead, which aliases
31 * the same call sites to the non-deprecated LP64 symbols.
32 */
33 #ifdef __APPLE__
34 #include <vecLib/cblas_new.h>
35 #else
36 #include <cblas.h>
37 #endif
38 #endif
39
40 #include "algo/simd_utils.h"
41 #include "algo/vecops.h"
42 #include "core/types.h"
43 #include "quant/matrix.h"
44
45 /*
46 * Simple xoshiro256** PRNG for reproducible random generation.
47 * Chosen for speed and quality - not cryptographic.
48 */
49 typedef struct
50 {
51 uint64_t s[4];
52 } Xoshiro256State;
53
54 static inline uint64_t
55 32128844 rotl(uint64_t x, int k)
56 {
57 32128844 return (x << k) | (x >> (64 - k));
58 }
59
60 static uint64_t
61 31922928 xoshiro256_next(Xoshiro256State *state)
62 {
63 31922928 uint64_t *s = state->s;
64 31922928 uint64_t result = rotl(s[1] * 5, 7) * 9;
65 31922928 uint64_t t = s[1] << 17;
66
67 31922928 s[2] ^= s[0];
68 31922928 s[3] ^= s[1];
69 31922928 s[1] ^= s[2];
70 31922928 s[0] ^= s[3];
71 31922928 s[2] ^= t;
72 31922928 s[3] = rotl(s[3], 45);
73
74 31922928 return result;
75 }
76
77 static void
78 248 xoshiro256_seed(Xoshiro256State *state, uint64_t seed)
79 {
80 /* SplitMix64 to expand seed into state */
81
2/2
✓ Branch 0 taken 992 times.
✓ Branch 1 taken 248 times.
 1240 for (int i = 0; i < 4; i++)
82 {
83 992 seed += 0x9e3779b97f4a7c15ULL;
84 992 uint64_t z = seed;
85 992 z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
86 992 z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
87 992 state->s[i] = z ^ (z >> 31);
88 }
89 248 }
90
91 /*
92 * Generate uniform random double in [0, 1)
93 */
94 static double
95 31922928 random_uniform(Xoshiro256State *state)
96 {
97 /* Use upper 53 bits for double precision */
98 47781434 uint64_t x = xoshiro256_next(state) >> 11;
99 31922928 return (double)x / (double)(1ULL << 53);
100 }
101
102 /*
103 * Generate standard normal random variable using Box-Muller transform.
104 * Returns two independent samples for efficiency.
105 */
106 static void
107 15961464 random_normal_pair(Xoshiro256State *state, double *n1, double *n2)
108 {
109 15961464 double u1 = random_uniform(state);
110 15961464 double u2 = random_uniform(state);
111
112 /* Avoid log(0) */
113
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 15961464 times.
 15961464 if (u1 < 1e-15)
114 ✗ u1 = 1e-15;
115
116 15961464 double r = sqrt(-2.0 * log(u1));
117 15961464 double theta = 2.0 * M_PI * u2;
118
119 15961464 *n1 = r * cos(theta);
120 15961464 *n2 = r * sin(theta);
121 15961464 }
122
123 /*
124 * Fill matrix with standard normal random values.
125 * Matrix is dim x dim, stored in row-major order.
126 */
127 static void
128 248 fill_gaussian_matrix(float *matrix, Dimension dim, uint64_t seed)
129 {
130 214 Xoshiro256State state;
131 248 xoshiro256_seed(&state, seed);
132
133 248 size_t n = (size_t)dim * dim;
134 248 size_t i = 0;
135
136 /* Generate pairs of normal values */
137
2/2
✓ Branch 0 taken 15961256 times.
✓ Branch 1 taken 248 times.
 15961504 while (i + 1 < n)
138 {
139 15858312 double n1, n2;
140 15961256 random_normal_pair(&state, &n1, &n2);
141 15961256 matrix[i] = (float)n1;
142 15961256 matrix[i + 1] = (float)n2;
143 15961256 i += 2;
144 }
145
146 /* Handle odd size */
147
2/2
✓ Branch 0 taken 208 times.
✓ Branch 1 taken 40 times.
 248 if (i < n)
148 {
149 194 double n1, n2;
150 208 random_normal_pair(&state, &n1, &n2);
151 194 (void)n2;
152 208 matrix[i] = (float)n1;
153 }
154 248 }
155
156 /*
157 * QR decomposition using modified Gram-Schmidt.
158 *
159 * Input: A (dim x dim matrix, row-major)
160 * Output: Q (orthogonal matrix, row-major)
161 *
162 * The input matrix A is overwritten with Q.
163 */
164 static void
165 248 gram_schmidt_qr(float *A, Dimension dim)
166 {
167 /* Process columns */
168
2/2
✓ Branch 0 taken 19428 times.
✓ Branch 1 taken 248 times.
 19676 for (Dimension j = 0; j < dim; j++)
169 {
170 /* Get pointer to column j (stored as row j in transposed view) */
171 /* For row-major, column j elements are at A[0*dim+j], A[1*dim+j], ...
172 */
173
174 /* Compute norm of column j */
175 1694 float norm = 0.0f;
176
2/2
✓ Branch 0 taken 31922720 times.
✓ Branch 1 taken 19428 times.
 31942148 for (Dimension i = 0; i < dim; i++)
177 31922720 norm += A[i * dim + j] * A[i * dim + j];
178 19428 norm = sqrtf(norm);
179
180 /* Handle near-zero columns (shouldn't happen with Gaussian init) */
181
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 19428 times.
 19428 if (norm < 1e-10f)
182 ✗ norm = 1.0f;
183
184 /* Normalize column j */
185
2/2
✓ Branch 0 taken 31922720 times.
✓ Branch 1 taken 19428 times.
 31942148 for (Dimension i = 0; i < dim; i++)
186 31922720 A[i * dim + j] /= norm;
187
188 /* Orthogonalize remaining columns against column j */
189
2/2
✓ Branch 0 taken 15951646 times.
✓ Branch 1 taken 19428 times.
 15971074 for (Dimension k = j + 1; k < dim; k++)
190 {
191 /* Compute dot product of columns j and k */
192 102104 float dot = 0.0f;
193
2/2
✓ Branch 0 taken 30708916994 times.
✓ Branch 1 taken 15951646 times.
 30724868640 for (Dimension i = 0; i < dim; i++)
194 30708916994 dot += A[i * dim + j] * A[i * dim + k];
195
196 /* Subtract projection: col_k = col_k - dot * col_j */
197
2/2
✓ Branch 0 taken 30708916994 times.
✓ Branch 1 taken 15951646 times.
 30724868640 for (Dimension i = 0; i < dim; i++)
198 30708916994 A[i * dim + k] -= dot * A[i * dim + j];
199 }
200 }
201 248 }
202
203 /*
204 * Generate random orthogonal matrix using QR decomposition.
205 *
206 * Creates a Gaussian random matrix and orthogonalizes it via QR.
207 * The resulting Q matrix is uniformly distributed over O(n)
208 * (orthogonal group) according to Haar measure.
209 *
210 * Parameters:
211 * matrix: Output buffer (dim * dim floats, row-major)
212 * dim: Matrix dimension
213 * seed: Random seed for reproducibility
214 *
215 * Returns 0 on success, -1 on failure.
216 */
217 int
218 248 vs_random_orthogonal_matrix(float *matrix, Dimension dim, uint64_t seed)
219 {
220
2/4
✓ Branch 0 taken 248 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 34 times.
 248 if (matrix == NULL || dim == 0)
221 ✗ return -1;
222
223 /* Fill with Gaussian random values */
224 248 fill_gaussian_matrix(matrix, dim, seed);
225
226 /* QR decomposition to get orthogonal matrix */
227 248 gram_schmidt_qr(matrix, dim);
228
229 248 return 0;
230 }
231
232 /*
233 * Verify matrix is approximately orthogonal (for testing).
234 *
235 * Checks that M * M^T ≈ I within tolerance.
236 *
237 * Returns 1 if orthogonal, 0 if not.
238 */
239 int
240 4 vs_matrix_is_orthogonal(const float *matrix, Dimension dim, float tolerance)
241 {
242
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (matrix == NULL || dim == 0)
243 ✗ return 0;
244
245 /* Check M * M^T = I */
246
2/2
✓ Branch 0 taken 144 times.
✓ Branch 1 taken 4 times.
 148 for (Dimension i = 0; i < dim; i++)
247 {
248
2/2
✓ Branch 0 taken 8320 times.
✓ Branch 1 taken 144 times.
 8464 for (Dimension j = 0; j < dim; j++)
249 {
250 /* Compute (i,j) element of M * M^T = dot(row_i, row_j) */
251 ✗ float dot =
252 8320 vs_dot_product(matrix + i * dim, matrix + j * dim, dim);
253
254 /* Expected value: 1 on diagonal, 0 elsewhere */
255
2/2
✓ Branch 0 taken 144 times.
✓ Branch 1 taken 8176 times.
 8320 float expected = (i == j) ? 1.0f : 0.0f;
256 8320 float diff = fabsf(dot - expected);
257
258
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8320 times.
 8320 if (diff > tolerance)
259 ✗ return 0;
260 }
261 }
262
263 4 return 1;
264 }
265
266 /*
267 * Matrix-vector multiplication: result = M^T * v
268 *
269 * Computes the product of the transpose of M with vector v.
270 * This is the common operation in RaBitQ for transforming residuals.
271 *
272 * Rewritten as: result = sum_j (v[j] * row_j)
273 * This accesses rows (contiguous in row-major), enabling auto-vectorization.
274 *
275 * Parameters:
276 * M: Input matrix (dim x dim, row-major)
277 * v: Input vector (dim elements)
278 * result: Output vector (dim elements)
279 * dim: Dimension
280 */
281 void
282 62612 vs_matrix_transpose_vector_mul(
283 const float *M, const float *v, float *result, Dimension dim)
284 {
285 /* Zero result */
286 62612 memset(result, 0, dim * sizeof(float));
287
288 /* Accumulate v[j] * row_j for each row j */
289
2/2
✓ Branch 0 taken 5852824 times.
✓ Branch 1 taken 62612 times.
 5915436 for (Dimension j = 0; j < dim; j++)
290 {
291 5852824 const float *row = M + j * dim;
292 5852824 float scale = v[j];
293
294 /* This inner loop auto-vectorizes well */
295
2/2
✓ Branch 0 taken 9905246360 times.
✓ Branch 1 taken 5852824 times.
 9911099184 for (Dimension i = 0; i < dim; i++)
296 9905246360 result[i] += scale * row[i];
297 }
298 62612 }
299
300 /*
301 * Matrix-vector multiplication: result = M * v
302 *
303 * Each output element is the dot product of a row with v.
304 * Rows are contiguous in row-major storage, enabling SIMD.
305 *
306 * Parameters:
307 * M: Input matrix (dim x dim, row-major)
308 * v: Input vector (dim elements)
309 * result: Output vector (dim elements)
310 * dim: Dimension
311 */
312 void
313 ✗ vs_matrix_vector_mul(
314 const float *M, const float *v, float *result, Dimension dim)
315 {
316 ✗ for (Dimension i = 0; i < dim; i++)
317 ✗ result[i] = vs_dot_product(M + i * dim, v, dim);
318 ✗ }
319
320 /*
321 * Batched matrix-vector multiplication: results = M^T * vectors
322 *
323 * Computes M^T * v for multiple vectors at once. This is more efficient
324 * than calling vs_matrix_transpose_vector_mul() repeatedly because:
325 * 1. Matrix M is loaded into cache once and reused for all vectors
326 * 2. Inner loop processes contiguous memory (good for SIMD)
327 *
328 * Memory layout:
329 * vectors: count vectors of dim elements each, contiguous (vectors[i*dim +
330 * j]) results: count output vectors, same layout
331 *
332 * Parameters:
333 * M: Input matrix (dim x dim, row-major)
334 * vectors: Input vectors (count * dim elements)
335 * results: Output vectors (count * dim elements)
336 * count: Number of vectors to process
337 * dim: Vector/matrix dimension
338 */
339 VS_TARGET_CLONES static void
340 ✗ matrix_transpose_mul_batch_inner(
341 const float *M_row,
342 const float *vectors,
343 float *results,
344 uint32_t count,
345 Dimension dim,
346 Dimension j)
347 {
348 /* Process each vector - inner loop is contiguous and vectorizes well */
349 ✗ for (uint32_t i = 0; i < count; i++)
350 {
351 ✗ float scale = vectors[i * dim + j];
352 ✗ float *result = results + i * dim;
353 ✗ const float *row = M_row;
354
355 ✗ for (Dimension k = 0; k < dim; k++)
356 ✗ result[k] += scale * row[k];
357 }
358 ✗ }
359
360 static void
361 ✗ matrix_transpose_mul_batch_builtin(
362 const float *M,
363 const float *vectors,
364 float *results,
365 uint32_t count,
366 Dimension dim)
367 {
368 /* Zero all results */
369 ✗ memset(results, 0, (size_t)count * dim * sizeof(float));
370
371 /* Process each row of M - row stays in cache while we update all vectors
372 */
373 ✗ for (Dimension j = 0; j < dim; j++)
374 {
375 ✗ const float *M_row = M + j * dim;
376 ✗ matrix_transpose_mul_batch_inner(
377 M_row, vectors, results, count, dim, j);
378 }
379 ✗ }
380
381 #ifdef VS_HAVE_CBLAS
382 static void
383 matrix_transpose_mul_batch_cblas(
384 const float *M,
385 const float *vectors,
386 float *results,
387 uint32_t count,
388 Dimension dim)
389 {
390 /*
391 * Use CBLAS sgemm for optimized matrix multiplication.
392 *
393 * We want: results[i] = M^T * vectors[i] for each vector i
394 * In row-major form: results = vectors * M
395 *
396 * sgemm: C = alpha * A * B + beta * C
397 * Where: A = vectors (count x dim), B = M (dim x dim), C = results
398 */
399 cblas_sgemm(
400 CblasRowMajor,
401 CblasNoTrans,
402 CblasNoTrans,
403 (int)count, /* M */
404 (int)dim, /* N */
405 (int)dim, /* K */
406 1.0f, /* alpha */
407 vectors, /* A */
408 (int)dim, /* lda */
409 M, /* B */
410 (int)dim, /* ldb */
411 0.0f, /* beta */
412 results, /* C */
413 (int)dim /* ldc */
414 );
415 }
416 #endif /* VS_HAVE_CBLAS */
417
418 /*
419 * Runtime selection of matrix multiplication implementation.
420 * Default: use CBLAS if available, else built-in.
421 */
422 static bool g_use_cblas = true;
423
424 void
425 ✗ vs_matrix_set_use_cblas(bool use_cblas)
426 {
427 ✗ g_use_cblas = use_cblas;
428 ✗ }
429
430 bool
431 2 vs_matrix_get_use_cblas(void)
432 {
433 #ifdef VS_HAVE_CBLAS
434 return g_use_cblas;
435 #else
436 2 return false;
437 #endif
438 }
439
440 const char *
441 ✗ vs_matrix_impl_name(void)
442 {
443 #ifdef VS_HAVE_CBLAS
444 return g_use_cblas ? "cblas" : "builtin";
445 #else
446 ✗ return "builtin";
447 #endif
448 }
449
450 void
451 ✗ vs_matrix_transpose_vector_mul_batch(
452 const float *M,
453 const float *vectors,
454 float *results,
455 uint32_t count,
456 Dimension dim)
457 {
458 #ifdef VS_HAVE_CBLAS
459 if (g_use_cblas)
460 {
461 matrix_transpose_mul_batch_cblas(M, vectors, results, count, dim);
462 return;
463 }
464 #endif
465 ✗ matrix_transpose_mul_batch_builtin(M, vectors, results, count, dim);
466 ✗ }
467