| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * posting_scan.c - Posting list scan with fused cluster function | ||
| 6 | * | ||
| 7 | * Scans posting data (paged or flat), computes batch RaBitQ | ||
| 8 | * distances, prunes via error bounds, and reranks survivors | ||
| 9 | * in a single function call per cluster. | ||
| 10 | * | ||
| 11 | * Optimizations: | ||
| 12 | * - Inline page access via stored base pointer (no vtable dispatch) | ||
| 13 | * - Score + prune + rerank fused into one function (no iterator | ||
| 14 | * overhead, compiler can optimize across all phases) | ||
| 15 | * - IP count padded to multiple of 4 (avoids tail kernel) | ||
| 16 | */ | ||
| 17 | |||
| 18 | #include "vs_config.h" | ||
| 19 | |||
| 20 | #include <string.h> | ||
| 21 | |||
| 22 | #include "core/log.h" | ||
| 23 | #include "core/memory.h" | ||
| 24 | #include "index/posting_scan.h" | ||
| 25 | #include "quant/fastscan.h" | ||
| 26 | |||
| 27 | #ifdef VS_SIMD_FULL | ||
| 28 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 29 | #include <immintrin.h> | ||
| 30 | |||
| 31 | #include "algo/simd_utils.h" | ||
| 32 | #endif | ||
| 33 | #endif | ||
| 34 | |||
| 35 | /* ---------------------------------------------------------------- | ||
| 36 | * Init / cleanup | ||
| 37 | * ---------------------------------------------------------------- */ | ||
| 38 | |||
| 39 | void | ||
| 40 | 18135 | prism_posting_scan_init( | |
| 41 | PrismPostingScan *scan, | ||
| 42 | VsStorage *storage, | ||
| 43 | char *page_base, | ||
| 44 | const RaBitQParams *params, | ||
| 45 | Dimension dim, | ||
| 46 | uint32_t max_entries_per_page) | ||
| 47 | { | ||
| 48 | 18135 | memset(scan, 0, sizeof(*scan)); | |
| 49 | 18135 | scan->storage = storage; | |
| 50 | 18135 | scan->page_base = page_base; | |
| 51 | 18135 | scan->params = params; | |
| 52 | 18135 | scan->dim = dim; | |
| 53 | 18135 | scan->packed_bytes = VS_RABITQ_BYTES(dim); | |
| 54 | 18135 | scan->cur_blkno = InvalidBlockNumber; | |
| 55 | |||
| 56 | /* Pre-allocate batch buffers — pad to multiple of 4 for IP kernel. | ||
| 57 | * max_entries_cap is the trusted per-page capacity these buffers are | ||
| 58 | * sized for; a page whose on-disk entry_count exceeds it (corruption) | ||
| 59 | * would overrun page_distances/page_scratch and is rejected on load. */ | ||
| 60 | 18135 | scan->max_entries_cap = max_entries_per_page; | |
| 61 | 18135 | uint32_t padded = (max_entries_per_page + 3) & ~3u; | |
| 62 | 18135 | scan->page_distances = vs_alloc(padded * sizeof(Distance)); | |
| 63 | 18135 | scan->page_scratch = vs_alloc(padded * sizeof(float)); | |
| 64 | 18135 | } | |
| 65 | |||
| 66 | void | ||
| 67 | 18133 | prism_posting_scan_cleanup(PrismPostingScan *scan) | |
| 68 | { | ||
| 69 | /* Release pinned page if still held */ | ||
| 70 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 18133 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
18133 | if (scan->cur_page != NULL && scan->storage != NULL && |
| 71 | ✗ | scan->page_base == NULL) | |
| 72 | { | ||
| 73 | ✗ | vs_storage_release_page(scan->storage, scan->cur_blkno); | |
| 74 | ✗ | scan->cur_page = NULL; | |
| 75 | } | ||
| 76 | |||
| 77 | 18133 | vs_free(scan->page_distances); | |
| 78 | 18133 | vs_free(scan->page_scratch); | |
| 79 | 18133 | scan->page_distances = NULL; | |
| 80 | 18133 | scan->page_scratch = NULL; | |
| 81 | |||
| 82 |
2/2✓ Branch 0 taken 235 times.
✓ Branch 1 taken 17898 times.
|
18133 | if (scan->fs_lut != NULL) |
| 83 | { | ||
| 84 | 235 | vs_free(scan->fs_lut); | |
| 85 | 235 | scan->fs_lut = NULL; | |
| 86 | } | ||
| 87 |
2/2✓ Branch 0 taken 235 times.
✓ Branch 1 taken 17898 times.
|
18133 | if (scan->fs_accum != NULL) |
| 88 | { | ||
| 89 | 235 | vs_free(scan->fs_accum); | |
| 90 | 235 | scan->fs_accum = NULL; | |
| 91 | } | ||
| 92 | 18133 | } | |
| 93 | |||
| 94 | void | ||
| 95 | 237 | prism_posting_scan_enable_fastscan(PrismPostingScan *scan, int lut_bits) | |
| 96 | { | ||
| 97 | 237 | scan->fs_lut_bits = lut_bits; | |
| 98 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 237 times.
|
237 | if (lut_bits == 8) |
| 99 | ✗ | scan->fs_lut = vs_alloc0(VS_FASTSCAN_LUT_BYTES(scan->dim)); | |
| 100 | else | ||
| 101 | 237 | scan->fs_lut = vs_alloc0(VS_FASTSCAN_LUT_HACC_BYTES(scan->dim)); | |
| 102 | 237 | scan->fs_accum = vs_alloc(VS_FASTSCAN_GROUP * sizeof(int32_t)); | |
| 103 | |||
| 104 | /* Resolve dispatch once per scan (see posting_scan.h). */ | ||
| 105 | 237 | scan->fs_accum_hacc = vs_fastscan_get_accumulate_hacc(); | |
| 106 | #if defined(VS_SIMD_FULL) && (defined(__x86_64__) || defined(_M_X64)) | ||
| 107 | 237 | scan->fs_has_avx512 = vs_has_all_simd(VS_SIMD_AVX512_BW); | |
| 108 | #else | ||
| 109 | scan->fs_has_avx512 = false; | ||
| 110 | #endif | ||
| 111 | 237 | scan->fs_max_groups_first = prism_fastscan_max_groups(scan->dim, true); | |
| 112 | 237 | scan->fs_max_groups_over = prism_fastscan_max_groups(scan->dim, false); | |
| 113 | 237 | } | |
| 114 | |||
| 115 | static bool advance_page(PrismPostingScan *scan); | ||
| 116 | |||
| 117 | /* ---------------------------------------------------------------- | ||
| 118 | * Per-cluster begin / end | ||
| 119 | * ---------------------------------------------------------------- */ | ||
| 120 | |||
| 121 | void | ||
| 122 | 44572 | prism_posting_scan_begin_cluster( | |
| 123 | PrismPostingScan *scan, | ||
| 124 | RaBitQQueryState *qstate, | ||
| 125 | BlockNumber posting_head) | ||
| 126 | { | ||
| 127 | 44572 | scan->qstate = qstate; | |
| 128 | 44572 | scan->cur_blkno = posting_head; | |
| 129 | 44572 | scan->cur_page = NULL; | |
| 130 | 44572 | scan->cur_content = NULL; | |
| 131 | 44572 | scan->cur_max_entries = 0; | |
| 132 | 44572 | scan->cur_count = 0; | |
| 133 | 44572 | scan->fs_lut_valid = false; | |
| 134 | 44572 | scan->pages_read = 0; | |
| 135 | 44572 | scan->pages_skipped = 0; | |
| 136 | 44572 | scan->entries_scanned = 0; | |
| 137 | 44572 | scan->entries_pruned = 0; | |
| 138 | |||
| 139 | /* Eagerly read the first page so pt_centroid is accessible | ||
| 140 | * via prism_posting_pt_centroid(scan->cur_page) before _cluster | ||
| 141 | * is called. */ | ||
| 142 |
1/2✓ Branch 0 taken 44572 times.
✗ Branch 1 not taken.
|
44572 | if (posting_head != InvalidBlockNumber) |
| 143 | 44572 | advance_page(scan); | |
| 144 | 44572 | } | |
| 145 | |||
| 146 | /* | ||
| 147 | * Return pt_centroid from the first page (must be called after | ||
| 148 | * begin_cluster). Returns NULL if no page was loaded. | ||
| 149 | */ | ||
| 150 | const float * | ||
| 151 | 44544 | prism_posting_scan_pt_centroid(const PrismPostingScan *scan) | |
| 152 | { | ||
| 153 |
2/2✓ Branch 0 taken 10486 times.
✓ Branch 1 taken 34058 times.
|
44544 | if (scan->cur_page == NULL) |
| 154 | ✗ | return NULL; | |
| 155 | 44544 | return prism_posting_pt_centroid(scan->cur_page); | |
| 156 | } | ||
| 157 | |||
| 158 | void | ||
| 159 | 220 | prism_posting_scan_begin_flat( | |
| 160 | PrismPostingScan *scan, RaBitQQueryState *qstate, char *flat_buf) | ||
| 161 | { | ||
| 162 | 220 | PrismFlatPostingHeader *hdr = prism_flat_posting_header(flat_buf); | |
| 163 | |||
| 164 | 220 | scan->qstate = qstate; | |
| 165 | 220 | scan->cur_blkno = InvalidBlockNumber; | |
| 166 | 220 | scan->cur_page = flat_buf; | |
| 167 |
0/2✗ Branch 0 not taken.
✗ Branch 1 not taken.
|
220 | scan->cur_content = prism_flat_posting_content(flat_buf); |
| 168 | 220 | scan->cur_max_entries = hdr->max_entries; | |
| 169 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 220 times.
|
220 | if (hdr->entry_count > scan->max_entries_cap) |
| 170 | ✗ | vs_error( | |
| 171 | VS_EXTENSION_NAME | ||
| 172 | ": flat posting page has an invalid entry count (%u > %u); " | ||
| 173 | "the index may be corrupted -- REINDEX it", | ||
| 174 | (unsigned)hdr->entry_count, | ||
| 175 | scan->max_entries_cap); | ||
| 176 | 220 | scan->cur_count = hdr->entry_count; | |
| 177 | 220 | scan->pages_read = 1; | |
| 178 | 220 | scan->pages_skipped = 0; | |
| 179 | 220 | scan->entries_scanned = 0; | |
| 180 | 220 | scan->entries_pruned = 0; | |
| 181 | 220 | } | |
| 182 | |||
| 183 | void | ||
| 184 | 44792 | prism_posting_scan_end_cluster(PrismPostingScan *scan) | |
| 185 | { | ||
| 186 | /* Release current page if held via storage vtable */ | ||
| 187 |
3/4✓ Branch 0 taken 9092 times.
✓ Branch 1 taken 35700 times.
✓ Branch 2 taken 9092 times.
✗ Branch 3 not taken.
|
44792 | if (scan->cur_page != NULL && scan->storage != NULL && |
| 188 |
2/2✓ Branch 0 taken 1016 times.
✓ Branch 1 taken 8076 times.
|
9092 | scan->page_base == NULL) |
| 189 | { | ||
| 190 | 1016 | vs_storage_release_page(scan->storage, scan->cur_blkno); | |
| 191 | } | ||
| 192 | |||
| 193 | 44792 | scan->cur_page = NULL; | |
| 194 | 44792 | scan->cur_content = NULL; | |
| 195 | 44792 | scan->qstate = NULL; | |
| 196 | 44792 | } | |
| 197 | |||
| 198 | /* ---------------------------------------------------------------- | ||
| 199 | * Advance to next page in chain | ||
| 200 | * ---------------------------------------------------------------- */ | ||
| 201 | |||
| 202 | static bool | ||
| 203 | 83286 | advance_page(PrismPostingScan *scan) | |
| 204 | { | ||
| 205 | /* Follow chain from current page */ | ||
| 206 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83286 times.
|
83286 | if (scan->cur_page != NULL) |
| 207 | { | ||
| 208 | ✗ | if (scan->storage != NULL) | |
| 209 | { | ||
| 210 | ✗ | scan->cur_blkno = prism_posting_opaque(scan->cur_page)->next_blkno; | |
| 211 | } | ||
| 212 | else | ||
| 213 | { | ||
| 214 | /* Flat mode: single page, no chain */ | ||
| 215 | ✗ | scan->cur_blkno = InvalidBlockNumber; | |
| 216 | } | ||
| 217 | ✗ | scan->cur_page = NULL; | |
| 218 | ✗ | scan->cur_content = NULL; | |
| 219 | } | ||
| 220 | |||
| 221 |
2/2✓ Branch 0 taken 35700 times.
✓ Branch 1 taken 47586 times.
|
83286 | if (scan->cur_blkno == InvalidBlockNumber) |
| 222 | 26230 | return false; | |
| 223 | |||
| 224 | /* Read next page — inline for ArrayPageStorage */ | ||
| 225 |
2/2✓ Branch 0 taken 36568 times.
✓ Branch 1 taken 11018 times.
|
47586 | if (scan->page_base != NULL) |
| 226 | 36568 | scan->cur_page = scan->page_base + (size_t)scan->cur_blkno * BLCKSZ; | |
| 227 | else | ||
| 228 | 11018 | scan->cur_page = vs_storage_read_page(scan->storage, scan->cur_blkno); | |
| 229 | |||
| 230 | 47586 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 231 | |||
| 232 | /* Validate page identity — catch corrupted chain pointers early */ | ||
| 233 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 47586 times.
|
47586 | if (opaque->page_id != PRISM_POSTING_PAGE_ID) |
| 234 | { | ||
| 235 | ✗ | vs_warn(VS_EXTENSION_NAME ": posting scan hit non-posting page " | |
| 236 | "(blkno=%u, page_id=0x%04X)", | ||
| 237 | scan->cur_blkno, | ||
| 238 | opaque->page_id); | ||
| 239 | /* Release the pin taken above before bailing (storage-backed only; | ||
| 240 | * in page_base mode the page is a borrowed pointer, not a pin). */ | ||
| 241 | ✗ | if (scan->storage != NULL && scan->page_base == NULL) | |
| 242 | ✗ | vs_storage_release_page(scan->storage, scan->cur_blkno); | |
| 243 | ✗ | scan->cur_page = NULL; | |
| 244 | ✗ | scan->cur_content = NULL; | |
| 245 | ✗ | return false; | |
| 246 | } | ||
| 247 | |||
| 248 | 58574 | scan->cur_content = | |
| 249 | 47586 | (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 250 | 44572 | ? prism_posting_content_first(scan->cur_page, scan->dim) | |
| 251 |
2/2✓ Branch 0 taken 44572 times.
✓ Branch 1 taken 3014 times.
|
47586 | : prism_posting_content(scan->cur_page); |
| 252 | 47586 | scan->cur_max_entries = opaque->max_entries; | |
| 253 | /* Reject an entry_count past what a valid page of this format can | ||
| 254 | * hold. AoS pages drive the page_distances/page_scratch batch (sized | ||
| 255 | * to max_entries_cap), so an over-count there is an out-of-bounds | ||
| 256 | * write; fastscan pages are scored per 32-vector group and hold more | ||
| 257 | * entries per page, so bound them by the fastscan capacity instead of | ||
| 258 | * the smaller AoS one. */ | ||
| 259 | 10988 | uint32_t page_cap; | |
| 260 |
2/2✓ Branch 0 taken 11274 times.
✓ Branch 1 taken 36312 times.
|
47586 | if (opaque->flags & PRISM_POSTING_PAGE_FASTSCAN) |
| 261 | 11274 | page_cap = (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 262 | 10174 | ? prism_fastscan_max_entries_first(scan->dim) | |
| 263 |
2/2✓ Branch 0 taken 10174 times.
✓ Branch 1 taken 1100 times.
|
11274 | : prism_fastscan_max_entries(scan->dim); |
| 264 | else | ||
| 265 | 36312 | page_cap = scan->max_entries_cap; | |
| 266 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 47586 times.
|
47586 | if (opaque->entry_count > page_cap) |
| 267 | ✗ | vs_error( | |
| 268 | VS_EXTENSION_NAME | ||
| 269 | ": posting page %u has an invalid entry count (%u > %u); " | ||
| 270 | "the index may be corrupted -- REINDEX it", | ||
| 271 | scan->cur_blkno, | ||
| 272 | (unsigned)opaque->entry_count, | ||
| 273 | page_cap); | ||
| 274 | 47586 | scan->cur_count = opaque->entry_count; | |
| 275 | 47586 | scan->pages_read++; | |
| 276 | 47586 | return true; | |
| 277 | } | ||
| 278 | |||
| 279 | /* ---------------------------------------------------------------- | ||
| 280 | * Score + prune a full cluster into topk | ||
| 281 | * | ||
| 282 | * Inserts survivors with approximate (estimated) distances and | ||
| 283 | * error bounds. No exact reranking — the caller does that after | ||
| 284 | * extracting candidates from topk. | ||
| 285 | * ---------------------------------------------------------------- */ | ||
| 286 | |||
| 287 | void | ||
| 288 | 26679 | prism_posting_scan_cluster(PrismPostingScan *scan, VsTopK *topk) | |
| 289 | { | ||
| 290 | 26679 | Dimension dim = scan->dim; | |
| 291 | 26679 | uint32_t entry_size = PRISM_POSTING_ENTRY_SIZE(dim); | |
| 292 | |||
| 293 | 26679 | float g_add = scan->qstate->g_add; | |
| 294 | 26679 | float sum_t = scan->qstate->sum_transformed; | |
| 295 | 26679 | float inv_sqrt_d = scan->qstate->inv_sqrt_d; | |
| 296 | 26679 | float g_error = scan->qstate->g_error; | |
| 297 | |||
| 298 | 26679 | Distance *distances = scan->page_distances; | |
| 299 | 26679 | float *scratch = scan->page_scratch; | |
| 300 | |||
| 301 | /* Process pages until chain is exhausted. | ||
| 302 | * For flat mode: single iteration (one page, no chain). */ | ||
| 303 | 4755 | for (;;) | |
| 304 | 28472 | { | |
| 305 | /* Advance to next page */ | ||
| 306 |
2/2✓ Branch 0 taken 28472 times.
✓ Branch 1 taken 26679 times.
|
55151 | if (scan->cur_page == NULL) |
| 307 | { | ||
| 308 |
2/2✓ Branch 1 taken 25499 times.
✓ Branch 2 taken 2973 times.
|
28472 | if (!advance_page(scan)) |
| 309 | 25256 | break; | |
| 310 | } | ||
| 311 | |||
| 312 | /* A page tombstoned by VACUUM (all entries dead) is skipped — no | ||
| 313 | * scoring kernel; the chain-follow below still advances past it. A | ||
| 314 | * retired chain also carries TOMBSTONED once reclaimed, but its | ||
| 315 | * entries moved rather than died, so it is still scored. */ | ||
| 316 |
3/3✓ Branch 1 taken 1658 times.
✓ Branch 2 taken 26810 times.
✓ Branch 3 taken 4 times.
|
28472 | if (!prism_posting_page_all_dead(prism_posting_opaque(scan->cur_page))) |
| 317 | { | ||
| 318 | 28460 | char *content = scan->cur_content; | |
| 319 | 28460 | uint32_t count = scan->cur_count; | |
| 320 | |||
| 321 | /* --- Score: batch IP over all entries on this page --- | ||
| 322 | * | ||
| 323 | * AoS layout: entry i's bits live at content + i * entry_size | ||
| 324 | * + PRISM_POSTING_ENTRY_BITS_OFFSET. The SIMD kernel just needs | ||
| 325 | * the first entry's bits pointer and a stride of entry_size. */ | ||
| 326 | 28460 | const uint8_t *bits_base = prism_posting_first_bits(content); | |
| 327 | |||
| 328 | 28460 | uint32_t padded = (count + 3) & ~3u; | |
| 329 | 28460 | vs_rabitq_inner_product_multi( | |
| 330 | 28460 | scan->qstate->transformed, | |
| 331 | bits_base, | ||
| 332 | entry_size, | ||
| 333 | dim, | ||
| 334 | padded, | ||
| 335 | scratch); | ||
| 336 | |||
| 337 | /* Convert raw IPs to distances, reading f_add/f_rescale per | ||
| 338 | * entry via the strided AoS accessor. */ | ||
| 339 |
3/3✓ Branch 0 taken 2749772 times.
✓ Branch 1 taken 84349 times.
✓ Branch 2 taken 1658 times.
|
2835779 | for (uint32_t i = 0; i < count; i++) |
| 340 | { | ||
| 341 | 57547 | PrismPostingEntryHeader *e = | |
| 342 | 2807319 | prism_posting_entry_at(content, i, dim); | |
| 343 | 2807319 | float final_dot = (2.0f * scratch[i] - sum_t) * inv_sqrt_d; | |
| 344 | 2807319 | distances[i] = e->f_add + g_add - | |
| 345 | 2807319 | 2.0f * e->f_rescale * final_dot; | |
| 346 | } | ||
| 347 | |||
| 348 | /* --- Prune + insert approximate distances --- */ | ||
| 349 |
2/2✓ Branch 0 taken 1583 times.
✓ Branch 1 taken 75 times.
|
28460 | Distance threshold = vs_topk_threshold(topk); |
| 350 | |||
| 351 |
2/2✓ Branch 0 taken 2807319 times.
✓ Branch 1 taken 28460 times.
|
2835779 | for (uint32_t i = 0; i < count; i++) |
| 352 | { | ||
| 353 | 57547 | PrismPostingEntryHeader *e = | |
| 354 |
2/2✓ Branch 0 taken 2412 times.
✓ Branch 1 taken 55135 times.
|
2807319 | prism_posting_entry_at(content, i, dim); |
| 355 | 2807319 | scan->entries_scanned++; | |
| 356 | |||
| 357 |
2/2✓ Branch 0 taken 2418 times.
✓ Branch 1 taken 2804901 times.
|
2807319 | if (e->meta.flags & PRISM_POSTING_FLAG_DELETED) |
| 358 | { | ||
| 359 | 2418 | scan->entries_pruned++; | |
| 360 | 2418 | continue; | |
| 361 | } | ||
| 362 | |||
| 363 | 2804901 | Distance est = distances[i]; | |
| 364 | 2804901 | Distance err = e->f_error * g_error; | |
| 365 | 2804901 | Distance lb = est - err; | |
| 366 | |||
| 367 |
2/2✓ Branch 0 taken 2339688 times.
✓ Branch 1 taken 465213 times.
|
2804901 | if (lb >= threshold) |
| 368 | { | ||
| 369 | 2339688 | scan->entries_pruned++; | |
| 370 | 2339688 | continue; | |
| 371 | } | ||
| 372 | |||
| 373 | 465213 | uint64_t id = prism_posting_encode_tid(&e->meta.tid); | |
| 374 | 465213 | vs_topk_insert(topk, est, err, id); | |
| 375 |
2/2✓ Branch 0 taken 2998 times.
✓ Branch 1 taken 1401 times.
|
521359 | threshold = vs_topk_threshold(topk); |
| 376 | } | ||
| 377 | } /* end: page not tombstoned */ | ||
| 378 | else | ||
| 379 | 12 | scan->pages_skipped++; | |
| 380 | |||
| 381 | /* Follow chain: read next_blkno, then release current */ | ||
| 382 | 28472 | BlockNumber prev_blkno = scan->cur_blkno; | |
| 383 | |||
| 384 |
2/2✓ Branch 0 taken 28252 times.
✓ Branch 1 taken 220 times.
|
28472 | if (scan->storage != NULL) |
| 385 | 28252 | scan->cur_blkno = prism_posting_opaque(scan->cur_page)->next_blkno; | |
| 386 | else | ||
| 387 | 220 | scan->cur_blkno = InvalidBlockNumber; | |
| 388 | |||
| 389 |
4/4✓ Branch 0 taken 28252 times.
✓ Branch 1 taken 220 times.
✓ Branch 2 taken 1690 times.
✓ Branch 3 taken 26562 times.
|
28472 | if (scan->storage != NULL && scan->page_base == NULL) |
| 390 | 1690 | vs_storage_release_page(scan->storage, prev_blkno); | |
| 391 | |||
| 392 | 28472 | scan->cur_page = NULL; | |
| 393 | 28472 | scan->cur_content = NULL; | |
| 394 | } | ||
| 395 | 26679 | } | |
| 396 | |||
| 397 | /* ---------------------------------------------------------------- | ||
| 398 | * Fastscan: vectorized distance + prune for 16 entries | ||
| 399 | * ---------------------------------------------------------------- */ | ||
| 400 | |||
| 401 | #ifdef VS_SIMD_FULL | ||
| 402 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 403 | |||
| 404 | /* | ||
| 405 | * Compute distances and lower bounds for 16 entries using AVX-512. | ||
| 406 | * Returns a bitmask of survivors (lb < threshold). | ||
| 407 | */ | ||
| 408 | VS_TARGET_AVX512 static inline __mmask16 | ||
| 409 | ✗ | fastscan_prune_16( | |
| 410 | const int32_t *accum, | ||
| 411 | const float *f_add, | ||
| 412 | const float *f_rescale, | ||
| 413 | const float *f_error, | ||
| 414 | float *est_out, | ||
| 415 | float *err_out, | ||
| 416 | float lut_scale, | ||
| 417 | float lut_bias, | ||
| 418 | float sum_t, | ||
| 419 | float inv_sqrt_d, | ||
| 420 | float g_add, | ||
| 421 | float g_error, | ||
| 422 | float threshold) | ||
| 423 | { | ||
| 424 | /* Convert 16 int32 accumulators to float */ | ||
| 425 | ✗ | __m512i acc32 = _mm512_loadu_si512((const __m512i *)accum); | |
| 426 | ✗ | __m512 acc_f = _mm512_cvtepi32_ps(acc32); | |
| 427 | |||
| 428 | /* binary_ip = acc * scale + bias */ | ||
| 429 | ✗ | __m512 scale_v = _mm512_set1_ps(lut_scale); | |
| 430 | ✗ | __m512 bias_v = _mm512_set1_ps(lut_bias); | |
| 431 | ✗ | __m512 ip = _mm512_fmadd_ps(acc_f, scale_v, bias_v); | |
| 432 | |||
| 433 | /* final_dot = (2*ip - sum_t) * inv_sqrt_d */ | ||
| 434 | ✗ | __m512 two_v = _mm512_set1_ps(2.0f); | |
| 435 | ✗ | __m512 sum_t_v = _mm512_set1_ps(sum_t); | |
| 436 | ✗ | __m512 inv_sqrt_v = _mm512_set1_ps(inv_sqrt_d); | |
| 437 | ✗ | __m512 final_dot = | |
| 438 | ✗ | _mm512_mul_ps(_mm512_fmsub_ps(two_v, ip, sum_t_v), inv_sqrt_v); | |
| 439 | |||
| 440 | /* est = f_add + g_add - 2 * f_rescale * final_dot */ | ||
| 441 | ✗ | __m512 fa = _mm512_loadu_ps(f_add); | |
| 442 | ✗ | __m512 fr = _mm512_loadu_ps(f_rescale); | |
| 443 | ✗ | __m512 gadd = _mm512_set1_ps(g_add); | |
| 444 | ✗ | __m512 est_v = _mm512_add_ps(fa, gadd); | |
| 445 | ✗ | est_v = _mm512_fnmadd_ps(_mm512_mul_ps(two_v, fr), final_dot, est_v); | |
| 446 | |||
| 447 | /* err = f_error * g_error */ | ||
| 448 | ✗ | __m512 fe = _mm512_loadu_ps(f_error); | |
| 449 | ✗ | __m512 gerr_v = _mm512_set1_ps(g_error); | |
| 450 | ✗ | __m512 err_v = _mm512_mul_ps(fe, gerr_v); | |
| 451 | |||
| 452 | /* lb = est - err; survivors = lb < threshold */ | ||
| 453 | ✗ | __m512 lb_v = _mm512_sub_ps(est_v, err_v); | |
| 454 | ✗ | __m512 thr_v = _mm512_set1_ps(threshold); | |
| 455 | ✗ | __mmask16 surv = _mm512_cmp_ps_mask(lb_v, thr_v, _CMP_LT_OS); | |
| 456 | |||
| 457 | ✗ | _mm512_storeu_ps(est_out, est_v); | |
| 458 | ✗ | _mm512_storeu_ps(err_out, err_v); | |
| 459 | |||
| 460 | ✗ | return surv; | |
| 461 | } | ||
| 462 | |||
| 463 | /* | ||
| 464 | * AVX-512 vectorized distance + prune for one 32-vector group. | ||
| 465 | * Must be in a target-attributed function so the compiler only | ||
| 466 | * emits AVX-512 instructions here, not in the caller. | ||
| 467 | */ | ||
| 468 | VS_TARGET_AVX512 static void | ||
| 469 | ✗ | fastscan_prune_group_avx512( | |
| 470 | PrismPostingScan *scan, | ||
| 471 | VsTopK *topk, | ||
| 472 | ItemPointerData *tids, | ||
| 473 | const float *f_add, | ||
| 474 | const float *f_rescale, | ||
| 475 | const float *f_error, | ||
| 476 | float lut_scale, | ||
| 477 | float lut_bias, | ||
| 478 | float sum_t, | ||
| 479 | float inv_sqrt_d, | ||
| 480 | float g_add, | ||
| 481 | float g_error, | ||
| 482 | Distance *threshold_p) | ||
| 483 | { | ||
| 484 | ✗ | Distance threshold = *threshold_p; | |
| 485 | ✗ | float est_buf[VS_FASTSCAN_GROUP]; | |
| 486 | ✗ | float err_buf[VS_FASTSCAN_GROUP]; | |
| 487 | |||
| 488 | ✗ | for (uint32_t h = 0; h < 2; h++) | |
| 489 | { | ||
| 490 | ✗ | uint32_t off = h * 16; | |
| 491 | ✗ | __mmask16 surv = fastscan_prune_16( | |
| 492 | ✗ | scan->fs_accum + off, | |
| 493 | ✗ | f_add + off, | |
| 494 | ✗ | f_rescale + off, | |
| 495 | ✗ | f_error + off, | |
| 496 | ✗ | est_buf + off, | |
| 497 | ✗ | err_buf + off, | |
| 498 | lut_scale, | ||
| 499 | lut_bias, | ||
| 500 | sum_t, | ||
| 501 | inv_sqrt_d, | ||
| 502 | g_add, | ||
| 503 | g_error, | ||
| 504 | threshold); | ||
| 505 | |||
| 506 | ✗ | if (surv == 0) | |
| 507 | { | ||
| 508 | ✗ | scan->entries_pruned += 16; | |
| 509 | ✗ | continue; | |
| 510 | } | ||
| 511 | |||
| 512 | ✗ | scan->entries_pruned += 16 - _mm_popcnt_u32(surv); | |
| 513 | |||
| 514 | ✗ | while (surv != 0) | |
| 515 | { | ||
| 516 | ✗ | uint32_t v = off + __builtin_ctz(surv); | |
| 517 | ✗ | surv &= surv - 1; | |
| 518 | |||
| 519 | ✗ | uint64_t id = prism_posting_encode_tid(&tids[v]); | |
| 520 | ✗ | vs_topk_insert(topk, est_buf[v], err_buf[v], id); | |
| 521 | ✗ | threshold = vs_topk_threshold(topk); | |
| 522 | } | ||
| 523 | } | ||
| 524 | |||
| 525 | ✗ | *threshold_p = threshold; | |
| 526 | ✗ | } | |
| 527 | |||
| 528 | #endif | ||
| 529 | #endif | ||
| 530 | |||
| 531 | /* ---------------------------------------------------------------- | ||
| 532 | * Fastscan variant: VPSHUFB kernel for fastscan-format pages | ||
| 533 | * | ||
| 534 | * Processes 32-vector groups via table-lookup accumulation. | ||
| 535 | * Falls back to the AoS kernel for non-fastscan pages (mixed | ||
| 536 | * chains). | ||
| 537 | * ---------------------------------------------------------------- */ | ||
| 538 | |||
| 539 | static void | ||
| 540 | 10240 | scan_fastscan_page(PrismPostingScan *scan, VsTopK *topk) | |
| 541 | { | ||
| 542 | 10240 | Dimension dim = scan->dim; | |
| 543 | 10240 | uint32_t count = scan->cur_count; | |
| 544 | 10240 | char *content = scan->cur_content; | |
| 545 | |||
| 546 | 10240 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 547 | |||
| 548 | /* Fall back to AoS kernel for non-fastscan pages */ | ||
| 549 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10240 times.
|
10240 | if (!(opaque->flags & PRISM_POSTING_PAGE_FASTSCAN)) |
| 550 | ✗ | return; /* caller handles AoS via the standard path */ | |
| 551 | |||
| 552 | /* Build LUT from transformed query (once per cluster) */ | ||
| 553 |
2/2✓ Branch 0 taken 9140 times.
✓ Branch 1 taken 1100 times.
|
10240 | if (!scan->fs_lut_valid) |
| 554 | { | ||
| 555 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 9140 times.
|
9140 | if (scan->fs_lut_bits == 8) |
| 556 | ✗ | vs_fastscan_build_lut( | |
| 557 | ✗ | scan->qstate->transformed, | |
| 558 | dim, | ||
| 559 | scan->fs_lut, | ||
| 560 | &scan->fs_lut_scale, | ||
| 561 | &scan->fs_lut_bias); | ||
| 562 | else | ||
| 563 | 9140 | vs_fastscan_build_lut_hacc( | |
| 564 | 9140 | scan->qstate->transformed, | |
| 565 | dim, | ||
| 566 | scan->fs_lut, | ||
| 567 | &scan->fs_lut_scale, | ||
| 568 | &scan->fs_lut_bias); | ||
| 569 | 9140 | scan->fs_lut_valid = true; | |
| 570 | } | ||
| 571 | 10240 | float lut_scale = scan->fs_lut_scale; | |
| 572 | 10240 | float lut_bias = scan->fs_lut_bias; | |
| 573 | |||
| 574 | 10240 | float g_add = scan->qstate->g_add; | |
| 575 | 10240 | float sum_t = scan->qstate->sum_transformed; | |
| 576 | 10240 | float inv_sqrt_d = scan->qstate->inv_sqrt_d; | |
| 577 | 10240 | float g_error = scan->qstate->g_error; | |
| 578 | 20480 | uint32_t max_groups = (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 579 | ? scan->fs_max_groups_first | ||
| 580 |
2/2✓ Branch 0 taken 9140 times.
✓ Branch 1 taken 1100 times.
|
10240 | : scan->fs_max_groups_over; |
| 581 | 10240 | uint32_t ngroups = (count + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP; | |
| 582 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1934 times.
|
10240 | if (ngroups > max_groups) |
| 583 | ✗ | ngroups = max_groups; | |
| 584 | |||
| 585 |
2/2✓ Branch 0 taken 7803 times.
✓ Branch 1 taken 503 times.
|
10240 | Distance threshold = vs_topk_threshold(topk); |
| 586 | |||
| 587 |
2/2✓ Branch 0 taken 27321 times.
✓ Branch 1 taken 10240 times.
|
37561 | for (uint32_t g = 0; g < ngroups; g++) |
| 588 | { | ||
| 589 | 27321 | uint32_t g_start = g * VS_FASTSCAN_GROUP; | |
| 590 | 27321 | uint32_t g_count = count - g_start; | |
| 591 |
2/2✓ Branch 0 taken 16338 times.
✓ Branch 1 taken 1934 times.
|
27321 | if (g_count > VS_FASTSCAN_GROUP) |
| 592 | 16338 | g_count = VS_FASTSCAN_GROUP; | |
| 593 | |||
| 594 | /* Access group data — compute base once */ | ||
| 595 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 9049 times.
|
27321 | char *gbase = prism_fastscan_group_base(content, g, dim); |
| 596 | |||
| 597 | 27321 | ItemPointerData *tids = (ItemPointerData *)gbase; | |
| 598 | 27321 | float *f_add = (float *)(gbase + | |
| 599 | VS_FASTSCAN_GROUP * sizeof(ItemPointerData)); | ||
| 600 | 27321 | float *f_rescale = f_add + VS_FASTSCAN_GROUP; | |
| 601 | 27321 | float *f_error = f_rescale + VS_FASTSCAN_GROUP; | |
| 602 | 27321 | uint8_t *codes = (uint8_t *)(f_error + VS_FASTSCAN_GROUP); | |
| 603 | |||
| 604 | /* Run VPSHUFB accumulate kernel */ | ||
| 605 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 27321 times.
|
27321 | if (scan->fs_lut_bits == 8) |
| 606 | { | ||
| 607 | ✗ | uint16_t acc16[VS_FASTSCAN_GROUP]; | |
| 608 | ✗ | vs_fastscan_accumulate(codes, scan->fs_lut, acc16, dim); | |
| 609 | ✗ | for (uint32_t v2 = 0; v2 < VS_FASTSCAN_GROUP; v2++) | |
| 610 | ✗ | scan->fs_accum[v2] = acc16[v2]; | |
| 611 | } | ||
| 612 | else | ||
| 613 | { | ||
| 614 | 27321 | scan->fs_accum_hacc(codes, scan->fs_lut, scan->fs_accum, dim); | |
| 615 | } | ||
| 616 | |||
| 617 | /* Prefetch next group's codes during prune phase */ | ||
| 618 |
2/2✓ Branch 0 taken 19838 times.
✓ Branch 1 taken 7483 times.
|
27321 | if (g + 1 < ngroups) |
| 619 | { | ||
| 620 | 3500 | uint8_t *next_codes = | |
| 621 | 19838 | prism_fastscan_group_codes(content, g + 1, dim); | |
| 622 | 19838 | uint32_t code_bytes = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 623 |
2/2✓ Branch 0 taken 30576 times.
✓ Branch 1 taken 19838 times.
|
50414 | for (uint32_t p = 0; p < code_bytes; p += 64) |
| 624 | 30576 | __builtin_prefetch(next_codes + p, 0, 1); | |
| 625 | } | ||
| 626 | |||
| 627 | /* Vectorized distance + prune: compute 16 distances at a | ||
| 628 | * time, mask survivors, iterate only the few that pass. */ | ||
| 629 | 27321 | scan->entries_scanned += g_count; | |
| 630 | |||
| 631 | #ifdef VS_SIMD_FULL | ||
| 632 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 633 |
4/4✓ Branch 0 taken 22494 times.
✓ Branch 1 taken 4827 times.
✓ Branch 2 taken 3853 times.
✓ Branch 3 taken 17298 times.
|
27321 | if (vs_likely(g_count == VS_FASTSCAN_GROUP) && scan->fs_has_avx512) |
| 634 | { | ||
| 635 | ✗ | fastscan_prune_group_avx512( | |
| 636 | scan, | ||
| 637 | topk, | ||
| 638 | tids, | ||
| 639 | f_add, | ||
| 640 | f_rescale, | ||
| 641 | f_error, | ||
| 642 | lut_scale, | ||
| 643 | lut_bias, | ||
| 644 | sum_t, | ||
| 645 | inv_sqrt_d, | ||
| 646 | g_add, | ||
| 647 | g_error, | ||
| 648 | &threshold); | ||
| 649 | } | ||
| 650 | else | ||
| 651 | #endif | ||
| 652 | #endif | ||
| 653 | { | ||
| 654 | /* Scalar fallback for partial groups */ | ||
| 655 |
2/2✓ Branch 0 taken 767173 times.
✓ Branch 1 taken 27321 times.
|
794494 | for (uint32_t v = 0; v < g_count; v++) |
| 656 | { | ||
| 657 | 767173 | float binary_ip = (float)scan->fs_accum[v] * lut_scale + | |
| 658 | lut_bias; | ||
| 659 | 767173 | float final_dot = (2.0f * binary_ip - sum_t) * inv_sqrt_d; | |
| 660 | 767173 | Distance est = f_add[v] + g_add - | |
| 661 | 767173 | 2.0f * f_rescale[v] * final_dot; | |
| 662 | 767173 | Distance err = f_error[v] * g_error; | |
| 663 | 767173 | Distance lb = est - err; | |
| 664 | |||
| 665 |
2/2✓ Branch 0 taken 649746 times.
✓ Branch 1 taken 117427 times.
|
767173 | if (lb >= threshold) |
| 666 | { | ||
| 667 | 649746 | scan->entries_pruned++; | |
| 668 | 649746 | continue; | |
| 669 | } | ||
| 670 | |||
| 671 | 117427 | uint64_t id = prism_posting_encode_tid(&tids[v]); | |
| 672 | 117427 | vs_topk_insert(topk, est, err, id); | |
| 673 |
2/2✓ Branch 0 taken 16502 times.
✓ Branch 1 taken 6655 times.
|
133929 | threshold = vs_topk_threshold(topk); |
| 674 | } | ||
| 675 | } | ||
| 676 | } | ||
| 677 | } | ||
| 678 | |||
| 679 | void | ||
| 680 | 9142 | prism_posting_scan_cluster_fastscan(PrismPostingScan *scan, VsTopK *topk) | |
| 681 | { | ||
| 682 | /* Process pages until chain is exhausted */ | ||
| 683 | 16472 | for (;;) | |
| 684 | 1936 | { | |
| 685 |
2/2✓ Branch 0 taken 10242 times.
✓ Branch 1 taken 9142 times.
|
19384 | if (scan->cur_page == NULL) |
| 686 | { | ||
| 687 |
2/2✓ Branch 1 taken 1233 times.
✓ Branch 2 taken 9009 times.
|
10242 | if (!advance_page(scan)) |
| 688 | 974 | break; | |
| 689 | } | ||
| 690 | |||
| 691 | 10363 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 692 | |||
| 693 |
2/3✗ Branch 0 not taken.
✓ Branch 1 taken 8427 times.
✓ Branch 2 taken 1936 times.
|
10363 | if (prism_posting_page_all_dead(opaque)) |
| 694 | { | ||
| 695 | /* All entries dead — skip scoring; the chain-follow advances. | ||
| 696 | * A retired chain is not that: it carries TOMBSTONED too once | ||
| 697 | * reclaimed, but its entries were rewritten elsewhere, and a | ||
| 698 | * scan that followed a stale leaf pointer into it must still | ||
| 699 | * see them. */ | ||
| 700 | 2 | scan->pages_skipped++; | |
| 701 | } | ||
| 702 |
2/2✓ Branch 0 taken 10240 times.
✓ Branch 1 taken 121 times.
|
10361 | else if (opaque->flags & PRISM_POSTING_PAGE_FASTSCAN) |
| 703 | 10240 | scan_fastscan_page(scan, topk); | |
| 704 | else | ||
| 705 | 121 | prism_posting_scan_cluster(scan, topk); | |
| 706 | |||
| 707 | /* If AoS fallback consumed the entire chain, we're done */ | ||
| 708 |
4/4✓ Branch 0 taken 121 times.
✓ Branch 1 taken 10242 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 119 times.
|
10363 | if (scan->cur_page == NULL && scan->cur_blkno == InvalidBlockNumber) |
| 709 | 2 | break; | |
| 710 | |||
| 711 | /* Follow chain: read next_blkno, then release current */ | ||
| 712 |
2/2✓ Branch 0 taken 1936 times.
✓ Branch 1 taken 8306 times.
|
10242 | if (scan->cur_page != NULL) |
| 713 | { | ||
| 714 | 10242 | BlockNumber prev_blkno = scan->cur_blkno; | |
| 715 | 10242 | scan->cur_blkno = opaque->next_blkno; | |
| 716 | |||
| 717 |
3/4✓ Branch 0 taken 10242 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8312 times.
✓ Branch 3 taken 1930 times.
|
10242 | if (scan->storage != NULL && scan->page_base == NULL) |
| 718 | 8312 | vs_storage_release_page(scan->storage, prev_blkno); | |
| 719 | |||
| 720 | 10242 | scan->cur_page = NULL; | |
| 721 | 10242 | scan->cur_content = NULL; | |
| 722 | } | ||
| 723 | } | ||
| 724 | 9142 | } | |
| 725 |