| 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 | 18112 | 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 | 18112 | memset(scan, 0, sizeof(*scan)); | |
| 49 | 18112 | scan->storage = storage; | |
| 50 | 18112 | scan->page_base = page_base; | |
| 51 | 18112 | scan->params = params; | |
| 52 | 18112 | scan->dim = dim; | |
| 53 | 18112 | scan->packed_bytes = VS_RABITQ_BYTES(dim); | |
| 54 | 18112 | 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 | 18112 | scan->max_entries_cap = max_entries_per_page; | |
| 61 | 18112 | uint32_t padded = (max_entries_per_page + 3) & ~3u; | |
| 62 | 18112 | scan->page_distances = vs_alloc(padded * sizeof(Distance)); | |
| 63 | 18112 | scan->page_scratch = vs_alloc(padded * sizeof(float)); | |
| 64 | 18112 | } | |
| 65 | |||
| 66 | void | ||
| 67 | 18110 | prism_posting_scan_cleanup(PrismPostingScan *scan) | |
| 68 | { | ||
| 69 | /* Release pinned page if still held */ | ||
| 70 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 18110 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
18110 | 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 | 18110 | vs_free(scan->page_distances); | |
| 78 | 18110 | vs_free(scan->page_scratch); | |
| 79 | 18110 | scan->page_distances = NULL; | |
| 80 | 18110 | scan->page_scratch = NULL; | |
| 81 | |||
| 82 |
2/2✓ Branch 0 taken 245 times.
✓ Branch 1 taken 17865 times.
|
18110 | if (scan->fs_lut != NULL) |
| 83 | { | ||
| 84 | 245 | vs_free(scan->fs_lut); | |
| 85 | 245 | scan->fs_lut = NULL; | |
| 86 | } | ||
| 87 |
2/2✓ Branch 0 taken 245 times.
✓ Branch 1 taken 17865 times.
|
18110 | if (scan->fs_accum != NULL) |
| 88 | { | ||
| 89 | 245 | vs_free(scan->fs_accum); | |
| 90 | 245 | scan->fs_accum = NULL; | |
| 91 | } | ||
| 92 | 18110 | } | |
| 93 | |||
| 94 | void | ||
| 95 | 247 | prism_posting_scan_enable_fastscan(PrismPostingScan *scan, int lut_bits) | |
| 96 | { | ||
| 97 | 247 | scan->fs_lut_bits = lut_bits; | |
| 98 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 247 times.
|
247 | if (lut_bits == 8) |
| 99 | ✗ | scan->fs_lut = vs_alloc0(VS_FASTSCAN_LUT_BYTES(scan->dim)); | |
| 100 | else | ||
| 101 | 247 | scan->fs_lut = vs_alloc0(VS_FASTSCAN_LUT_HACC_BYTES(scan->dim)); | |
| 102 | 247 | scan->fs_accum = vs_alloc(VS_FASTSCAN_GROUP * sizeof(int32_t)); | |
| 103 | |||
| 104 | /* Resolve dispatch once per scan (see posting_scan.h). */ | ||
| 105 | 247 | scan->fs_accum_hacc = vs_fastscan_get_accumulate_hacc(); | |
| 106 | #if defined(VS_SIMD_FULL) && (defined(__x86_64__) || defined(_M_X64)) | ||
| 107 | 247 | scan->fs_has_avx512 = vs_has_all_simd(VS_SIMD_AVX512_BW); | |
| 108 | #else | ||
| 109 | scan->fs_has_avx512 = false; | ||
| 110 | #endif | ||
| 111 | 247 | scan->fs_max_groups_first = prism_fastscan_max_groups(scan->dim, true); | |
| 112 | 247 | scan->fs_max_groups_over = prism_fastscan_max_groups(scan->dim, false); | |
| 113 | 247 | } | |
| 114 | |||
| 115 | static bool advance_page(PrismPostingScan *scan); | ||
| 116 | |||
| 117 | /* ---------------------------------------------------------------- | ||
| 118 | * Per-cluster begin / end | ||
| 119 | * ---------------------------------------------------------------- */ | ||
| 120 | |||
| 121 | void | ||
| 122 | 44542 | prism_posting_scan_begin_cluster( | |
| 123 | PrismPostingScan *scan, | ||
| 124 | RaBitQQueryState *qstate, | ||
| 125 | BlockNumber posting_head) | ||
| 126 | { | ||
| 127 | 44542 | scan->qstate = qstate; | |
| 128 | 44542 | scan->cur_blkno = posting_head; | |
| 129 | 44542 | scan->cur_page = NULL; | |
| 130 | 44542 | scan->cur_content = NULL; | |
| 131 | 44542 | scan->cur_max_entries = 0; | |
| 132 | 44542 | scan->cur_count = 0; | |
| 133 | 44542 | scan->fs_lut_valid = false; | |
| 134 | 44542 | scan->pages_read = 0; | |
| 135 | 44542 | scan->pages_skipped = 0; | |
| 136 | 44542 | scan->entries_scanned = 0; | |
| 137 | 44542 | 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 44542 times.
✗ Branch 1 not taken.
|
44542 | if (posting_head != InvalidBlockNumber) |
| 143 | 44542 | advance_page(scan); | |
| 144 | 44542 | } | |
| 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 | 44500 | prism_posting_scan_pt_centroid(const PrismPostingScan *scan) | |
| 152 | { | ||
| 153 |
2/2✓ Branch 0 taken 10502 times.
✓ Branch 1 taken 33998 times.
|
44500 | if (scan->cur_page == NULL) |
| 154 | ✗ | return NULL; | |
| 155 | 44500 | 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 | 44762 | prism_posting_scan_end_cluster(PrismPostingScan *scan) | |
| 185 | { | ||
| 186 | /* Release current page if held via storage vtable */ | ||
| 187 |
3/4✓ Branch 0 taken 9016 times.
✓ Branch 1 taken 35746 times.
✓ Branch 2 taken 9016 times.
✗ Branch 3 not taken.
|
44762 | if (scan->cur_page != NULL && scan->storage != NULL && |
| 188 |
2/2✓ Branch 0 taken 1016 times.
✓ Branch 1 taken 8000 times.
|
9016 | scan->page_base == NULL) |
| 189 | { | ||
| 190 | 1016 | vs_storage_release_page(scan->storage, scan->cur_blkno); | |
| 191 | } | ||
| 192 | |||
| 193 | 44762 | scan->cur_page = NULL; | |
| 194 | 44762 | scan->cur_content = NULL; | |
| 195 | 44762 | scan->qstate = NULL; | |
| 196 | 44762 | } | |
| 197 | |||
| 198 | /* ---------------------------------------------------------------- | ||
| 199 | * Advance to next page in chain | ||
| 200 | * ---------------------------------------------------------------- */ | ||
| 201 | |||
| 202 | static bool | ||
| 203 | 83291 | advance_page(PrismPostingScan *scan) | |
| 204 | { | ||
| 205 | /* Follow chain from current page */ | ||
| 206 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 83291 times.
|
83291 | 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 35746 times.
✓ Branch 1 taken 47545 times.
|
83291 | if (scan->cur_blkno == InvalidBlockNumber) |
| 222 | 26260 | return false; | |
| 223 | |||
| 224 | /* Read next page — inline for ArrayPageStorage. A scan over a | ||
| 225 | * non-paged posting format carries neither a page base nor a storage | ||
| 226 | * (see the flat branch of the standalone query context); it reaches | ||
| 227 | * its single page through begin_flat and has no chain to follow. */ | ||
| 228 |
2/2✓ Branch 0 taken 36508 times.
✓ Branch 1 taken 11037 times.
|
47545 | if (scan->page_base != NULL) |
| 229 | 36508 | scan->cur_page = scan->page_base + (size_t)scan->cur_blkno * BLCKSZ; | |
| 230 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 10993 times.
|
11037 | else if (scan->storage != NULL) |
| 231 | 11037 | scan->cur_page = vs_storage_read_page(scan->storage, scan->cur_blkno); | |
| 232 | else | ||
| 233 | ✗ | return false; | |
| 234 | |||
| 235 | 47545 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 236 | |||
| 237 | /* Validate page identity — catch corrupted chain pointers early */ | ||
| 238 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 47545 times.
|
47545 | if (opaque->page_id != PRISM_POSTING_PAGE_ID) |
| 239 | { | ||
| 240 | ✗ | vs_warn(VS_EXTENSION_NAME ": posting scan hit non-posting page " | |
| 241 | "(blkno=%u, page_id=0x%04X)", | ||
| 242 | scan->cur_blkno, | ||
| 243 | opaque->page_id); | ||
| 244 | /* Release the pin taken above before bailing (storage-backed only; | ||
| 245 | * in page_base mode the page is a borrowed pointer, not a pin). */ | ||
| 246 | ✗ | if (scan->storage != NULL && scan->page_base == NULL) | |
| 247 | ✗ | vs_storage_release_page(scan->storage, scan->cur_blkno); | |
| 248 | ✗ | scan->cur_page = NULL; | |
| 249 | ✗ | scan->cur_content = NULL; | |
| 250 | ✗ | return false; | |
| 251 | } | ||
| 252 | |||
| 253 | 58538 | scan->cur_content = | |
| 254 | 47545 | (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 255 | 44542 | ? prism_posting_content_first(scan->cur_page, scan->dim) | |
| 256 |
2/2✓ Branch 0 taken 44542 times.
✓ Branch 1 taken 3003 times.
|
47545 | : prism_posting_content(scan->cur_page); |
| 257 | 47545 | scan->cur_max_entries = opaque->max_entries; | |
| 258 | /* Reject an entry_count past what a valid page of this format can | ||
| 259 | * hold. AoS pages drive the page_distances/page_scratch batch (sized | ||
| 260 | * to max_entries_cap), so an over-count there is an out-of-bounds | ||
| 261 | * write; fastscan pages are scored per 32-vector group and hold more | ||
| 262 | * entries per page, so bound them by the fastscan capacity instead of | ||
| 263 | * the smaller AoS one. */ | ||
| 264 | 10993 | uint32_t page_cap; | |
| 265 |
2/2✓ Branch 0 taken 11276 times.
✓ Branch 1 taken 36269 times.
|
47545 | if (opaque->flags & PRISM_POSTING_PAGE_FASTSCAN) |
| 266 | 11276 | page_cap = (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 267 | 10180 | ? prism_fastscan_max_entries_first(scan->dim) | |
| 268 |
2/2✓ Branch 0 taken 10180 times.
✓ Branch 1 taken 1096 times.
|
11276 | : prism_fastscan_max_entries(scan->dim); |
| 269 | else | ||
| 270 | 36269 | page_cap = scan->max_entries_cap; | |
| 271 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 47545 times.
|
47545 | if (opaque->entry_count > page_cap) |
| 272 | ✗ | vs_error( | |
| 273 | VS_EXTENSION_NAME | ||
| 274 | ": posting page %u has an invalid entry count (%u > %u); " | ||
| 275 | "the index may be corrupted -- REINDEX it", | ||
| 276 | scan->cur_blkno, | ||
| 277 | (unsigned)opaque->entry_count, | ||
| 278 | page_cap); | ||
| 279 | 47545 | scan->cur_count = opaque->entry_count; | |
| 280 | 47545 | scan->pages_read++; | |
| 281 | 47545 | return true; | |
| 282 | } | ||
| 283 | |||
| 284 | /* ---------------------------------------------------------------- | ||
| 285 | * Score + prune a full cluster into topk | ||
| 286 | * | ||
| 287 | * Inserts survivors with approximate (estimated) distances and | ||
| 288 | * error bounds. No exact reranking — the caller does that after | ||
| 289 | * extracting candidates from topk. | ||
| 290 | * ---------------------------------------------------------------- */ | ||
| 291 | |||
| 292 | void | ||
| 293 | 26692 | prism_posting_scan_cluster(PrismPostingScan *scan, VsTopK *topk) | |
| 294 | { | ||
| 295 | 26692 | Dimension dim = scan->dim; | |
| 296 | 26692 | uint32_t entry_size = PRISM_POSTING_ENTRY_SIZE(dim); | |
| 297 | |||
| 298 | 26692 | float g_add = scan->qstate->g_add; | |
| 299 | 26692 | float sum_t = scan->qstate->sum_transformed; | |
| 300 | 26692 | float inv_sqrt_d = scan->qstate->inv_sqrt_d; | |
| 301 | 26692 | float g_error = scan->qstate->g_error; | |
| 302 | |||
| 303 | 26692 | Distance *distances = scan->page_distances; | |
| 304 | 26692 | float *scratch = scan->page_scratch; | |
| 305 | |||
| 306 | /* Process pages until chain is exhausted. | ||
| 307 | * For flat mode: single iteration (one page, no chain). */ | ||
| 308 | 4734 | for (;;) | |
| 309 | 28485 | { | |
| 310 | /* Advance to next page */ | ||
| 311 |
2/2✓ Branch 0 taken 28485 times.
✓ Branch 1 taken 26692 times.
|
55177 | if (scan->cur_page == NULL) |
| 312 | { | ||
| 313 |
2/2✓ Branch 1 taken 25519 times.
✓ Branch 2 taken 2966 times.
|
28485 | if (!advance_page(scan)) |
| 314 | 25276 | break; | |
| 315 | } | ||
| 316 | |||
| 317 | /* A page tombstoned by VACUUM (all entries dead) is skipped — no | ||
| 318 | * scoring kernel; the chain-follow below still advances past it. A | ||
| 319 | * retired chain also carries TOMBSTONED once reclaimed, but its | ||
| 320 | * entries moved rather than died, so it is still scored. */ | ||
| 321 |
3/3✓ Branch 1 taken 1651 times.
✓ Branch 2 taken 26828 times.
✓ Branch 3 taken 6 times.
|
28485 | if (!prism_posting_page_all_dead(prism_posting_opaque(scan->cur_page))) |
| 322 | { | ||
| 323 | 28471 | char *content = scan->cur_content; | |
| 324 | 28471 | uint32_t count = scan->cur_count; | |
| 325 | |||
| 326 | /* --- Score: batch IP over all entries on this page --- | ||
| 327 | * | ||
| 328 | * AoS layout: entry i's bits live at content + i * entry_size | ||
| 329 | * + PRISM_POSTING_ENTRY_BITS_OFFSET. The SIMD kernel just needs | ||
| 330 | * the first entry's bits pointer and a stride of entry_size. */ | ||
| 331 | 28471 | const uint8_t *bits_base = prism_posting_first_bits(content); | |
| 332 | |||
| 333 | 28471 | uint32_t padded = (count + 3) & ~3u; | |
| 334 | 28471 | vs_rabitq_inner_product_multi( | |
| 335 | 28471 | scan->qstate->transformed, | |
| 336 | bits_base, | ||
| 337 | entry_size, | ||
| 338 | dim, | ||
| 339 | padded, | ||
| 340 | scratch); | ||
| 341 | |||
| 342 | /* Convert raw IPs to distances, reading f_add/f_rescale per | ||
| 343 | * entry via the strided AoS accessor. */ | ||
| 344 |
3/3✓ Branch 0 taken 2751898 times.
✓ Branch 1 taken 84324 times.
✓ Branch 2 taken 1651 times.
|
2837873 | for (uint32_t i = 0; i < count; i++) |
| 345 | { | ||
| 346 | 57504 | PrismPostingEntryHeader *e = | |
| 347 | 2809402 | prism_posting_entry_at(content, i, dim); | |
| 348 | 2809402 | float final_dot = (2.0f * scratch[i] - sum_t) * inv_sqrt_d; | |
| 349 | 2809402 | distances[i] = e->f_add + g_add - | |
| 350 | 2809402 | 2.0f * e->f_rescale * final_dot; | |
| 351 | } | ||
| 352 | |||
| 353 | /* --- Prune + insert approximate distances --- */ | ||
| 354 |
2/2✓ Branch 0 taken 1577 times.
✓ Branch 1 taken 74 times.
|
28471 | Distance threshold = vs_topk_threshold(topk); |
| 355 | |||
| 356 |
2/2✓ Branch 0 taken 2809402 times.
✓ Branch 1 taken 28471 times.
|
2837873 | for (uint32_t i = 0; i < count; i++) |
| 357 | { | ||
| 358 | 57504 | PrismPostingEntryHeader *e = | |
| 359 |
2/2✓ Branch 0 taken 2412 times.
✓ Branch 1 taken 55092 times.
|
2809402 | prism_posting_entry_at(content, i, dim); |
| 360 | 2809402 | scan->entries_scanned++; | |
| 361 | |||
| 362 |
2/2✓ Branch 0 taken 2434 times.
✓ Branch 1 taken 2806968 times.
|
2809402 | if (e->meta.flags & PRISM_POSTING_FLAG_DELETED) |
| 363 | { | ||
| 364 | 2434 | scan->entries_pruned++; | |
| 365 | 2434 | continue; | |
| 366 | } | ||
| 367 | |||
| 368 | 2806968 | Distance est = distances[i]; | |
| 369 | 2806968 | Distance err = e->f_error * g_error; | |
| 370 | 2806968 | Distance lb = est - err; | |
| 371 | |||
| 372 |
2/2✓ Branch 0 taken 2330793 times.
✓ Branch 1 taken 476175 times.
|
2806968 | if (lb >= threshold) |
| 373 | { | ||
| 374 | 2330793 | scan->entries_pruned++; | |
| 375 | 2330793 | continue; | |
| 376 | } | ||
| 377 | |||
| 378 | 476175 | uint64_t id = prism_posting_encode_tid(&e->meta.tid); | |
| 379 | 476175 | vs_topk_insert(topk, est, err, id); | |
| 380 |
2/2✓ Branch 0 taken 2992 times.
✓ Branch 1 taken 1397 times.
|
532282 | threshold = vs_topk_threshold(topk); |
| 381 | } | ||
| 382 | } /* end: page not tombstoned */ | ||
| 383 | else | ||
| 384 | 14 | scan->pages_skipped++; | |
| 385 | |||
| 386 | /* Follow chain: read next_blkno, then release current */ | ||
| 387 | 28485 | BlockNumber prev_blkno = scan->cur_blkno; | |
| 388 | |||
| 389 |
2/2✓ Branch 0 taken 28265 times.
✓ Branch 1 taken 220 times.
|
28485 | if (scan->storage != NULL) |
| 390 | 28265 | scan->cur_blkno = prism_posting_opaque(scan->cur_page)->next_blkno; | |
| 391 | else | ||
| 392 | 220 | scan->cur_blkno = InvalidBlockNumber; | |
| 393 | |||
| 394 |
4/4✓ Branch 0 taken 28265 times.
✓ Branch 1 taken 220 times.
✓ Branch 2 taken 1687 times.
✓ Branch 3 taken 26578 times.
|
28485 | if (scan->storage != NULL && scan->page_base == NULL) |
| 395 | 1687 | vs_storage_release_page(scan->storage, prev_blkno); | |
| 396 | |||
| 397 | 28485 | scan->cur_page = NULL; | |
| 398 | 28485 | scan->cur_content = NULL; | |
| 399 | } | ||
| 400 | 26692 | } | |
| 401 | |||
| 402 | /* ---------------------------------------------------------------- | ||
| 403 | * Fastscan: vectorized distance + prune for 16 entries | ||
| 404 | * ---------------------------------------------------------------- */ | ||
| 405 | |||
| 406 | #ifdef VS_SIMD_FULL | ||
| 407 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 408 | |||
| 409 | /* | ||
| 410 | * Compute distances and lower bounds for 16 entries using AVX-512. | ||
| 411 | * Returns a bitmask of survivors (lb < threshold). | ||
| 412 | */ | ||
| 413 | VS_TARGET_AVX512 static inline __mmask16 | ||
| 414 | 42256 | fastscan_prune_16( | |
| 415 | const int32_t *accum, | ||
| 416 | const float *f_add, | ||
| 417 | const float *f_rescale, | ||
| 418 | const float *f_error, | ||
| 419 | float *est_out, | ||
| 420 | float *err_out, | ||
| 421 | float lut_scale, | ||
| 422 | float lut_bias, | ||
| 423 | float sum_t, | ||
| 424 | float inv_sqrt_d, | ||
| 425 | float g_add, | ||
| 426 | float g_error, | ||
| 427 | float threshold, | ||
| 428 | __mmask16 live) | ||
| 429 | { | ||
| 430 | /* Convert 16 int32 accumulators to float */ | ||
| 431 | 42256 | __m512i acc32 = _mm512_loadu_si512((const __m512i *)accum); | |
| 432 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 acc_f = _mm512_cvtepi32_ps(acc32); |
| 433 | |||
| 434 | /* binary_ip = acc * scale + bias */ | ||
| 435 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 scale_v = _mm512_set1_ps(lut_scale); |
| 436 | 42256 | __m512 bias_v = _mm512_set1_ps(lut_bias); | |
| 437 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 ip = _mm512_fmadd_ps(acc_f, scale_v, bias_v); |
| 438 | |||
| 439 | /* final_dot = (2*ip - sum_t) * inv_sqrt_d */ | ||
| 440 | 42256 | __m512 two_v = _mm512_set1_ps(2.0f); | |
| 441 | 42256 | __m512 sum_t_v = _mm512_set1_ps(sum_t); | |
| 442 | 42256 | __m512 inv_sqrt_v = _mm512_set1_ps(inv_sqrt_d); | |
| 443 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
7636 | __m512 final_dot = |
| 444 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
76876 | _mm512_mul_ps(_mm512_fmsub_ps(two_v, ip, sum_t_v), inv_sqrt_v); |
| 445 | |||
| 446 | /* est = f_add + g_add - 2 * f_rescale * final_dot */ | ||
| 447 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 fa = _mm512_loadu_ps(f_add); |
| 448 | 42256 | __m512 fr = _mm512_loadu_ps(f_rescale); | |
| 449 | 42256 | __m512 gadd = _mm512_set1_ps(g_add); | |
| 450 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 est_v = _mm512_add_ps(fa, gadd); |
| 451 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
76876 | est_v = _mm512_fnmadd_ps(_mm512_mul_ps(two_v, fr), final_dot, est_v); |
| 452 | |||
| 453 | /* err = f_error * g_error */ | ||
| 454 | 42256 | __m512 fe = _mm512_loadu_ps(f_error); | |
| 455 | 42256 | __m512 gerr_v = _mm512_set1_ps(g_error); | |
| 456 | 42256 | __m512 err_v = _mm512_mul_ps(fe, gerr_v); | |
| 457 | |||
| 458 | /* lb = est - err; survivors = lb < threshold, restricted to live lanes -- | ||
| 459 | * the masked compare clears a dead lane's result bit regardless of score. | ||
| 460 | */ | ||
| 461 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __m512 lb_v = _mm512_sub_ps(est_v, err_v); |
| 462 | 42256 | __m512 thr_v = _mm512_set1_ps(threshold); | |
| 463 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | __mmask16 surv = _mm512_mask_cmp_ps_mask(live, lb_v, thr_v, _CMP_LT_OS); |
| 464 | |||
| 465 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
7636 | _mm512_storeu_ps(est_out, est_v); |
| 466 | 7636 | _mm512_storeu_ps(err_out, err_v); | |
| 467 | |||
| 468 | 42256 | return surv; | |
| 469 | } | ||
| 470 | |||
| 471 | /* | ||
| 472 | * AVX-512 vectorized distance + prune for one 32-vector group. | ||
| 473 | * Must be in a target-attributed function so the compiler only | ||
| 474 | * emits AVX-512 instructions here, not in the caller. | ||
| 475 | */ | ||
| 476 | VS_TARGET_AVX512 static void | ||
| 477 | 21128 | fastscan_prune_group_avx512( | |
| 478 | PrismPostingScan *scan, | ||
| 479 | VsTopK *topk, | ||
| 480 | ItemPointerData *tids, | ||
| 481 | const float *f_add, | ||
| 482 | const float *f_rescale, | ||
| 483 | const float *f_error, | ||
| 484 | float lut_scale, | ||
| 485 | float lut_bias, | ||
| 486 | float sum_t, | ||
| 487 | float inv_sqrt_d, | ||
| 488 | float g_add, | ||
| 489 | float g_error, | ||
| 490 | uint32_t tombstone_mask, | ||
| 491 | Distance *threshold_p) | ||
| 492 | { | ||
| 493 | 21128 | Distance threshold = *threshold_p; | |
| 494 | 3818 | float est_buf[VS_FASTSCAN_GROUP]; | |
| 495 | 3818 | float err_buf[VS_FASTSCAN_GROUP]; | |
| 496 | 21128 | uint32_t live_mask = ~tombstone_mask; | |
| 497 | |||
| 498 |
2/2✓ Branch 0 taken 42256 times.
✓ Branch 1 taken 21128 times.
|
63384 | for (uint32_t h = 0; h < 2; h++) |
| 499 | { | ||
| 500 | 42256 | uint32_t off = h * 16; | |
| 501 | 42256 | __mmask16 live = (__mmask16)((live_mask >> off) & 0xFFFFu); | |
| 502 | 42256 | __mmask16 surv = fastscan_prune_16( | |
| 503 | 42256 | scan->fs_accum + off, | |
| 504 | 34620 | f_add + off, | |
| 505 | 34620 | f_rescale + off, | |
| 506 | 34620 | f_error + off, | |
| 507 | 34620 | est_buf + off, | |
| 508 |
2/2✓ Branch 0 taken 6439 times.
✓ Branch 1 taken 1197 times.
|
42256 | err_buf + off, |
| 509 | lut_scale, | ||
| 510 | lut_bias, | ||
| 511 | sum_t, | ||
| 512 | inv_sqrt_d, | ||
| 513 | g_add, | ||
| 514 | g_error, | ||
| 515 | threshold, | ||
| 516 | live); | ||
| 517 | |||
| 518 |
2/2✓ Branch 0 taken 24241 times.
✓ Branch 1 taken 18015 times.
|
42256 | if (surv == 0) |
| 519 | { | ||
| 520 | 24241 | scan->entries_pruned += 16; | |
| 521 | 24241 | continue; | |
| 522 | } | ||
| 523 | |||
| 524 | 18015 | scan->entries_pruned += 16 - _mm_popcnt_u32(surv); | |
| 525 | |||
| 526 |
2/2✓ Branch 0 taken 93920 times.
✓ Branch 1 taken 16818 times.
|
111935 | while (surv != 0) |
| 527 | { | ||
| 528 | 105509 | uint32_t v = off + __builtin_ctz(surv); | |
| 529 | 105509 | surv &= surv - 1; | |
| 530 | |||
| 531 | 105509 | uint64_t id = prism_posting_encode_tid(&tids[v]); | |
| 532 | 105509 | vs_topk_insert(topk, est_buf[v], err_buf[v], id); | |
| 533 |
4/4✓ Branch 0 taken 10143 times.
✓ Branch 1 taken 1446 times.
✓ Branch 2 taken 11589 times.
✓ Branch 3 taken 1197 times.
|
118295 | threshold = vs_topk_threshold(topk); |
| 534 | } | ||
| 535 | } | ||
| 536 | |||
| 537 | 21128 | *threshold_p = threshold; | |
| 538 | 21128 | } | |
| 539 | |||
| 540 | #endif | ||
| 541 | #endif | ||
| 542 | |||
| 543 | /* ---------------------------------------------------------------- | ||
| 544 | * Fastscan variant: VPSHUFB kernel for fastscan-format pages | ||
| 545 | * | ||
| 546 | * Processes 32-vector groups via table-lookup accumulation. | ||
| 547 | * Falls back to the AoS kernel for non-fastscan pages (mixed | ||
| 548 | * chains). | ||
| 549 | * ---------------------------------------------------------------- */ | ||
| 550 | |||
| 551 | static void | ||
| 552 | 10262 | scan_fastscan_page(PrismPostingScan *scan, VsTopK *topk) | |
| 553 | { | ||
| 554 | 10262 | Dimension dim = scan->dim; | |
| 555 | 10262 | uint32_t count = scan->cur_count; | |
| 556 | 10262 | char *content = scan->cur_content; | |
| 557 | |||
| 558 | 10262 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 559 | |||
| 560 | /* Fall back to AoS kernel for non-fastscan pages */ | ||
| 561 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10262 times.
|
10262 | if (!(opaque->flags & PRISM_POSTING_PAGE_FASTSCAN)) |
| 562 | ✗ | return; /* caller handles AoS via the standard path */ | |
| 563 | |||
| 564 | /* Build LUT from transformed query (once per cluster) */ | ||
| 565 |
2/2✓ Branch 0 taken 9166 times.
✓ Branch 1 taken 1096 times.
|
10262 | if (!scan->fs_lut_valid) |
| 566 | { | ||
| 567 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 9166 times.
|
9166 | if (scan->fs_lut_bits == 8) |
| 568 | ✗ | vs_fastscan_build_lut( | |
| 569 | ✗ | scan->qstate->transformed, | |
| 570 | dim, | ||
| 571 | scan->fs_lut, | ||
| 572 | &scan->fs_lut_scale, | ||
| 573 | &scan->fs_lut_bias); | ||
| 574 | else | ||
| 575 | 9166 | vs_fastscan_build_lut_hacc( | |
| 576 | 9166 | scan->qstate->transformed, | |
| 577 | dim, | ||
| 578 | scan->fs_lut, | ||
| 579 | &scan->fs_lut_scale, | ||
| 580 | &scan->fs_lut_bias); | ||
| 581 | 9166 | scan->fs_lut_valid = true; | |
| 582 | } | ||
| 583 | 10262 | float lut_scale = scan->fs_lut_scale; | |
| 584 | 10262 | float lut_bias = scan->fs_lut_bias; | |
| 585 | |||
| 586 | 10262 | float g_add = scan->qstate->g_add; | |
| 587 | 10262 | float sum_t = scan->qstate->sum_transformed; | |
| 588 | 10262 | float inv_sqrt_d = scan->qstate->inv_sqrt_d; | |
| 589 | 10262 | float g_error = scan->qstate->g_error; | |
| 590 | 20524 | uint32_t max_groups = (opaque->flags & PRISM_POSTING_PAGE_FIRST) | |
| 591 | ? scan->fs_max_groups_first | ||
| 592 |
2/2✓ Branch 0 taken 9166 times.
✓ Branch 1 taken 1096 times.
|
10262 | : scan->fs_max_groups_over; |
| 593 | 10262 | uint32_t ngroups = (count + VS_FASTSCAN_GROUP - 1) / VS_FASTSCAN_GROUP; | |
| 594 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1944 times.
|
10262 | if (ngroups > max_groups) |
| 595 | ✗ | ngroups = max_groups; | |
| 596 | |||
| 597 |
2/2✓ Branch 0 taken 7808 times.
✓ Branch 1 taken 510 times.
|
10262 | Distance threshold = vs_topk_threshold(topk); |
| 598 | |||
| 599 |
2/2✓ Branch 0 taken 27339 times.
✓ Branch 1 taken 10262 times.
|
37601 | for (uint32_t g = 0; g < ngroups; g++) |
| 600 | { | ||
| 601 | 27339 | uint32_t g_start = g * VS_FASTSCAN_GROUP; | |
| 602 | 27339 | uint32_t g_count = count - g_start; | |
| 603 |
2/2✓ Branch 0 taken 16350 times.
✓ Branch 1 taken 1944 times.
|
27339 | if (g_count > VS_FASTSCAN_GROUP) |
| 604 | 16350 | g_count = VS_FASTSCAN_GROUP; | |
| 605 | |||
| 606 | /* Access group data — compute base once */ | ||
| 607 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 9045 times.
|
27339 | char *gbase = prism_fastscan_group_base(content, g, dim); |
| 608 | |||
| 609 | 27339 | ItemPointerData *tids = (ItemPointerData *)gbase; | |
| 610 | 27339 | float *f_add = (float *)(gbase + | |
| 611 | VS_FASTSCAN_GROUP * sizeof(ItemPointerData)); | ||
| 612 | 27339 | float *f_rescale = f_add + VS_FASTSCAN_GROUP; | |
| 613 | 27339 | float *f_error = f_rescale + VS_FASTSCAN_GROUP; | |
| 614 | 27339 | uint8_t *codes = (uint8_t *)(f_error + VS_FASTSCAN_GROUP); | |
| 615 | 27339 | uint32_t tombstone_mask = *(uint32_t *)(codes + | |
| 616 | 27339 | VS_FASTSCAN_GROUP_BYTES(dim)); | |
| 617 | |||
| 618 | /* Run VPSHUFB accumulate kernel */ | ||
| 619 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 27339 times.
|
27339 | if (scan->fs_lut_bits == 8) |
| 620 | { | ||
| 621 | ✗ | uint16_t acc16[VS_FASTSCAN_GROUP]; | |
| 622 | ✗ | vs_fastscan_accumulate(codes, scan->fs_lut, acc16, dim); | |
| 623 | ✗ | for (uint32_t v2 = 0; v2 < VS_FASTSCAN_GROUP; v2++) | |
| 624 | ✗ | scan->fs_accum[v2] = acc16[v2]; | |
| 625 | } | ||
| 626 | else | ||
| 627 | { | ||
| 628 | 27339 | scan->fs_accum_hacc(codes, scan->fs_lut, scan->fs_accum, dim); | |
| 629 | } | ||
| 630 | |||
| 631 | /* Prefetch next group's codes during prune phase */ | ||
| 632 |
2/2✓ Branch 0 taken 19822 times.
✓ Branch 1 taken 7517 times.
|
27339 | if (g + 1 < ngroups) |
| 633 | { | ||
| 634 | 3472 | uint8_t *next_codes = | |
| 635 | 19822 | prism_fastscan_group_codes(content, g + 1, dim); | |
| 636 | 19822 | uint32_t code_bytes = VS_FASTSCAN_GROUP_BYTES(dim); | |
| 637 |
2/2✓ Branch 0 taken 30600 times.
✓ Branch 1 taken 19822 times.
|
50422 | for (uint32_t p = 0; p < code_bytes; p += 64) |
| 638 | 30600 | __builtin_prefetch(next_codes + p, 0, 1); | |
| 639 | } | ||
| 640 | |||
| 641 | /* Vectorized distance + prune: compute 16 distances at a | ||
| 642 | * time, mask survivors, iterate only the few that pass. */ | ||
| 643 | 27339 | scan->entries_scanned += g_count; | |
| 644 | |||
| 645 | #ifdef VS_SIMD_FULL | ||
| 646 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 647 |
4/4✓ Branch 0 taken 22537 times.
✓ Branch 1 taken 4802 times.
✓ Branch 2 taken 17310 times.
✓ Branch 3 taken 3818 times.
|
27339 | if (vs_likely(g_count == VS_FASTSCAN_GROUP) && scan->fs_has_avx512) |
| 648 | { | ||
| 649 | 21128 | fastscan_prune_group_avx512( | |
| 650 | scan, | ||
| 651 | topk, | ||
| 652 | tids, | ||
| 653 | f_add, | ||
| 654 | f_rescale, | ||
| 655 | f_error, | ||
| 656 | lut_scale, | ||
| 657 | lut_bias, | ||
| 658 | sum_t, | ||
| 659 | inv_sqrt_d, | ||
| 660 | g_add, | ||
| 661 | g_error, | ||
| 662 | tombstone_mask, | ||
| 663 | &threshold); | ||
| 664 | } | ||
| 665 | else | ||
| 666 | #endif | ||
| 667 | #endif | ||
| 668 | { | ||
| 669 | /* Scalar fallback for partial groups (and non-AVX512 CPUs, | ||
| 670 | * NEON included): no masked-compare primitive to fold the | ||
| 671 | * exclusion into, so a dead lane is skipped explicitly, the | ||
| 672 | * same as a DELETED AoS entry below. */ | ||
| 673 |
2/2✓ Branch 0 taken 90895 times.
✓ Branch 1 taken 6211 times.
|
97106 | for (uint32_t v = 0; v < g_count; v++) |
| 674 | { | ||
| 675 |
2/3✗ Branch 0 not taken.
✓ Branch 1 taken 67579 times.
✓ Branch 2 taken 23316 times.
|
90895 | if (prism_fastscan_lane_is_tombstoned(tombstone_mask, v)) |
| 676 | { | ||
| 677 | 4 | scan->entries_pruned++; | |
| 678 | 4 | continue; | |
| 679 | } | ||
| 680 | |||
| 681 | 90891 | float binary_ip = (float)scan->fs_accum[v] * lut_scale + | |
| 682 | lut_bias; | ||
| 683 | 90891 | float final_dot = (2.0f * binary_ip - sum_t) * inv_sqrt_d; | |
| 684 | 90891 | Distance est = f_add[v] + g_add - | |
| 685 | 90891 | 2.0f * f_rescale[v] * final_dot; | |
| 686 | 90891 | Distance err = f_error[v] * g_error; | |
| 687 | 90891 | Distance lb = est - err; | |
| 688 | |||
| 689 |
2/2✓ Branch 0 taken 77012 times.
✓ Branch 1 taken 13879 times.
|
90891 | if (lb >= threshold) |
| 690 | { | ||
| 691 | 77012 | scan->entries_pruned++; | |
| 692 | 77012 | continue; | |
| 693 | } | ||
| 694 | |||
| 695 | 13879 | uint64_t id = prism_posting_encode_tid(&tids[v]); | |
| 696 | 13879 | vs_topk_insert(topk, est, err, id); | |
| 697 |
2/2✓ Branch 0 taken 6330 times.
✓ Branch 1 taken 5209 times.
|
20209 | threshold = vs_topk_threshold(topk); |
| 698 | } | ||
| 699 | } | ||
| 700 | } | ||
| 701 | } | ||
| 702 | |||
| 703 | void | ||
| 704 | 9168 | prism_posting_scan_cluster_fastscan(PrismPostingScan *scan, VsTopK *topk) | |
| 705 | { | ||
| 706 | /* Process pages until chain is exhausted */ | ||
| 707 | 16500 | for (;;) | |
| 708 | 1946 | { | |
| 709 |
2/2✓ Branch 0 taken 10264 times.
✓ Branch 1 taken 9168 times.
|
19432 | if (scan->cur_page == NULL) |
| 710 | { | ||
| 711 |
2/2✓ Branch 1 taken 1232 times.
✓ Branch 2 taken 9032 times.
|
10264 | if (!advance_page(scan)) |
| 712 | 984 | break; | |
| 713 | } | ||
| 714 | |||
| 715 | 10378 | PrismPostingPageOpaque *opaque = prism_posting_opaque(scan->cur_page); | |
| 716 | |||
| 717 |
2/3✗ Branch 0 not taken.
✓ Branch 1 taken 8432 times.
✓ Branch 2 taken 1946 times.
|
10378 | if (prism_posting_page_all_dead(opaque)) |
| 718 | { | ||
| 719 | /* All entries dead — skip scoring; the chain-follow advances. | ||
| 720 | * A retired chain is not that: it carries TOMBSTONED too once | ||
| 721 | * reclaimed, but its entries were rewritten elsewhere, and a | ||
| 722 | * scan that followed a stale leaf pointer into it must still | ||
| 723 | * see them. */ | ||
| 724 | 2 | scan->pages_skipped++; | |
| 725 | } | ||
| 726 |
2/2✓ Branch 0 taken 10262 times.
✓ Branch 1 taken 114 times.
|
10376 | else if (opaque->flags & PRISM_POSTING_PAGE_FASTSCAN) |
| 727 | 10262 | scan_fastscan_page(scan, topk); | |
| 728 | else | ||
| 729 | 114 | prism_posting_scan_cluster(scan, topk); | |
| 730 | |||
| 731 | /* If AoS fallback consumed the entire chain, we're done */ | ||
| 732 |
4/4✓ Branch 0 taken 114 times.
✓ Branch 1 taken 10264 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 112 times.
|
10378 | if (scan->cur_page == NULL && scan->cur_blkno == InvalidBlockNumber) |
| 733 | 2 | break; | |
| 734 | |||
| 735 | /* Follow chain: read next_blkno, then release current */ | ||
| 736 |
2/2✓ Branch 0 taken 1946 times.
✓ Branch 1 taken 8318 times.
|
10264 | if (scan->cur_page != NULL) |
| 737 | { | ||
| 738 | 10264 | BlockNumber prev_blkno = scan->cur_blkno; | |
| 739 | 10264 | scan->cur_blkno = opaque->next_blkno; | |
| 740 | |||
| 741 |
3/4✓ Branch 0 taken 10264 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8334 times.
✓ Branch 3 taken 1930 times.
|
10264 | if (scan->storage != NULL && scan->page_base == NULL) |
| 742 | 8334 | vs_storage_release_page(scan->storage, prev_blkno); | |
| 743 | |||
| 744 | 10264 | scan->cur_page = NULL; | |
| 745 | 10264 | scan->cur_content = NULL; | |
| 746 | } | ||
| 747 | } | ||
| 748 | 9168 | } | |
| 749 |