| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * fastscan.c - VPSHUFB-based fast scan for RaBitQ | ||
| 6 | * | ||
| 7 | * Implements the fastscan approach from the RaBitQ Library: | ||
| 8 | * - LUT construction with kPos bit-to-dimension mapping | ||
| 9 | * - Code packing with kPerm0 vector interleaving | ||
| 10 | * - Even/odd byte accumulation (no explicit widening) | ||
| 11 | * - SIMD dispatch (AVX-512 > AVX2 > scalar) | ||
| 12 | */ | ||
| 13 | |||
| 14 | #include "vs_config.h" | ||
| 15 | |||
| 16 | #include <limits.h> | ||
| 17 | #include <math.h> | ||
| 18 | #include <stdatomic.h> | ||
| 19 | #include <stdint.h> | ||
| 20 | #include <string.h> | ||
| 21 | |||
| 22 | #include "algo/simd_utils.h" | ||
| 23 | #include "core/platform.h" | ||
| 24 | #include "quant/fastscan.h" | ||
| 25 | #include "quant/rabitq.h" | ||
| 26 | |||
| 27 | /* ---------------------------------------------------------------- | ||
| 28 | * kPos: maps 4-bit code to dimension index for LUT construction. | ||
| 29 | * | ||
| 30 | * For code j, kPos[j] gives the dimension index of the lowest | ||
| 31 | * set bit. Used in the incremental LUT build: | ||
| 32 | * lut[j] = lut[j - LOWBIT(j)] + query[kPos[j]] | ||
| 33 | * ---------------------------------------------------------------- */ | ||
| 34 | |||
| 35 | /* LUT entry mapping (implicit in build_lut code): | ||
| 36 | * bit 0 → s0 (dim 0), bit 1 → s1 (dim 1), | ||
| 37 | * bit 2 → s2 (dim 2), bit 3 → s3 (dim 3). | ||
| 38 | * Entry j = sum of s_i for each set bit i in j. */ | ||
| 39 | |||
| 40 | /* ---------------------------------------------------------------- | ||
| 41 | * kPerm0: vector interleaving for even/odd byte accumulation. | ||
| 42 | * | ||
| 43 | * Maps logical vector index (0-15) to byte position within a | ||
| 44 | * 16-byte block. Even positions get vectors 0-7, odd positions | ||
| 45 | * get vectors 8-15. This enables uint16 accumulation without | ||
| 46 | * explicit widening: even bytes accumulate one group, odd bytes | ||
| 47 | * another. | ||
| 48 | * ---------------------------------------------------------------- */ | ||
| 49 | |||
| 50 | static const int kPerm0[16] = { | ||
| 51 | 0, | ||
| 52 | 8, | ||
| 53 | 1, | ||
| 54 | 9, | ||
| 55 | 2, | ||
| 56 | 10, | ||
| 57 | 3, | ||
| 58 | 11, | ||
| 59 | 4, | ||
| 60 | 12, | ||
| 61 | 5, | ||
| 62 | 13, | ||
| 63 | 6, | ||
| 64 | 14, | ||
| 65 | 7, | ||
| 66 | 15, | ||
| 67 | }; | ||
| 68 | |||
| 69 | /* Lowest set bit: j & (-j) */ | ||
| 70 | #define LOWBIT(j) ((j) & (-(j))) | ||
| 71 | |||
| 72 | /* ---------------------------------------------------------------- | ||
| 73 | * LUT construction | ||
| 74 | * ---------------------------------------------------------------- */ | ||
| 75 | |||
| 76 | static void | ||
| 77 | 44 | fastscan_build_lut_scalar( | |
| 78 | const float *transformed, | ||
| 79 | Dimension dim, | ||
| 80 | uint8_t *lut_out, | ||
| 81 | float *delta_out, | ||
| 82 | float *bias_out) | ||
| 83 | { | ||
| 84 | 44 | uint32_t nsq = VS_FASTSCAN_NSQ(dim); | |
| 85 | 44 | uint32_t nsq_pairs = VS_FASTSCAN_NSQ_PAIRS(dim); | |
| 86 | |||
| 87 | /* Compute global min/max analytically from transformed values. | ||
| 88 | * Min LUT entry per sq = min(0, sum of negative values). | ||
| 89 | * Max LUT entry per sq = max(0, sum of positive values). | ||
| 90 | * Global min = sum of all sq mins. Global max = sum of all sq maxes. */ | ||
| 91 | 44 | float global_min = 0.0f; | |
| 92 | 44 | float global_max = 0.0f; | |
| 93 |
2/2✓ Branch 0 taken 17710 times.
✓ Branch 1 taken 44 times.
|
17754 | for (Dimension d = 0; d < dim; d++) |
| 94 | { | ||
| 95 |
2/2✓ Branch 0 taken 8540 times.
✓ Branch 1 taken 9170 times.
|
17710 | if (transformed[d] > 0) |
| 96 | 8540 | global_max += transformed[d]; | |
| 97 | else | ||
| 98 | 9170 | global_min += transformed[d]; | |
| 99 | } | ||
| 100 | |||
| 101 | 44 | float range = global_max - global_min; | |
| 102 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 44 times.
|
44 | if (range < VS_FASTSCAN_MIN_RANGE) |
| 103 | ✗ | range = VS_FASTSCAN_MIN_RANGE; | |
| 104 | |||
| 105 | 44 | float delta = range / (float)UINT8_MAX; | |
| 106 | 44 | float inv_delta = 1.0f / delta; | |
| 107 | 44 | *delta_out = delta; | |
| 108 | 44 | *bias_out = global_min * (float)nsq; | |
| 109 | |||
| 110 | /* Build uint8 LUT directly (no intermediate float buffer). | ||
| 111 | * Pre-scale the bias and query values so each entry is just | ||
| 112 | * additions + truncation to uint8. */ | ||
| 113 | 44 | float bias_scaled = -global_min * inv_delta + 0.5f; | |
| 114 | |||
| 115 | 44 | memset(lut_out, 0, nsq_pairs * 2 * 16); | |
| 116 | |||
| 117 | 44 | const float *q = transformed; | |
| 118 |
2/2✓ Branch 0 taken 4428 times.
✓ Branch 1 taken 44 times.
|
4472 | for (uint32_t sq = 0; sq < nsq; sq++) |
| 119 | { | ||
| 120 | 4428 | uint8_t *out = lut_out + sq * 16; | |
| 121 | |||
| 122 | /* Pre-scale the 4 values for this subquantizer */ | ||
| 123 | 4428 | Dimension base = sq * 4; | |
| 124 |
1/2✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
|
4428 | float s0 = (base + 0 < dim) ? q[0] * inv_delta : 0.0f; |
| 125 |
1/2✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
|
4428 | float s1 = (base + 1 < dim) ? q[1] * inv_delta : 0.0f; |
| 126 |
1/2✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
|
4428 | float s2 = (base + 2 < dim) ? q[2] * inv_delta : 0.0f; |
| 127 |
2/2✓ Branch 0 taken 4426 times.
✓ Branch 1 taken 2 times.
|
4428 | float s3 = (base + 3 < dim) ? q[3] * inv_delta : 0.0f; |
| 128 | |||
| 129 | /* Build 16 entries via incremental sums. | ||
| 130 | * Uses kPos mapping: bit 0→dim 0, bit 1→dim 1, etc. */ | ||
| 131 | 4428 | float p01 = s0 + s1; | |
| 132 | 4428 | float p23 = s2 + s3; | |
| 133 | |||
| 134 | 4428 | float f0 = bias_scaled; | |
| 135 | 4428 | float f4 = f0 + s2; | |
| 136 | 4428 | float f8 = f0 + s3; | |
| 137 | 4428 | float f12 = f0 + p23; | |
| 138 | |||
| 139 | #define Q(v) \ | ||
| 140 | ((uint8_t)((int)(v) < 0 ? 0 \ | ||
| 141 | : ((int)(v) > UINT8_MAX ? UINT8_MAX : (int)(v)))) | ||
| 142 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[0] = Q(f0); |
| 143 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[1] = Q(f0 + s0); |
| 144 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[2] = Q(f0 + s1); |
| 145 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[3] = Q(f0 + p01); |
| 146 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[4] = Q(f4); |
| 147 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[5] = Q(f4 + s0); |
| 148 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[6] = Q(f4 + s1); |
| 149 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[7] = Q(f4 + p01); |
| 150 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[8] = Q(f8); |
| 151 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[9] = Q(f8 + s0); |
| 152 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[10] = Q(f8 + s1); |
| 153 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[11] = Q(f8 + p01); |
| 154 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[12] = Q(f12); |
| 155 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[13] = Q(f12 + s0); |
| 156 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[14] = Q(f12 + s1); |
| 157 |
2/4✓ Branch 0 taken 4428 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4428 times.
✗ Branch 3 not taken.
|
4428 | out[15] = Q(f12 + p01); |
| 158 | #undef Q | ||
| 159 | 4428 | q += 4; | |
| 160 | } | ||
| 161 | 44 | } | |
| 162 | |||
| 163 | /* ---------------------------------------------------------------- | ||
| 164 | * High-accuracy LUT construction (uint16 → split lo/hi tables) | ||
| 165 | * ---------------------------------------------------------------- */ | ||
| 166 | |||
| 167 | static void | ||
| 168 | 394425 | fastscan_build_lut_hacc_scalar( | |
| 169 | const float *transformed, | ||
| 170 | Dimension dim, | ||
| 171 | uint8_t *lut_out, | ||
| 172 | float *delta_out, | ||
| 173 | float *bias_out) | ||
| 174 | { | ||
| 175 | 394425 | uint32_t nsq = VS_FASTSCAN_NSQ(dim); | |
| 176 | |||
| 177 | /* Same min/max as uint8 version */ | ||
| 178 | 394425 | float global_min = 0.0f; | |
| 179 | 394425 | float global_max = 0.0f; | |
| 180 |
2/2✓ Branch 0 taken 26908804 times.
✓ Branch 1 taken 394425 times.
|
27303229 | for (Dimension d = 0; d < dim; d++) |
| 181 | { | ||
| 182 |
2/2✓ Branch 0 taken 13421820 times.
✓ Branch 1 taken 13486984 times.
|
26908804 | if (transformed[d] > 0) |
| 183 | 13421820 | global_max += transformed[d]; | |
| 184 | else | ||
| 185 | 13486984 | global_min += transformed[d]; | |
| 186 | } | ||
| 187 | |||
| 188 | 394425 | float range = global_max - global_min; | |
| 189 |
2/2✓ Branch 0 taken 31 times.
✓ Branch 1 taken 394394 times.
|
394425 | if (range < VS_FASTSCAN_MIN_RANGE) |
| 190 | 31 | range = VS_FASTSCAN_MIN_RANGE; | |
| 191 | |||
| 192 | 394425 | float delta = range / (float)UINT16_MAX; | |
| 193 | 394425 | float inv_delta = 1.0f / delta; | |
| 194 | 394425 | *delta_out = delta; | |
| 195 | 394425 | *bias_out = global_min * (float)nsq; | |
| 196 | |||
| 197 | 394425 | float bias_scaled = -global_min * inv_delta + 0.5f; | |
| 198 | |||
| 199 | /* Build uint16 LUT entries, then split into lo/hi byte tables. | ||
| 200 | * Layout: for each group of 4 sq (matching 64B code blocks), | ||
| 201 | * [lo_table_64B, hi_table_64B]. */ | ||
| 202 | 394425 | uint32_t lut_bytes = VS_FASTSCAN_LUT_HACC_BYTES(dim); | |
| 203 | 394425 | memset(lut_out, 0, lut_bytes); | |
| 204 | |||
| 205 | 394425 | const float *q = transformed; | |
| 206 |
2/2✓ Branch 0 taken 6735365 times.
✓ Branch 1 taken 394425 times.
|
7129790 | for (uint32_t sq = 0; sq < nsq; sq++) |
| 207 | { | ||
| 208 | 6735365 | Dimension base = sq * 4; | |
| 209 |
1/2✓ Branch 0 taken 6735365 times.
✗ Branch 1 not taken.
|
6735365 | float s0 = (base + 0 < dim) ? q[0] * inv_delta : 0.0f; |
| 210 |
1/2✓ Branch 0 taken 6735365 times.
✗ Branch 1 not taken.
|
6735365 | float s1 = (base + 1 < dim) ? q[1] * inv_delta : 0.0f; |
| 211 |
1/2✓ Branch 0 taken 6735365 times.
✗ Branch 1 not taken.
|
6735365 | float s2 = (base + 2 < dim) ? q[2] * inv_delta : 0.0f; |
| 212 |
2/2✓ Branch 0 taken 6702709 times.
✓ Branch 1 taken 32656 times.
|
6735365 | float s3 = (base + 3 < dim) ? q[3] * inv_delta : 0.0f; |
| 213 | |||
| 214 | 6735365 | float p01 = s0 + s1; | |
| 215 | 6735365 | float p23 = s2 + s3; | |
| 216 | |||
| 217 | 6735365 | float f0 = bias_scaled; | |
| 218 | 6735365 | float f4 = f0 + s2; | |
| 219 | 6735365 | float f8 = f0 + s3; | |
| 220 | 6735365 | float f12 = f0 + p23; | |
| 221 | |||
| 222 | 6735365 | float entries[16] = { | |
| 223 | f0, | ||
| 224 | 6735365 | f0 + s0, | |
| 225 | 6735365 | f0 + s1, | |
| 226 | 6735365 | f0 + p01, | |
| 227 | f4, | ||
| 228 | 6735365 | f4 + s0, | |
| 229 | 6735365 | f4 + s1, | |
| 230 | 6735365 | f4 + p01, | |
| 231 | f8, | ||
| 232 | 6735365 | f8 + s0, | |
| 233 | 6735365 | f8 + s1, | |
| 234 | 6735365 | f8 + p01, | |
| 235 | f12, | ||
| 236 | 6735365 | f12 + s0, | |
| 237 | 6735365 | f12 + s1, | |
| 238 | 6735365 | f12 + p01, | |
| 239 | }; | ||
| 240 | |||
| 241 | /* Quantize to uint16 and split into lo/hi bytes. | ||
| 242 | * Group of 4 sq → 128B block: [lo_64B, hi_64B]. | ||
| 243 | * Within each 64B: 4 × 16-byte tables. */ | ||
| 244 | 6735365 | uint32_t group4 = sq / 4; | |
| 245 | 6735365 | uint32_t pos_in_group = sq % 4; | |
| 246 | 6735365 | uint8_t *lo_base = lut_out + group4 * 128 + pos_in_group * 16; | |
| 247 | 6735365 | uint8_t *hi_base = lo_base + 64; | |
| 248 | |||
| 249 |
2/2✓ Branch 0 taken 107765840 times.
✓ Branch 1 taken 6735365 times.
|
114501205 | for (uint32_t j = 0; j < 16; j++) |
| 250 | { | ||
| 251 | 107765840 | float scaled = entries[j]; | |
| 252 | 107765840 | int val = (int)scaled; | |
| 253 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 704224 times.
|
107765840 | if (val < 0) |
| 254 | ✗ | val = 0; | |
| 255 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 704224 times.
|
107765840 | if (val > UINT16_MAX) |
| 256 | ✗ | val = UINT16_MAX; | |
| 257 | 107765840 | lo_base[j] = (uint8_t)(val & 0xFF); | |
| 258 | 107765840 | hi_base[j] = (uint8_t)((val >> 8) & 0xFF); | |
| 259 | } | ||
| 260 | |||
| 261 | 6735365 | q += 4; | |
| 262 | } | ||
| 263 | 394425 | } | |
| 264 | |||
| 265 | /* ---------------------------------------------------------------- | ||
| 266 | * High-accuracy scalar accumulate (reference) | ||
| 267 | * ---------------------------------------------------------------- */ | ||
| 268 | |||
| 269 | static void | ||
| 270 | 20 | fastscan_accumulate_hacc_scalar( | |
| 271 | const uint8_t *codes, | ||
| 272 | const uint8_t *lut, | ||
| 273 | int32_t *accum, | ||
| 274 | Dimension dim) | ||
| 275 | { | ||
| 276 | 20 | uint32_t code_length = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 277 | |||
| 278 | 20 | int32_t lo_accum[VS_FASTSCAN_GROUP] = {0}; | |
| 279 | 20 | int32_t hi_accum[VS_FASTSCAN_GROUP] = {0}; | |
| 280 | |||
| 281 |
2/2✓ Branch 0 taken 2648 times.
✓ Branch 1 taken 20 times.
|
2668 | for (uint32_t i = 0; i < code_length; i += 16) |
| 282 | { | ||
| 283 | 2648 | const uint8_t *c = codes + i; | |
| 284 | /* lo/hi tables are interleaved per group of 4 sq (64B each). | ||
| 285 | * Compute which 128B block and offset within it. */ | ||
| 286 | 2648 | uint32_t sq_idx = i / 16; | |
| 287 | 2648 | uint32_t grp4 = sq_idx / 4; | |
| 288 | 2648 | uint32_t pos = sq_idx % 4; | |
| 289 | 2648 | const uint8_t *lo_tab = lut + grp4 * 128 + pos * 16; | |
| 290 | 2648 | const uint8_t *hi_tab = lo_tab + 64; | |
| 291 | |||
| 292 |
2/2✓ Branch 0 taken 42368 times.
✓ Branch 1 taken 2648 times.
|
45016 | for (uint32_t j = 0; j < 16; j++) |
| 293 | { | ||
| 294 | 42368 | uint8_t byte_val = c[j]; | |
| 295 | 42368 | uint8_t lo_nib = byte_val & 0x0F; | |
| 296 | 42368 | uint8_t hi_nib = byte_val >> 4; | |
| 297 | |||
| 298 | 42368 | lo_accum[kPerm0[j]] += lo_tab[lo_nib]; | |
| 299 | 42368 | lo_accum[kPerm0[j] + 16] += lo_tab[hi_nib]; | |
| 300 | |||
| 301 | 42368 | hi_accum[kPerm0[j]] += hi_tab[lo_nib]; | |
| 302 | 42368 | hi_accum[kPerm0[j] + 16] += hi_tab[hi_nib]; | |
| 303 | } | ||
| 304 | } | ||
| 305 | |||
| 306 |
2/2✓ Branch 0 taken 640 times.
✓ Branch 1 taken 20 times.
|
660 | for (uint32_t v = 0; v < VS_FASTSCAN_GROUP; v++) |
| 307 | 640 | accum[v] = lo_accum[v] + (hi_accum[v] << 8); | |
| 308 | 20 | } | |
| 309 | |||
| 310 | /* ---------------------------------------------------------------- | ||
| 311 | * Code packing with kPerm0 interleaving | ||
| 312 | * | ||
| 313 | * Each column (8 dims = 1 byte of 1-bit code) produces two | ||
| 314 | * 16-byte blocks: lower nibble sq first, upper nibble sq second. | ||
| 315 | * This matches the sequential LUT order (sq0, sq1, sq2, ...). | ||
| 316 | * | ||
| 317 | * Within each 16-byte block, vectors are interleaved with kPerm0: | ||
| 318 | * byte j packs kPerm0[j] vector (lo nibble) and kPerm0[j]+16 | ||
| 319 | * vector (hi nibble). This enables the even/odd byte accumulation | ||
| 320 | * trick in the SIMD kernel. | ||
| 321 | * | ||
| 322 | * Memory layout per group per column: | ||
| 323 | * bytes 0-15: lower nibble sq (sq0, sq2, ...) with kPerm0 | ||
| 324 | * bytes 16-31: upper nibble sq (sq1, sq3, ...) with kPerm0 | ||
| 325 | * ---------------------------------------------------------------- */ | ||
| 326 | |||
| 327 | uint32_t | ||
| 328 | 17986 | vs_fastscan_pack_codes( | |
| 329 | const uint8_t *bits_1bit, | ||
| 330 | uint32_t count, | ||
| 331 | Dimension dim, | ||
| 332 | uint8_t *codes_out) | ||
| 333 | { | ||
| 334 | 17986 | uint32_t packed_bytes = (dim + 7) / 8; | |
| 335 | 17986 | uint32_t ngroups = (count + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP; | |
| 336 | |||
| 337 | 17986 | memset(codes_out, 0, (size_t)ngroups * VS_FASTSCAN_GROUP_BYTES(dim)); | |
| 338 | |||
| 339 |
2/2✓ Branch 0 taken 18026 times.
✓ Branch 1 taken 17986 times.
|
36012 | for (uint32_t g = 0; g < ngroups; g++) |
| 340 | { | ||
| 341 | 18026 | uint32_t g_start = g * VS_FASTSCAN_GROUP; | |
| 342 | |||
| 343 |
2/2✓ Branch 0 taken 141613 times.
✓ Branch 1 taken 18026 times.
|
159639 | for (uint32_t col = 0; col < packed_bytes; col++) |
| 344 | { | ||
| 345 | /* Collect one byte (8 dims) from each of 32 vectors */ | ||
| 346 | uint8_t raw[VS_FASTSCAN_GROUP]; | ||
| 347 |
2/2✓ Branch 0 taken 4531616 times.
✓ Branch 1 taken 141613 times.
|
4673229 | for (uint32_t v = 0; v < VS_FASTSCAN_GROUP; v++) |
| 348 | { | ||
| 349 | 4531616 | uint32_t vi = g_start + v; | |
| 350 |
2/2✓ Branch 0 taken 3294741 times.
✓ Branch 1 taken 1236875 times.
|
7826357 | raw[v] = (vi < count) |
| 351 | 3294741 | ? bits_1bit[(size_t)vi * packed_bytes + col] | |
| 352 | : 0; | ||
| 353 | } | ||
| 354 | |||
| 355 | /* Split into upper and lower nibbles */ | ||
| 356 | uint8_t upper[VS_FASTSCAN_GROUP]; /* dims 4-7 = sq1 */ | ||
| 357 | uint8_t lower[VS_FASTSCAN_GROUP]; /* dims 0-3 = sq0 */ | ||
| 358 |
2/2✓ Branch 0 taken 4531616 times.
✓ Branch 1 taken 141613 times.
|
4673229 | for (uint32_t v = 0; v < VS_FASTSCAN_GROUP; v++) |
| 359 | { | ||
| 360 | 4531616 | upper[v] = raw[v] >> 4; | |
| 361 | 4531616 | lower[v] = raw[v] & 0x0F; | |
| 362 | } | ||
| 363 | |||
| 364 | /* Pack with kPerm0 interleaving. | ||
| 365 | * Lower nibble sq first (to match sequential LUT). */ | ||
| 366 | 141613 | uint8_t *out = codes_out + | |
| 367 | 141613 | (size_t)g * VS_FASTSCAN_GROUP_BYTES(dim) + | |
| 368 | 141613 | (size_t)col * VS_FASTSCAN_GROUP; | |
| 369 | |||
| 370 |
2/2✓ Branch 0 taken 2265808 times.
✓ Branch 1 taken 141613 times.
|
2407421 | for (uint32_t j = 0; j < 16; j++) |
| 371 | { | ||
| 372 | 2265808 | out[j] = lower[kPerm0[j]] | (lower[kPerm0[j] + 16] << 4); | |
| 373 | 2265808 | out[j + 16] = upper[kPerm0[j]] | (upper[kPerm0[j] + 16] << 4); | |
| 374 | } | ||
| 375 | } | ||
| 376 | } | ||
| 377 | |||
| 378 | 17986 | return ngroups; | |
| 379 | } | ||
| 380 | |||
| 381 | /* | ||
| 382 | * Inverse of vs_fastscan_pack_codes: reconstruct the per-vector 1-bit | ||
| 383 | * RaBitQ codes from the packed fastscan layout. Used when merging a | ||
| 384 | * fastscan posting page back into a builder (the parallel build folds | ||
| 385 | * workers' trailing partial pages into the head). bits_out must hold | ||
| 386 | * count * packed_bytes bytes. | ||
| 387 | */ | ||
| 388 | void | ||
| 389 | 2 | vs_fastscan_unpack_codes( | |
| 390 | const uint8_t *codes, uint32_t count, Dimension dim, uint8_t *bits_out) | ||
| 391 | { | ||
| 392 | 2 | uint32_t packed_bytes = (dim + 7) / 8; | |
| 393 | 2 | uint32_t ngroups = (count + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP; | |
| 394 | |||
| 395 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
|
6 | for (uint32_t g = 0; g < ngroups; g++) |
| 396 | { | ||
| 397 | 4 | uint32_t g_start = g * VS_FASTSCAN_GROUP; | |
| 398 | |||
| 399 |
2/2✓ Branch 0 taken 48 times.
✓ Branch 1 taken 4 times.
|
52 | for (uint32_t col = 0; col < packed_bytes; col++) |
| 400 | { | ||
| 401 | 48 | const uint8_t *in = codes + | |
| 402 | 48 | (size_t)g * VS_FASTSCAN_GROUP_BYTES(dim) + | |
| 403 | 48 | (size_t)col * VS_FASTSCAN_GROUP; | |
| 404 | |||
| 405 | /* Undo the kPerm0 nibble interleaving (inverse of the pack). */ | ||
| 406 | ✗ | uint8_t lower[VS_FASTSCAN_GROUP]; | |
| 407 | ✗ | uint8_t upper[VS_FASTSCAN_GROUP]; | |
| 408 |
2/2✓ Branch 0 taken 768 times.
✓ Branch 1 taken 48 times.
|
816 | for (uint32_t j = 0; j < 16; j++) |
| 409 | { | ||
| 410 | 768 | lower[kPerm0[j]] = in[j] & 0x0F; | |
| 411 | 768 | lower[kPerm0[j] + 16] = in[j] >> 4; | |
| 412 | 768 | upper[kPerm0[j]] = in[j + 16] & 0x0F; | |
| 413 | 768 | upper[kPerm0[j] + 16] = in[j + 16] >> 4; | |
| 414 | } | ||
| 415 | |||
| 416 |
2/2✓ Branch 0 taken 1536 times.
✓ Branch 1 taken 48 times.
|
1584 | for (uint32_t v = 0; v < VS_FASTSCAN_GROUP; v++) |
| 417 | { | ||
| 418 | 1536 | uint32_t vi = g_start + v; | |
| 419 |
2/2✓ Branch 0 taken 960 times.
✓ Branch 1 taken 576 times.
|
1536 | if (vi < count) |
| 420 | 960 | bits_out[(size_t)vi * packed_bytes + col] = | |
| 421 | 960 | (uint8_t)((upper[v] << 4) | lower[v]); | |
| 422 | } | ||
| 423 | } | ||
| 424 | } | ||
| 425 | 2 | } | |
| 426 | |||
| 427 | /* ---------------------------------------------------------------- | ||
| 428 | * Scalar accumulate kernel (reference implementation) | ||
| 429 | * | ||
| 430 | * Even/odd byte accumulation: add uint8 results as uint16, | ||
| 431 | * track upper byte contamination, subtract at end. | ||
| 432 | * ---------------------------------------------------------------- */ | ||
| 433 | |||
| 434 | static void | ||
| 435 | 20 | fastscan_accumulate_scalar( | |
| 436 | const uint8_t *codes, | ||
| 437 | const uint8_t *lut, | ||
| 438 | uint16_t *accum, | ||
| 439 | Dimension dim) | ||
| 440 | { | ||
| 441 | 20 | uint32_t code_length = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 442 | |||
| 443 | 20 | memset(accum, 0, VS_FASTSCAN_GROUP * sizeof(uint16_t)); | |
| 444 | |||
| 445 | /* Process 64 bytes per iteration: 4 sq blocks of 16 bytes each. | ||
| 446 | * Each 16-byte block has codes for one sq, with kPerm0 | ||
| 447 | * interleaving: byte j packs kPerm0[j] (lo) and kPerm0[j]+16 (hi). | ||
| 448 | * The corresponding LUT 16-byte block is the table for that sq. */ | ||
| 449 |
2/2✓ Branch 0 taken 2648 times.
✓ Branch 1 taken 20 times.
|
2668 | for (uint32_t i = 0; i < code_length; i += 16) |
| 450 | { | ||
| 451 | 2648 | const uint8_t *c = codes + i; | |
| 452 | 2648 | const uint8_t *tab = lut + i; | |
| 453 | |||
| 454 |
2/2✓ Branch 0 taken 42368 times.
✓ Branch 1 taken 2648 times.
|
45016 | for (uint32_t j = 0; j < 16; j++) |
| 455 | { | ||
| 456 | 42368 | uint8_t byte_val = c[j]; | |
| 457 | 42368 | uint8_t lo_nib = byte_val & 0x0F; | |
| 458 | 42368 | uint8_t hi_nib = byte_val >> 4; | |
| 459 | |||
| 460 | /* VPSHUFB per 128-bit lane: lookup in this sq's table */ | ||
| 461 | 42368 | accum[kPerm0[j]] += tab[lo_nib]; | |
| 462 | 42368 | accum[kPerm0[j] + 16] += tab[hi_nib]; | |
| 463 | } | ||
| 464 | } | ||
| 465 | 20 | } | |
| 466 | |||
| 467 | /* ---------------------------------------------------------------- | ||
| 468 | * SIMD dispatch | ||
| 469 | * ---------------------------------------------------------------- */ | ||
| 470 | |||
| 471 | typedef void (*VsFastscanAccumulateFn)( | ||
| 472 | const uint8_t *, const uint8_t *, uint16_t *, Dimension); | ||
| 473 | |||
| 474 | typedef void (*VsFastscanBuildLutFn)( | ||
| 475 | const float *, Dimension, uint8_t *, float *, float *); | ||
| 476 | |||
| 477 | static VsFastscanAccumulateFn g_fastscan_accumulate_fn = NULL; | ||
| 478 | static VsFastscanAccumulateHaccFn g_fastscan_accumulate_hacc_fn = NULL; | ||
| 479 | static VsFastscanBuildLutFn g_fastscan_build_lut_fn = NULL; | ||
| 480 | static VsFastscanBuildLutFn g_fastscan_build_lut_hacc_fn = NULL; | ||
| 481 | static atomic_bool g_fastscan_initialized = false; | ||
| 482 | |||
| 483 | void | ||
| 484 | 36 | vs_fastscan_reset_simd(void) | |
| 485 | { | ||
| 486 | 36 | atomic_store(&g_fastscan_initialized, false); | |
| 487 | 36 | } | |
| 488 | |||
| 489 | void | ||
| 490 | 151 | vs_fastscan_init_simd(void) | |
| 491 | { | ||
| 492 |
2/2✓ Branch 0 taken 110 times.
✓ Branch 1 taken 41 times.
|
151 | if (atomic_load(&g_fastscan_initialized)) |
| 493 | 2 | return; | |
| 494 | |||
| 495 | #ifdef VS_SIMD_FULL | ||
| 496 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 497 | 149 | SimdCapability caps = vs_detect_simd(); | |
| 498 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 149 times.
|
149 | if ((caps & VS_SIMD_AVX512_BW) == VS_SIMD_AVX512_BW) |
| 499 | { | ||
| 500 | ✗ | g_fastscan_accumulate_fn = vs_fastscan_accumulate_avx512; | |
| 501 | ✗ | g_fastscan_accumulate_hacc_fn = vs_fastscan_accumulate_hacc_avx512; | |
| 502 | ✗ | g_fastscan_build_lut_fn = vs_fastscan_build_lut_avx512; | |
| 503 | ✗ | g_fastscan_build_lut_hacc_fn = vs_fastscan_build_lut_hacc_avx512; | |
| 504 | ✗ | atomic_store(&g_fastscan_initialized, true); | |
| 505 | ✗ | return; | |
| 506 | } | ||
| 507 |
2/2✓ Branch 0 taken 143 times.
✓ Branch 1 taken 6 times.
|
149 | if (caps & SIMD_AVX2) |
| 508 | { | ||
| 509 | 143 | g_fastscan_accumulate_fn = vs_fastscan_accumulate_avx2; | |
| 510 | 143 | g_fastscan_accumulate_hacc_fn = vs_fastscan_accumulate_hacc_avx2; | |
| 511 | 143 | g_fastscan_build_lut_fn = fastscan_build_lut_scalar; | |
| 512 | 143 | g_fastscan_build_lut_hacc_fn = fastscan_build_lut_hacc_scalar; | |
| 513 | 143 | atomic_store(&g_fastscan_initialized, true); | |
| 514 | 143 | return; | |
| 515 | } | ||
| 516 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 517 | SimdCapability caps = vs_detect_simd(); | ||
| 518 | if (caps & SIMD_NEON) | ||
| 519 | { | ||
| 520 | g_fastscan_accumulate_fn = vs_fastscan_accumulate_neon; | ||
| 521 | g_fastscan_accumulate_hacc_fn = vs_fastscan_accumulate_hacc_neon; | ||
| 522 | g_fastscan_build_lut_fn = fastscan_build_lut_scalar; | ||
| 523 | g_fastscan_build_lut_hacc_fn = fastscan_build_lut_hacc_scalar; | ||
| 524 | atomic_store(&g_fastscan_initialized, true); | ||
| 525 | return; | ||
| 526 | } | ||
| 527 | #endif | ||
| 528 | #endif | ||
| 529 | |||
| 530 | 6 | g_fastscan_accumulate_fn = fastscan_accumulate_scalar; | |
| 531 | 6 | g_fastscan_accumulate_hacc_fn = fastscan_accumulate_hacc_scalar; | |
| 532 | 6 | g_fastscan_build_lut_fn = fastscan_build_lut_scalar; | |
| 533 | 6 | g_fastscan_build_lut_hacc_fn = fastscan_build_lut_hacc_scalar; | |
| 534 | 6 | atomic_store(&g_fastscan_initialized, true); | |
| 535 | } | ||
| 536 | |||
| 537 | void | ||
| 538 | 44 | vs_fastscan_build_lut( | |
| 539 | const float *transformed, | ||
| 540 | Dimension dim, | ||
| 541 | uint8_t *lut_out, | ||
| 542 | float *delta_out, | ||
| 543 | float *bias_out) | ||
| 544 | { | ||
| 545 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 42 times.
|
44 | if (vs_unlikely(!atomic_load(&g_fastscan_initialized))) |
| 546 | 2 | vs_fastscan_init_simd(); | |
| 547 | 44 | g_fastscan_build_lut_fn(transformed, dim, lut_out, delta_out, bias_out); | |
| 548 | 44 | } | |
| 549 | |||
| 550 | void | ||
| 551 | 394425 | vs_fastscan_build_lut_hacc( | |
| 552 | const float *transformed, | ||
| 553 | Dimension dim, | ||
| 554 | uint8_t *lut_out, | ||
| 555 | float *delta_out, | ||
| 556 | float *bias_out) | ||
| 557 | { | ||
| 558 |
2/2✓ Branch 0 taken 85 times.
✓ Branch 1 taken 394340 times.
|
394425 | if (vs_unlikely(!atomic_load(&g_fastscan_initialized))) |
| 559 | 85 | vs_fastscan_init_simd(); | |
| 560 | 394425 | g_fastscan_build_lut_hacc_fn( | |
| 561 | transformed, dim, lut_out, delta_out, bias_out); | ||
| 562 | 394425 | } | |
| 563 | |||
| 564 | void | ||
| 565 | 68 | vs_fastscan_accumulate( | |
| 566 | const uint8_t *codes, | ||
| 567 | const uint8_t *lut, | ||
| 568 | uint16_t *accum, | ||
| 569 | Dimension dim) | ||
| 570 | { | ||
| 571 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 68 times.
|
68 | if (vs_unlikely(!atomic_load(&g_fastscan_initialized))) |
| 572 | ✗ | vs_fastscan_init_simd(); | |
| 573 | 68 | g_fastscan_accumulate_fn(codes, lut, accum, dim); | |
| 574 | 68 | } | |
| 575 | |||
| 576 | VsFastscanAccumulateHaccFn | ||
| 577 | 237 | vs_fastscan_get_accumulate_hacc(void) | |
| 578 | { | ||
| 579 |
2/2✓ Branch 0 taken 26 times.
✓ Branch 1 taken 211 times.
|
237 | if (vs_unlikely(!atomic_load(&g_fastscan_initialized))) |
| 580 | 26 | vs_fastscan_init_simd(); | |
| 581 | 237 | return g_fastscan_accumulate_hacc_fn; | |
| 582 | } | ||
| 583 | |||
| 584 | void | ||
| 585 | 3670261 | vs_fastscan_accumulate_hacc( | |
| 586 | const uint8_t *codes, | ||
| 587 | const uint8_t *lut, | ||
| 588 | int32_t *accum, | ||
| 589 | Dimension dim) | ||
| 590 | { | ||
| 591 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 3670261 times.
|
3670261 | if (vs_unlikely(!atomic_load(&g_fastscan_initialized))) |
| 592 | ✗ | vs_fastscan_init_simd(); | |
| 593 | 3670261 | g_fastscan_accumulate_hacc_fn(codes, lut, accum, dim); | |
| 594 | 3670261 | } | |
| 595 | |||
| 596 | const char * | ||
| 597 | ✗ | vs_fastscan_impl_name(void) | |
| 598 | { | ||
| 599 | ✗ | if (!atomic_load(&g_fastscan_initialized)) | |
| 600 | ✗ | vs_fastscan_init_simd(); | |
| 601 | |||
| 602 | #ifdef VS_SIMD_FULL | ||
| 603 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 604 | ✗ | if (g_fastscan_accumulate_fn == vs_fastscan_accumulate_avx512) | |
| 605 | ✗ | return "avx512"; | |
| 606 | ✗ | if (g_fastscan_accumulate_fn == vs_fastscan_accumulate_avx2) | |
| 607 | ✗ | return "avx2"; | |
| 608 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 609 | if (g_fastscan_accumulate_fn == vs_fastscan_accumulate_neon) | ||
| 610 | return "neon"; | ||
| 611 | #endif | ||
| 612 | #endif | ||
| 613 | |||
| 614 | ✗ | return "scalar"; | |
| 615 | } | ||
| 616 | |||
| 617 | /* ---------------------------------------------------------------- | ||
| 618 | * Batch distance computation | ||
| 619 | * ---------------------------------------------------------------- */ | ||
| 620 | |||
| 621 | void | ||
| 622 | 2 | vs_fastscan_distance_batch( | |
| 623 | const RaBitQQueryState *qstate, | ||
| 624 | const float *f_add, | ||
| 625 | const float *f_rescale, | ||
| 626 | const uint8_t *codes, | ||
| 627 | uint32_t ngroups, | ||
| 628 | uint32_t count, | ||
| 629 | Dimension dim, | ||
| 630 | float *distances, | ||
| 631 | uint8_t *lut_buf, | ||
| 632 | uint16_t *accum_buf) | ||
| 633 | { | ||
| 634 |
4/8✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
|
2 | if (qstate == NULL || codes == NULL || distances == NULL || count == 0) |
| 635 | ✗ | return; | |
| 636 | |||
| 637 | /* Build LUT from transformed query */ | ||
| 638 | ✗ | float lut_delta, lut_bias; | |
| 639 | 2 | vs_fastscan_build_lut( | |
| 640 | 2 | qstate->transformed, dim, lut_buf, &lut_delta, &lut_bias); | |
| 641 | |||
| 642 | 2 | uint32_t group_bytes = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 643 | 2 | float g_add = qstate->g_add; | |
| 644 | 2 | float sum_t = qstate->sum_transformed; | |
| 645 | 2 | float inv_sqrt_d = qstate->inv_sqrt_d; | |
| 646 | |||
| 647 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
|
4 | for (uint32_t g = 0; g < ngroups; g++) |
| 648 | { | ||
| 649 | 2 | const uint8_t *group_codes = codes + (size_t)g * group_bytes; | |
| 650 | 2 | uint32_t g_start = g * VS_FASTSCAN_GROUP; | |
| 651 | 4 | uint32_t g_count = (g_start + VS_FASTSCAN_GROUP <= count) | |
| 652 | ? VS_FASTSCAN_GROUP | ||
| 653 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
|
2 | : count - g_start; |
| 654 | |||
| 655 | 2 | vs_fastscan_accumulate(group_codes, lut_buf, accum_buf, dim); | |
| 656 | |||
| 657 |
2/3✓ Branch 0 taken 64 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
|
66 | for (uint32_t v = 0; v < g_count; v++) |
| 658 | { | ||
| 659 | 64 | uint32_t vi = g_start + v; | |
| 660 | |||
| 661 | /* De-quantize: ip = accum * delta + nsq * vl */ | ||
| 662 | 64 | float binary_ip = (float)accum_buf[v] * lut_delta + lut_bias; | |
| 663 | |||
| 664 | 64 | float final_dot = (2.0f * binary_ip - sum_t) * inv_sqrt_d; | |
| 665 | 64 | distances[vi] = f_add[vi] + g_add - | |
| 666 | 64 | 2.0f * f_rescale[vi] * final_dot; | |
| 667 | } | ||
| 668 | } | ||
| 669 | } | ||
| 670 |