| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * parallel_build_worker.c - Parallel index build, per-participant bodies | ||
| 6 | * | ||
| 7 | * The leader runs these as participant 0 alongside the workers: | ||
| 8 | * Sampling: cooperative heap scan, each participant fills its own | ||
| 9 | * sample slot. | ||
| 10 | * K-means: leader seeds the initial centroids, then iterative | ||
| 11 | * assignment + accumulation over the per-participant samples, | ||
| 12 | * barrier-synchronized with a leader reduce. | ||
| 13 | * Subtrees: batched largest-first schedule; each participant clusters | ||
| 14 | * its batch's child into a ring slot the leader consumes. | ||
| 15 | * Refine (when subsampled): route the full table page-backed and | ||
| 16 | * accumulate per-leaf means into the shared tiled accumulator. | ||
| 17 | * Posting: cooperative heap scan, page-backed routing + RaBitQ encode | ||
| 18 | * into the cluster-keyed shared sort. | ||
| 19 | */ | ||
| 20 | |||
| 21 | #ifdef VS_STANDALONE | ||
| 22 | #include "standalone/parallel_ctx.h" /* ParallelWorkerNumber */ | ||
| 23 | #include "standalone/pg_compat.h" | ||
| 24 | #else | ||
| 25 | #include <postgres.h> | ||
| 26 | |||
| 27 | #include <access/parallel.h> | ||
| 28 | #include <catalog/index.h> | ||
| 29 | #include <miscadmin.h> | ||
| 30 | #include <storage/barrier.h> | ||
| 31 | #include <storage/bufmgr.h> | ||
| 32 | #include <storage/shm_toc.h> | ||
| 33 | #include <utils/rel.h> | ||
| 34 | #include <utils/wait_event.h> | ||
| 35 | #endif | ||
| 36 | |||
| 37 | #include <inttypes.h> | ||
| 38 | #include <math.h> | ||
| 39 | |||
| 40 | #include "algo/hkmeans.h" | ||
| 41 | #include "algo/kmeans_internal.h" | ||
| 42 | #include "algo/vecops.h" | ||
| 43 | #include "core/log.h" | ||
| 44 | #include "core/memory.h" | ||
| 45 | #include "index/build_progress.h" | ||
| 46 | #include "index/index_build.h" | ||
| 47 | #include "index/parallel_build.h" | ||
| 48 | #include "index/posting_build.h" | ||
| 49 | #include "quant/rabitq.h" | ||
| 50 | |||
| 51 | /* ---------------------------------------------------------------- | ||
| 52 | * Phase 1: Sampling callback | ||
| 53 | * | ||
| 54 | * Collects vectors into a per-worker sample slot (stride-subsampled, | ||
| 55 | * normalized for cosine). Initial centroids are seeded later by the | ||
| 56 | * leader from the pooled samples. | ||
| 57 | * ---------------------------------------------------------------- */ | ||
| 58 | |||
| 59 | void | ||
| 60 | 286242 | prism_sample_cb(void *state, ItemPointerData tid, const float *vec) | |
| 61 | { | ||
| 62 | 286242 | SampleCbState *sc = (SampleCbState *)state; | |
| 63 | |||
| 64 | 109480 | (void)tid; | |
| 65 | |||
| 66 | 286242 | sc->seen++; | |
| 67 | /* The sampling pass walks the whole heap; surface that in the progress | ||
| 68 | * view in coarse batches (the leader publishes exact counts at phase | ||
| 69 | * boundaries). */ | ||
| 70 |
2/2✓ Branch 0 taken 210 times.
✓ Branch 1 taken 286032 times.
|
286242 | if ((sc->seen & 1023) == 0) |
| 71 | 210 | prism_build_progress_incr_tuples(1024); | |
| 72 | |||
| 73 | /* Stride-based subsampling */ | ||
| 74 |
2/2✓ Branch 0 taken 22927 times.
✓ Branch 1 taken 263315 times.
|
286242 | if (sc->stride_counter > 0) |
| 75 | { | ||
| 76 | 22927 | sc->stride_counter--; | |
| 77 | 22927 | return; | |
| 78 | } | ||
| 79 | 263315 | sc->stride_counter = sc->stride - 1; | |
| 80 | |||
| 81 |
2/2✓ Branch 0 taken 89649 times.
✓ Branch 1 taken 173666 times.
|
263315 | if (sc->count >= sc->max_samples) |
| 82 | 20000 | return; | |
| 83 | |||
| 84 | 226411 | Dimension dim = sc->dim; | |
| 85 | 226411 | float *dest = sc->samples + (size_t)sc->count * dim; | |
| 86 | |||
| 87 | /* Normalize for cosine. Zero-norm vectors are dropped from the | ||
| 88 | * sample: they carry no direction, so under cosine they sit at the | ||
| 89 | * origin -- equidistant from every real point -- and k-means drags | ||
| 90 | * centroids toward them, warping the encode references (and so the | ||
| 91 | * distance estimates) for every vector in the affected subtree. | ||
| 92 | * They stay in the posting lists like any other row; they just | ||
| 93 | * don't get a vote on the clustering. The check runs on the source | ||
| 94 | * so a dropped vector costs no copy, and the scale writes straight | ||
| 95 | * into the sample slot, normalizing and copying in one pass. */ | ||
| 96 |
2/2✓ Branch 0 taken 5493 times.
✓ Branch 1 taken 220918 times.
|
226411 | if (sc->metric == DISTANCE_COSINE) |
| 97 | { | ||
| 98 | 5493 | float norm_sq = vs_l2_norm_squared(vec, dim); | |
| 99 | |||
| 100 |
2/2✓ Branch 0 taken 3692 times.
✓ Branch 1 taken 1801 times.
|
5493 | if (norm_sq == 0.0f) |
| 101 | 80 | return; | |
| 102 | 5412 | vec32_scale(vec, 1.0f / sqrtf(norm_sq), dest, dim); | |
| 103 | } | ||
| 104 | else | ||
| 105 | 220918 | memcpy(dest, vec, dim * sizeof(float)); | |
| 106 | |||
| 107 | 226330 | sc->count++; | |
| 108 | } | ||
| 109 | |||
| 110 | /* ---------------------------------------------------------------- | ||
| 111 | * Phase 2: K-means assignment — thin wrappers around shared kernel | ||
| 112 | * ---------------------------------------------------------------- */ | ||
| 113 | |||
| 114 | void | ||
| 115 | 3389 | prism_km_assign_and_accumulate( | |
| 116 | const float *samples, | ||
| 117 | uint32_t nsamples, | ||
| 118 | const float *centroids, | ||
| 119 | const float *norms_c, | ||
| 120 | uint32_t nlist, | ||
| 121 | Dimension dim, | ||
| 122 | DistanceMetric metric, | ||
| 123 | float *out_sums, | ||
| 124 | uint32_t *out_cnts, | ||
| 125 | float *out_cost) | ||
| 126 | { | ||
| 127 | 3389 | memset(out_sums, 0, (size_t)nlist * dim * sizeof(float)); | |
| 128 | 3389 | memset(out_cnts, 0, nlist * sizeof(uint32_t)); | |
| 129 | 3389 | *out_cost = 0.0f; | |
| 130 | 3389 | kmeans_assign_accumulate( | |
| 131 | samples, | ||
| 132 | NULL, | ||
| 133 | 0, | ||
| 134 | nsamples, | ||
| 135 | centroids, | ||
| 136 | norms_c, | ||
| 137 | nlist, | ||
| 138 | dim, | ||
| 139 | metric, | ||
| 140 | NULL, | ||
| 141 | 0, | ||
| 142 | out_sums, | ||
| 143 | out_cnts, | ||
| 144 | out_cost); | ||
| 145 | 3389 | } | |
| 146 | |||
| 147 | /* | ||
| 148 | * Build ONE root-child's subtree into `slot` (a DSM ring slot indexed by | ||
| 149 | * participant, not by child, so the batched streaming build keeps only | ||
| 150 | * nparticipants subtrees resident). The subtree is clustered from this child's | ||
| 151 | * root-assigned samples, indexed in place (no contiguous per-child copy). | ||
| 152 | */ | ||
| 153 | void | ||
| 154 | 245 | prism_build_child_subtree( | |
| 155 | uint32_t child, | ||
| 156 | uint32_t child_count, | ||
| 157 | int nparticipants, | ||
| 158 | PrismDsmSamples *dsm_samples, | ||
| 159 | PrismDsmRootAssign *dsm_ra, | ||
| 160 | const float *root_cents, | ||
| 161 | uint32_t nlist, | ||
| 162 | uint32_t fan_out, | ||
| 163 | Dimension dim, | ||
| 164 | DistanceMetric metric, | ||
| 165 | uint32_t km_max_iterations, | ||
| 166 | char *slot, | ||
| 167 | uint64_t slot_size) | ||
| 168 | { | ||
| 169 | 245 | uint32_t nlist_c = (nlist + fan_out - 1) / fan_out; | |
| 170 | /* The caller's histogram already counted this child's samples (one pass | ||
| 171 | * over the assignments instead of one per child). */ | ||
| 172 | 245 | uint32_t cc = child_count; | |
| 173 | |||
| 174 | 245 | KMeansOptions opts = VS_KMEANS_OPTIONS_DEFAULT; | |
| 175 | 245 | opts.max_iterations = km_max_iterations; | |
| 176 | 245 | opts.algorithm = KMEANS_ALGO_LLOYD; | |
| 177 | 245 | opts.initial_centroids = NULL; | |
| 178 | |||
| 179 | 245 | const size_t vec_nbytes = (size_t)dim * sizeof(float); | |
| 180 | |||
| 181 | 245 | HKMeansResult *sub = NULL; | |
| 182 |
2/2✓ Branch 0 taken 3 times.
✓ Branch 1 taken 242 times.
|
245 | if (cc == 0) |
| 183 | { | ||
| 184 | 3 | float *seed = vs_alloc(vec_nbytes); | |
| 185 | 3 | memcpy(seed, root_cents + (size_t)child * dim, vec_nbytes); | |
| 186 | 3 | sub = vs_hkmeans_f32( | |
| 187 | seed, 1, NULL, dim, nlist_c, fan_out, metric, &opts); | ||
| 188 | 3 | vs_free(seed); | |
| 189 | } | ||
| 190 | else | ||
| 191 | { | ||
| 192 | 242 | const float *vbase = prism_dsm_worker_samples(dsm_samples, 0); | |
| 193 | 242 | uint32_t mpw = dsm_samples->max_per_worker; | |
| 194 | 242 | uint32_t *idx = vs_alloc((size_t)cc * sizeof(uint32_t)); | |
| 195 | 242 | uint32_t g = 0; | |
| 196 |
3/3✓ Branch 0 taken 376 times.
✓ Branch 1 taken 461 times.
✓ Branch 2 taken 110 times.
|
947 | for (int t = 0; t < nparticipants; t++) |
| 197 | { | ||
| 198 | 705 | const uint32_t *ra = prism_dsm_root_assignments(dsm_ra, t); | |
| 199 | 705 | uint32_t n = prism_dsm_sample_counts(dsm_samples)[t]; | |
| 200 |
2/2✓ Branch 0 taken 1258104 times.
✓ Branch 1 taken 705 times.
|
1258809 | for (uint32_t i = 0; i < n; i++) |
| 201 |
2/2✓ Branch 0 taken 147769 times.
✓ Branch 1 taken 1110335 times.
|
1258104 | if (ra[i] == child) |
| 202 | 147769 | idx[g++] = (uint32_t)t * mpw + i; | |
| 203 | } | ||
| 204 | 242 | sub = vs_hkmeans_f32( | |
| 205 | vbase, cc, idx, dim, nlist_c, fan_out, metric, &opts); | ||
| 206 | 242 | vs_free(idx); | |
| 207 | } | ||
| 208 | |||
| 209 |
1/2✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
|
245 | if (sub != NULL) |
| 210 | { | ||
| 211 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 245 times.
|
245 | if ((uint64_t)sub->total_size > slot_size) |
| 212 | ✗ | vs_error( | |
| 213 | "prism: subtree blob %u exceeds slot " | ||
| 214 | "(%u > %" PRIu64 ")", | ||
| 215 | child, | ||
| 216 | sub->total_size, | ||
| 217 | slot_size); | ||
| 218 | 245 | memcpy(slot, sub, sub->total_size); | |
| 219 | 245 | vs_free(sub); | |
| 220 | } | ||
| 221 | 245 | } | |
| 222 | |||
| 223 | void | ||
| 224 | 197 | prism_pbuild_count_children( | |
| 225 | PrismDsmSamples *dsm_samples, | ||
| 226 | PrismDsmRootAssign *dsm_ra, | ||
| 227 | int nparticipants, | ||
| 228 | uint32_t km_k, | ||
| 229 | uint32_t *out_counts) | ||
| 230 | { | ||
| 231 | 197 | memset(out_counts, 0, (size_t)km_k * sizeof(uint32_t)); | |
| 232 |
2/2✓ Branch 0 taken 587 times.
✓ Branch 1 taken 197 times.
|
784 | for (int t = 0; t < nparticipants; t++) |
| 233 | { | ||
| 234 | 587 | const uint32_t *ra = prism_dsm_root_assignments(dsm_ra, t); | |
| 235 | 587 | uint32_t n = prism_dsm_sample_counts(dsm_samples)[t]; | |
| 236 |
2/2✓ Branch 0 taken 617718 times.
✓ Branch 1 taken 587 times.
|
618305 | for (uint32_t i = 0; i < n; i++) |
| 237 | 617718 | out_counts[ra[i]]++; | |
| 238 | } | ||
| 239 | 197 | } | |
| 240 | |||
| 241 | void | ||
| 242 | 145 | prism_pbuild_stream_subtrees( | |
| 243 | int participant_id, | ||
| 244 | int nparticipants, | ||
| 245 | PrismDsmSamples *dsm_samples, | ||
| 246 | PrismDsmRootAssign *dsm_ra, | ||
| 247 | const float *root_cents, | ||
| 248 | uint32_t km_k, | ||
| 249 | uint32_t nlist, | ||
| 250 | uint32_t fan_out, | ||
| 251 | Dimension dim, | ||
| 252 | DistanceMetric metric, | ||
| 253 | uint32_t km_max_iterations, | ||
| 254 | char *subtrees_base, | ||
| 255 | uint64_t slot_size, | ||
| 256 | Barrier *barrier, | ||
| 257 | PrismBatchCb batch_cb, | ||
| 258 | void *cb_arg, | ||
| 259 | uint32_t *out_child_order) | ||
| 260 | { | ||
| 261 | 145 | uint32_t np = (uint32_t)nparticipants; | |
| 262 | 145 | uint32_t nbatches = (km_k + np - 1) / np; | |
| 263 | 57 | char *slot = | |
| 264 | 145 | prism_dsm_child_subtree(subtrees_base, participant_id, slot_size); | |
| 265 | |||
| 266 | /* One histogram pass over the root assignments feeds every child's | ||
| 267 | * sample count (prism_build_child_subtree needs it, and counting per | ||
| 268 | * child would re-scan the assignments km_k times). */ | ||
| 269 | 145 | uint32_t *child_count = vs_alloc0((size_t)km_k * sizeof(uint32_t)); | |
| 270 | 145 | prism_pbuild_count_children( | |
| 271 | dsm_samples, dsm_ra, nparticipants, km_k, child_count); | ||
| 272 | |||
| 273 | /* Schedule children largest-first (LPT): every batch waits for its | ||
| 274 | * slowest subtree at two barriers, so the skewed children must land in | ||
| 275 | * the full batches, leaving the tail batch the small ones. Counts are | ||
| 276 | * identical for every participant (same shared assignments) and ties | ||
| 277 | * break on the child id, so all participants and the leader's blob | ||
| 278 | * replay derive the same order with no coordination. */ | ||
| 279 | 145 | uint32_t *order = vs_alloc((size_t)km_k * sizeof(uint32_t)); | |
| 280 |
3/3✓ Branch 0 taken 376 times.
✓ Branch 1 taken 428 times.
✓ Branch 2 taken 57 times.
|
861 | for (uint32_t i = 0; i < km_k; i++) |
| 281 | 716 | order[i] = i; | |
| 282 |
2/2✓ Branch 0 taken 571 times.
✓ Branch 1 taken 145 times.
|
716 | for (uint32_t i = 1; i < km_k; i++) |
| 283 | { | ||
| 284 | 571 | uint32_t id = order[i]; | |
| 285 | 571 | uint32_t j = i; | |
| 286 |
6/6✓ Branch 0 taken 1347 times.
✓ Branch 1 taken 199 times.
✓ Branch 2 taken 975 times.
✓ Branch 3 taken 372 times.
✓ Branch 4 taken 63 times.
✓ Branch 5 taken 145 times.
|
1546 | while (j > 0 && (child_count[order[j - 1]] < child_count[id] || |
| 287 |
2/2✓ Branch 0 taken 14 times.
✓ Branch 1 taken 213 times.
|
227 | (child_count[order[j - 1]] == child_count[id] && |
| 288 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
|
14 | order[j - 1] > id))) |
| 289 | { | ||
| 290 | 975 | order[j] = order[j - 1]; | |
| 291 | 975 | j--; | |
| 292 | } | ||
| 293 | 571 | order[j] = id; | |
| 294 | } | ||
| 295 |
2/2✓ Branch 0 taken 51 times.
✓ Branch 1 taken 94 times.
|
145 | if (out_child_order != NULL) |
| 296 | 51 | memcpy(out_child_order, order, (size_t)km_k * sizeof(uint32_t)); | |
| 297 | |||
| 298 |
2/2✓ Branch 0 taken 293 times.
✓ Branch 1 taken 143 times.
|
436 | for (uint32_t b = 0; b < nbatches; b++) |
| 299 | { | ||
| 300 | 293 | uint32_t my_idx = b * np + (uint32_t)participant_id; | |
| 301 |
2/2✓ Branch 0 taken 245 times.
✓ Branch 1 taken 48 times.
|
293 | if (my_idx < km_k) |
| 302 | { | ||
| 303 | 245 | uint32_t my_child = order[my_idx]; | |
| 304 | 245 | prism_build_child_subtree( | |
| 305 | my_child, | ||
| 306 | 245 | child_count[my_child], | |
| 307 | nparticipants, | ||
| 308 | dsm_samples, | ||
| 309 | dsm_ra, | ||
| 310 | root_cents, | ||
| 311 | nlist, | ||
| 312 | fan_out, | ||
| 313 | dim, | ||
| 314 | metric, | ||
| 315 | km_max_iterations, | ||
| 316 | slot, | ||
| 317 | slot_size); | ||
| 318 | } | ||
| 319 | |||
| 320 | /* All participants have built this batch's subtrees into their slots. | ||
| 321 | */ | ||
| 322 | 293 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 323 | |||
| 324 | /* Leader consumes the batch before the slots are reused. */ | ||
| 325 |
3/4✓ Branch 0 taken 106 times.
✓ Branch 1 taken 185 times.
✓ Branch 2 taken 60 times.
✗ Branch 3 not taken.
|
291 | if (participant_id == 0 && batch_cb != NULL) |
| 326 | { | ||
| 327 | 106 | uint32_t base = b * np; | |
| 328 | 106 | uint32_t bs = km_k - base; | |
| 329 |
2/2✓ Branch 0 taken 28 times.
✓ Branch 1 taken 32 times.
|
106 | if (bs > np) |
| 330 | 28 | bs = np; | |
| 331 | 106 | batch_cb(cb_arg, order + base, bs, subtrees_base, slot_size); | |
| 332 | } | ||
| 333 | |||
| 334 | /* Leader done with the batch; slots free for the next batch. */ | ||
| 335 | 291 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 336 | } | ||
| 337 | |||
| 338 | 143 | vs_free(order); | |
| 339 | 143 | vs_free(child_count); | |
| 340 | 143 | } | |
| 341 | |||
| 342 | /* ---------------------------------------------------------------- | ||
| 343 | * Phases 1, 2, 2b: per-participant execution, shared by leader and workers | ||
| 344 | * | ||
| 345 | * The leader runs as participant 0 and calls these exactly as a worker does, | ||
| 346 | * rather than duplicating the bodies. The phase barriers live inside each | ||
| 347 | * function, so the leader and workers stay in lockstep by construction; the | ||
| 348 | * leader-only coordination (seeding the initial centroids, reducing the | ||
| 349 | * per-iteration accumulators) is gated on participant_id == 0 and runs at the | ||
| 350 | * barrier rendezvous while the other participants wait. | ||
| 351 | * ---------------------------------------------------------------- */ | ||
| 352 | |||
| 353 | void | ||
| 354 | 326 | prism_pbuild_exec_sampling( | |
| 355 | int participant_id, | ||
| 356 | Relation heap, | ||
| 357 | Relation index, | ||
| 358 | struct IndexInfo *index_info, | ||
| 359 | PrismBuildShared *shared, | ||
| 360 | PrismDsmSamples *dsm_samples, | ||
| 361 | Barrier *barrier) | ||
| 362 | { | ||
| 363 | 326 | Dimension dim = shared->dim; | |
| 364 | |||
| 365 | /* Stride to subsample ~max_samples_per_worker from this participant's | ||
| 366 | * share of the heap. */ | ||
| 367 | 326 | double est_rows = RelationGetNumberOfBlocks(heap) * | |
| 368 | 326 | (BLCKSZ / (double)(dim * sizeof(float) + 32)); | |
| 369 | 326 | double est_per = est_rows / shared->nparticipants; | |
| 370 | 326 | uint32_t stride = 1; | |
| 371 |
2/2✓ Branch 0 taken 99 times.
✓ Branch 1 taken 227 times.
|
326 | if (est_per > shared->max_samples_per_worker) |
| 372 | 99 | stride = (uint32_t)(est_per / shared->max_samples_per_worker); | |
| 373 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 196 times.
|
281 | if (stride < 1) |
| 374 | 130 | stride = 1; | |
| 375 | |||
| 376 | 326 | SampleCbState sc = { | |
| 377 | 326 | .samples = prism_dsm_worker_samples(dsm_samples, participant_id), | |
| 378 | .count = 0, | ||
| 379 | .seen = 0, | ||
| 380 | 196 | .max_samples = shared->max_samples_per_worker, | |
| 381 | .stride = stride, | ||
| 382 | .stride_counter = 0, | ||
| 383 | .dim = dim, | ||
| 384 | 326 | .metric = shared->metric, | |
| 385 | }; | ||
| 386 | |||
| 387 | /* progress: only the leader (participant 0) drives the PG progress view. | ||
| 388 | * The scan count is discarded: the posting pass is the one full scan that | ||
| 389 | * feeds shared->reltuples, so the extra passes must not double-count. */ | ||
| 390 | 326 | (void)prism_build_scan( | |
| 391 | heap, | ||
| 392 | index, | ||
| 393 | index_info, | ||
| 394 | shared, | ||
| 395 | true, | ||
| 396 | participant_id == 0, | ||
| 397 | prism_sample_cb, | ||
| 398 | &sc); | ||
| 399 | |||
| 400 | 326 | prism_dsm_sample_counts(dsm_samples)[participant_id] = sc.count; | |
| 401 | 326 | prism_dsm_sample_seen(dsm_samples)[participant_id] = sc.seen; | |
| 402 | |||
| 403 | /* Flush the sub-batch remainder so the participant's contribution to | ||
| 404 | * the progress view is exact (and so every participant makes at least | ||
| 405 | * one flush per scan, whatever share the work-stealing gave it). */ | ||
| 406 | 326 | prism_build_progress_incr_tuples((int64_t)(sc.seen & 1023)); | |
| 407 | |||
| 408 | /* Barrier: all participants done sampling. */ | ||
| 409 | 326 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 410 | 326 | } | |
| 411 | |||
| 412 | uint32_t | ||
| 413 | 326 | prism_pbuild_exec_kmeans( | |
| 414 | int participant_id, | ||
| 415 | PrismBuildShared *shared, | ||
| 416 | PrismDsmSamples *dsm_samples, | ||
| 417 | char *centroids_base, | ||
| 418 | char *km_workers_base, | ||
| 419 | Barrier *barrier) | ||
| 420 | { | ||
| 421 | 326 | Dimension dim = shared->dim; | |
| 422 | 326 | uint32_t km_k = shared->km_k; | |
| 423 | 326 | int nparticipants = shared->nparticipants; | |
| 424 | 326 | float *cents = prism_dsm_centroids(centroids_base); | |
| 425 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
326 | float *norms_c = prism_dsm_norms_c(centroids_base, km_k, dim); |
| 426 | /* One seed vector, and the whole km_k-centroid array. */ | ||
| 427 | 326 | const size_t vec_nbytes = (size_t)dim * sizeof(float); | |
| 428 | 326 | const size_t cents_nbytes = (size_t)km_k * vec_nbytes; | |
| 429 | |||
| 430 | /* | ||
| 431 | * Leader-only seed: pick km_k initial centroids from the pooled samples, | ||
| 432 | * spread evenly across the concatenation of every participant's slot. | ||
| 433 | * Centralizing it (rather than slicing by participant) keeps init | ||
| 434 | * independent of how many workers launched and robust to a sample-starved | ||
| 435 | * one. The result varies run to run (work-stealing scan order), so the | ||
| 436 | * parallel build is not bit-reproducible. | ||
| 437 | */ | ||
| 438 |
2/2✓ Branch 0 taken 130 times.
✓ Branch 1 taken 196 times.
|
326 | if (participant_id == 0) |
| 439 | { | ||
| 440 | 86 | uint32_t total_ns = 0; | |
| 441 |
2/2✓ Branch 0 taken 326 times.
✓ Branch 1 taken 130 times.
|
456 | for (int t = 0; t < nparticipants; t++) |
| 442 | 326 | total_ns += prism_dsm_sample_counts(dsm_samples)[t]; | |
| 443 | |||
| 444 |
3/4✓ Branch 0 taken 128 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 86 times.
✗ Branch 3 not taken.
|
130 | uint32_t step = (km_k > 0 && total_ns >= km_k) ? total_ns / km_k : 1; |
| 445 |
2/2✓ Branch 0 taken 750 times.
✓ Branch 1 taken 130 times.
|
880 | for (uint32_t i = 0; i < km_k; i++) |
| 446 | { | ||
| 447 |
2/2✓ Branch 0 taken 749 times.
✓ Branch 1 taken 1 times.
|
750 | uint32_t gidx = (total_ns > 0) ? (i * step) % total_ns : 0; |
| 448 | |||
| 449 | 750 | int t = 0; | |
| 450 | 750 | uint32_t base = 0; | |
| 451 |
2/2✓ Branch 0 taken 1554 times.
✓ Branch 1 taken 1 times.
|
1555 | while (t < nparticipants && |
| 452 |
3/3✓ Branch 0 taken 576 times.
✓ Branch 1 taken 790 times.
✓ Branch 2 taken 188 times.
|
1554 | base + prism_dsm_sample_counts(dsm_samples)[t] <= gidx) |
| 453 | { | ||
| 454 | 805 | base += prism_dsm_sample_counts(dsm_samples)[t]; | |
| 455 | 805 | t++; | |
| 456 | } | ||
| 457 |
2/2✓ Branch 0 taken 749 times.
✓ Branch 1 taken 1 times.
|
750 | if (t < nparticipants) |
| 458 | 749 | memcpy(cents + (size_t)i * dim, | |
| 459 | 749 | prism_dsm_worker_samples(dsm_samples, t) + | |
| 460 | 749 | (size_t)(gidx - base) * dim, | |
| 461 | vec_nbytes); | ||
| 462 | } | ||
| 463 | |||
| 464 |
2/2✓ Branch 0 taken 121 times.
✓ Branch 1 taken 9 times.
|
130 | if (shared->metric == DISTANCE_L2) |
| 465 |
2/2✓ Branch 0 taken 693 times.
✓ Branch 1 taken 121 times.
|
814 | for (uint32_t j = 0; j < km_k; j++) |
| 466 | 693 | norms_c[j] = vs_l2_norm_squared(cents + (size_t)j * dim, dim); | |
| 467 | } | ||
| 468 | |||
| 469 | /* Barrier: initial centroids + norms ready; guards the reads below against | ||
| 470 | * the leader's seed write. */ | ||
| 471 | 326 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 472 | |||
| 473 | 326 | float *my_samples = prism_dsm_worker_samples(dsm_samples, participant_id); | |
| 474 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
326 | uint32_t my_n = prism_dsm_sample_counts(dsm_samples)[participant_id]; |
| 475 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
326 | float *my_sums = prism_dsm_km_worker_sums( |
| 476 | km_workers_base, km_k, dim, participant_id); | ||
| 477 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
326 | uint32_t *my_cnts = prism_dsm_km_worker_cnts( |
| 478 | km_workers_base, km_k, dim, participant_id); | ||
| 479 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 86 times.
|
326 | float *my_cost = prism_dsm_km_worker_cost( |
| 480 | km_workers_base, km_k, dim, participant_id); | ||
| 481 | |||
| 482 | /* | ||
| 483 | * Working set for the reduce. Leader-only, and every field is the | ||
| 484 | * same size on every iteration, so it is allocated once here rather | ||
| 485 | * than inside the loop. KMeansReduce documents what each field | ||
| 486 | * holds; the locals are kept writable for the fills and the frees, | ||
| 487 | * while the struct exposes to the merge only what it may read. | ||
| 488 | */ | ||
| 489 | 326 | float *prev_centroids = NULL; | |
| 490 | 326 | const float **worker_vector_sums = NULL; | |
| 491 | 326 | const uint32_t **worker_vector_counts = NULL; | |
| 492 | 326 | float *worker_costs = NULL; | |
| 493 | 326 | uint32_t *cluster_vector_counts = NULL; | |
| 494 | 326 | KMeansReduce reduce = {0}; | |
| 495 | |||
| 496 |
2/2✓ Branch 0 taken 130 times.
✓ Branch 1 taken 196 times.
|
326 | if (participant_id == 0) |
| 497 | { | ||
| 498 | 130 | prev_centroids = vs_alloc(cents_nbytes); | |
| 499 | 130 | worker_vector_sums = vs_alloc((size_t)nparticipants * sizeof(float *)); | |
| 500 | 130 | worker_vector_counts = vs_alloc( | |
| 501 | 86 | (size_t)nparticipants * sizeof(uint32_t *)); | |
| 502 | 130 | worker_costs = vs_alloc((size_t)nparticipants * sizeof(float)); | |
| 503 | 130 | cluster_vector_counts = vs_alloc((size_t)km_k * sizeof(uint32_t)); | |
| 504 | |||
| 505 | /* The accumulators live in shared memory at addresses that do not | ||
| 506 | * change between iterations, so they are resolved once here. Only | ||
| 507 | * the costs are re-read each round, inside the loop. */ | ||
| 508 |
3/3✓ Branch 0 taken 196 times.
✓ Branch 1 taken 216 times.
✓ Branch 2 taken 44 times.
|
456 | for (int t = 0; t < nparticipants; t++) |
| 509 | { | ||
| 510 | 522 | worker_vector_sums[t] = | |
| 511 | 326 | prism_dsm_km_worker_sums(km_workers_base, km_k, dim, t); | |
| 512 | 326 | worker_vector_counts[t] = | |
| 513 | 326 | prism_dsm_km_worker_cnts(km_workers_base, km_k, dim, t); | |
| 514 | } | ||
| 515 | |||
| 516 | 130 | reduce = (KMeansReduce){ | |
| 517 | .worker_vector_sums = worker_vector_sums, | ||
| 518 | .worker_vector_counts = worker_vector_counts, | ||
| 519 | .worker_costs = worker_costs, | ||
| 520 | 130 | .nworkers = (uint32_t)nparticipants, | |
| 521 | .prev_centroids = prev_centroids, | ||
| 522 | .cluster_vector_counts = cluster_vector_counts, | ||
| 523 | }; | ||
| 524 | } | ||
| 525 | 326 | Size km_sz = prism_dsm_km_workers_size(nparticipants, km_k, dim); | |
| 526 | |||
| 527 | 326 | uint32_t iters = 0; | |
| 528 |
2/2✓ Branch 0 taken 3389 times.
✓ Branch 1 taken 105 times.
|
3494 | for (uint32_t iter = 0; iter < shared->km_max_iterations; iter++) |
| 529 | { | ||
| 530 | 3389 | iters++; | |
| 531 | |||
| 532 | /* Every participant assigns + accumulates over its own samples. */ | ||
| 533 | 3389 | prism_km_assign_and_accumulate( | |
| 534 | my_samples, | ||
| 535 | my_n, | ||
| 536 | cents, | ||
| 537 | norms_c, | ||
| 538 | km_k, | ||
| 539 | dim, | ||
| 540 | shared->metric, | ||
| 541 | my_sums, | ||
| 542 | my_cnts, | ||
| 543 | my_cost); | ||
| 544 | |||
| 545 | /* Barrier: all accumulators written; the leader reduces. */ | ||
| 546 | 3389 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 547 | |||
| 548 |
2/2✓ Branch 0 taken 1200 times.
✓ Branch 1 taken 2189 times.
|
3389 | if (participant_id == 0) |
| 549 | { | ||
| 550 | 1200 | memcpy(prev_centroids, cents, cents_nbytes); | |
| 551 | |||
| 552 | /* Only the costs change between iterations; the accumulator | ||
| 553 | * addresses were resolved once above. */ | ||
| 554 |
2/2✓ Branch 0 taken 3389 times.
✓ Branch 1 taken 1200 times.
|
4589 | for (int t = 0; t < nparticipants; t++) |
| 555 | 3389 | worker_costs[t] = *prism_dsm_km_worker_cost( | |
| 556 | km_workers_base, km_k, dim, t); | ||
| 557 | |||
| 558 | 498 | float total_cost; | |
| 559 | 1200 | float shift_sq = kmeans_merge_centroids( | |
| 560 | cents, | ||
| 561 | norms_c, | ||
| 562 | km_k, | ||
| 563 | dim, | ||
| 564 | shared->metric, | ||
| 565 | &reduce, | ||
| 566 | &total_cost); | ||
| 567 | |||
| 568 | 1200 | float tol_sq = shared->km_tolerance * shared->km_tolerance; | |
| 569 | 1200 | shared->km_converged = (shift_sq < tol_sq); | |
| 570 | 1200 | memset(km_workers_base, 0, km_sz); | |
| 571 | } | ||
| 572 | |||
| 573 | /* Barrier: updated centroids + convergence flag visible to all. */ | ||
| 574 | 3389 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 575 | |||
| 576 |
2/2✓ Branch 0 taken 1540 times.
✓ Branch 1 taken 1849 times.
|
3389 | if (shared->km_converged) |
| 577 | 128 | break; | |
| 578 | } | ||
| 579 | |||
| 580 |
2/2✓ Branch 0 taken 130 times.
✓ Branch 1 taken 196 times.
|
326 | if (participant_id == 0) |
| 581 | { | ||
| 582 | 130 | vs_free(cluster_vector_counts); | |
| 583 | 130 | vs_free(worker_costs); | |
| 584 | 130 | vs_free(worker_vector_counts); | |
| 585 | 130 | vs_free(worker_vector_sums); | |
| 586 | 130 | vs_free(prev_centroids); | |
| 587 | } | ||
| 588 | |||
| 589 | 326 | return iters; | |
| 590 | } | ||
| 591 | |||
| 592 | void | ||
| 593 | 326 | prism_pbuild_exec_root_assign( | |
| 594 | int participant_id, | ||
| 595 | PrismBuildShared *shared, | ||
| 596 | PrismDsmSamples *dsm_samples, | ||
| 597 | PrismDsmRootAssign *dsm_ra, | ||
| 598 | char *centroids_base, | ||
| 599 | Barrier *barrier) | ||
| 600 | { | ||
| 601 | 326 | Dimension dim = shared->dim; | |
| 602 | 326 | uint32_t km_k = shared->km_k; | |
| 603 | 326 | float *cents = prism_dsm_centroids(centroids_base); | |
| 604 | 326 | float *norms_c = prism_dsm_norms_c(centroids_base, km_k, dim); | |
| 605 | |||
| 606 | 522 | kmeans_assign( | |
| 607 | 326 | prism_dsm_worker_samples(dsm_samples, participant_id), | |
| 608 | 0, | ||
| 609 | 326 | prism_dsm_sample_counts(dsm_samples)[participant_id], | |
| 610 | cents, | ||
| 611 | norms_c, | ||
| 612 | km_k, | ||
| 613 | dim, | ||
| 614 | shared->metric, | ||
| 615 | prism_dsm_root_assignments(dsm_ra, participant_id)); | ||
| 616 | |||
| 617 | /* Barrier: all root assignments written; subtree gather can read them. */ | ||
| 618 | 326 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 619 | 326 | } | |
| 620 | |||
| 621 | /* ---------------------------------------------------------------- | ||
| 622 | * Phase 2.5: parallel full-table leaf-encode-reference refinement | ||
| 623 | * | ||
| 624 | * Page-backed mirror of the serial serial_refine_heads: route every row | ||
| 625 | * exactly as the query/insert do (prism_query_route k=1 over the centroid | ||
| 626 | * pages, then head -> leaf), accumulate per-leaf means into the tiled DSM | ||
| 627 | * accumulator, and rewrite each leaf's head-page pt_centroid to the full-table | ||
| 628 | * mean. No in-RAM tree. | ||
| 629 | * ---------------------------------------------------------------- */ | ||
| 630 | |||
| 631 | typedef struct RefineCbState | ||
| 632 | { | ||
| 633 | PrismBuildShared *shared; | ||
| 634 | PrismQueryState *qs; /* page-backed router (workers) */ | ||
| 635 | BlockNumber first_posting; /* head -> leaf: leaf = head - first_posting */ | ||
| 636 | double *sums; /* shared accumulator, indexed leaf - tile_lo */ | ||
| 637 | uint64_t *counts; /* shared accumulator */ | ||
| 638 | Dimension dim; | ||
| 639 | bool cosine; | ||
| 640 | uint32_t tile_lo; /* accumulate only leaves in [tile_lo, tile_hi) */ | ||
| 641 | uint32_t tile_hi; | ||
| 642 | float *scratch; /* per-participant normalized copy (cosine) */ | ||
| 643 | } RefineCbState; | ||
| 644 | |||
| 645 | static void | ||
| 646 | 12600 | prism_refine_cb(void *state, ItemPointerData tid, const float *vec) | |
| 647 | { | ||
| 648 | 12600 | RefineCbState *rs = (RefineCbState *)state; | |
| 649 | 12600 | Dimension dim = rs->dim; | |
| 650 | |||
| 651 | 12600 | (void)tid; | |
| 652 | |||
| 653 | 12600 | uint32_t idx; | |
| 654 | 25200 | const float *v = prism_refine_route_row( | |
| 655 | 12600 | rs->qs, | |
| 656 | rs->first_posting, | ||
| 657 | vec, | ||
| 658 | dim, | ||
| 659 | 12600 | rs->cosine, | |
| 660 | rs->scratch, | ||
| 661 | rs->tile_lo, | ||
| 662 | rs->tile_hi, | ||
| 663 | &idx); | ||
| 664 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 12600 times.
|
12600 | if (v == NULL) |
| 665 | ✗ | return; | |
| 666 | |||
| 667 | 12600 | uint32_t stripe = idx % PRISM_REFINE_LOCK_STRIPES; | |
| 668 | |||
| 669 | 12600 | prism_pbuild_accum_lock(rs->shared, stripe); | |
| 670 | 12600 | double *sum = rs->sums + (size_t)idx * dim; | |
| 671 |
2/2✓ Branch 0 taken 1228800 times.
✓ Branch 1 taken 12600 times.
|
1241400 | for (Dimension j = 0; j < dim; j++) |
| 672 | 1228800 | sum[j] += v[j]; | |
| 673 | 12600 | rs->counts[idx]++; | |
| 674 | 12600 | prism_pbuild_accum_unlock(rs->shared, stripe); | |
| 675 | } | ||
| 676 | |||
| 677 | void | ||
| 678 | 6 | prism_pbuild_exec_refine_paged( | |
| 679 | int participant_id, | ||
| 680 | Relation heap, | ||
| 681 | Relation index, | ||
| 682 | struct IndexInfo *index_info, | ||
| 683 | PrismBuildShared *shared, | ||
| 684 | struct PrismQueryState *qs, | ||
| 685 | BlockNumber first_posting, | ||
| 686 | PrismDsmRefineAccum *accum, | ||
| 687 | Barrier *barrier, | ||
| 688 | PrismLeafWriteFn write_head, | ||
| 689 | void *write_head_ctx) | ||
| 690 | { | ||
| 691 | 6 | Dimension dim = shared->dim; | |
| 692 | 6 | uint32_t nleaves = shared->nlist; /* actual leaf count (published) */ | |
| 693 | 6 | double *sums = prism_dsm_refine_sums(accum); | |
| 694 | 6 | uint64_t *counts = prism_dsm_refine_counts(accum); | |
| 695 | |||
| 696 | /* The accumulator holds at most accum->nleaves leaves (the bounded tile | ||
| 697 | * capacity), so leaves are processed in tiles, re-scanning the heap per | ||
| 698 | * tile. nleaves <= capacity is a single tile (the common case). Leader and | ||
| 699 | * workers derive `tile` identically from the DSM capacity, so the barrier | ||
| 700 | * sequence below stays in lockstep. */ | ||
| 701 | 6 | uint32_t tile = accum->nleaves; | |
| 702 | |||
| 703 | 12 | RefineCbState rs = { | |
| 704 | .shared = shared, | ||
| 705 | .qs = qs, | ||
| 706 | .first_posting = first_posting, | ||
| 707 | .sums = sums, | ||
| 708 | .counts = counts, | ||
| 709 | .dim = dim, | ||
| 710 | 6 | .cosine = (shared->metric == DISTANCE_COSINE), | |
| 711 | 6 | .scratch = vs_alloc((size_t)dim * sizeof(float)), | |
| 712 | }; | ||
| 713 | |||
| 714 | /* Single pass: routing reads only the centroid pages, which refine | ||
| 715 | * never rewrites, so the assignment is a fixed point and one pass | ||
| 716 | * reaches it. */ | ||
| 717 | { | ||
| 718 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6 times.
|
12 | for (uint32_t lo = 0; lo < nleaves; lo += tile) |
| 719 | { | ||
| 720 | 6 | uint32_t hi = (lo + tile < nleaves) ? lo + tile : nleaves; | |
| 721 | 6 | rs.tile_lo = lo; | |
| 722 | 6 | rs.tile_hi = hi; | |
| 723 | |||
| 724 | /* Leader zeroes the (tile-sized) accumulator and resets the scan. | ||
| 725 | */ | ||
| 726 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
|
6 | if (participant_id == 0) |
| 727 | { | ||
| 728 | 2 | memset(sums, 0, (size_t)(hi - lo) * dim * sizeof(double)); | |
| 729 | 2 | memset(counts, 0, (size_t)(hi - lo) * sizeof(uint64_t)); | |
| 730 | 2 | prism_pbuild_rescan(heap, shared); | |
| 731 | } | ||
| 732 | /* Barrier: accumulator cleared + scan reset before anyone scans. | ||
| 733 | */ | ||
| 734 | 6 | BarrierArriveAndWait( | |
| 735 | barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | ||
| 736 | |||
| 737 | /* Workers cooperatively scan the heap and accumulate page-backed; | ||
| 738 | * the leader does not route (it has no qs), it only | ||
| 739 | * clears/divides, mirroring the phase-3 division of labor. */ | ||
| 740 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
|
6 | if (participant_id != 0) |
| 741 | 4 | (void)prism_build_scan( | |
| 742 | heap, | ||
| 743 | index, | ||
| 744 | index_info, | ||
| 745 | shared, | ||
| 746 | true, | ||
| 747 | false, | ||
| 748 | prism_refine_cb, | ||
| 749 | &rs); | ||
| 750 | |||
| 751 | /* Barrier: every row accumulated before the leader divides. */ | ||
| 752 | 6 | BarrierArriveAndWait( | |
| 753 | barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | ||
| 754 | |||
| 755 | /* Leader rewrites this tile's leaf head pages = per-leaf means. */ | ||
| 756 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
|
6 | if (participant_id == 0) |
| 757 | 2 | prism_refine_write_means( | |
| 758 | sums, | ||
| 759 | counts, | ||
| 760 | lo, | ||
| 761 | hi, | ||
| 762 | dim, | ||
| 763 | rs.scratch, | ||
| 764 | write_head, | ||
| 765 | write_head_ctx); | ||
| 766 | /* Barrier: refined heads written before the next tile/pass. */ | ||
| 767 | 6 | BarrierArriveAndWait( | |
| 768 | barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | ||
| 769 | } | ||
| 770 | } | ||
| 771 | |||
| 772 | 6 | vs_free(rs.scratch); | |
| 773 | 6 | } | |
| 774 | |||
| 775 | /* ---------------------------------------------------------------- | ||
| 776 | * Phase 3: posting scan -> cluster-keyed sort (sort-seam path) | ||
| 777 | * | ||
| 778 | * Each worker routes every vector page-backed (the same prism_query_route the | ||
| 779 | * query and insert paths use), RaBitQ-encodes against the target list's head | ||
| 780 | * pt_centroid, and feeds the compact entry into the shared cluster-keyed | ||
| 781 | * sorter (primary + optional SOAR / boundary secondary), via the shared | ||
| 782 | * PrismBuildRouteCtx helper. The leader merges and builds the pages. Memory is | ||
| 783 | * bounded by maintenance_work_mem inside the sorter. | ||
| 784 | * ---------------------------------------------------------------- */ | ||
| 785 | static void | ||
| 786 | 285242 | route_scan_cb(void *state, ItemPointerData tid, const float *vec) | |
| 787 | { | ||
| 788 | 285242 | PrismBuildRouteCtx *ctx = (PrismBuildRouteCtx *)state; | |
| 789 | |||
| 790 | 285242 | prism_build_route_emit(ctx, vec, tid); | |
| 791 |
2/2✓ Branch 0 taken 208 times.
✓ Branch 1 taken 285034 times.
|
285242 | if (((uint64_t)ctx->indtuples & 1023) == 0) |
| 792 | 208 | prism_build_progress_incr_tuples(1024); | |
| 793 | 285242 | } | |
| 794 | |||
| 795 | /* ---------------------------------------------------------------- | ||
| 796 | * Worker entry point — multi-phase build | ||
| 797 | * ---------------------------------------------------------------- */ | ||
| 798 | |||
| 799 | void | ||
| 800 | 196 | prism_parallel_build_main(dsm_segment *seg, shm_toc *toc) | |
| 801 | { | ||
| 802 | 86 | PrismPBuildWorker w; | |
| 803 | 196 | prism_pbuild_worker_attach(toc, &w); | |
| 804 | |||
| 805 | 196 | PrismBuildShared *shared = w.shared; | |
| 806 | 196 | Barrier *barrier = w.barrier; | |
| 807 | 196 | Relation heapRel = w.heapRel; | |
| 808 | 196 | Relation indexRel = w.indexRel; | |
| 809 | 196 | int worker_id = w.worker_id; | |
| 810 | 196 | Dimension dim = w.dim; | |
| 811 | |||
| 812 | /* ---- Phases 1, 2, 2b: the shared per-participant bodies (the leader runs | ||
| 813 | * the very same code as participant 0). Each call includes its phase | ||
| 814 | * barrier(s). ---- */ | ||
| 815 | 196 | void *sample_seg = NULL; | |
| 816 | 86 | PrismDsmSamples *dsm_samples = | |
| 817 | 196 | prism_pbuild_samples_attach(toc, shared, &sample_seg); | |
| 818 | 196 | char *centroids_base = shm_toc_lookup(toc, PRISM_DSM_KEY_CENTROIDS, false); | |
| 819 | 196 | float *cents = prism_dsm_centroids(centroids_base); | |
| 820 | 86 | char *km_workers_base = | |
| 821 | 196 | shm_toc_lookup(toc, PRISM_DSM_KEY_KM_WORKERS, false); | |
| 822 | 86 | PrismDsmRootAssign *dsm_ra = | |
| 823 | 196 | shm_toc_lookup(toc, PRISM_DSM_KEY_ROOT_ASSIGN, false); | |
| 824 | 196 | IndexInfo *indexInfo = BuildIndexInfo(indexRel); | |
| 825 | #ifndef VS_STANDALONE | ||
| 826 | /* The leader marked the build concurrent and scans with an MVCC snapshot; | ||
| 827 | * the worker's freshly built IndexInfo defaults to non-concurrent, so it | ||
| 828 | * must be aligned or heapam's snapshot/OldestXmin assert trips in the scan | ||
| 829 | * (a valid OldestXmin paired with an MVCC snapshot). Matches nbtsort. */ | ||
| 830 | 86 | indexInfo->ii_Concurrent = shared->concurrent; | |
| 831 | #endif | ||
| 832 | |||
| 833 | 196 | prism_pbuild_exec_sampling( | |
| 834 | worker_id, | ||
| 835 | heapRel, | ||
| 836 | indexRel, | ||
| 837 | indexInfo, | ||
| 838 | shared, | ||
| 839 | dsm_samples, | ||
| 840 | barrier); | ||
| 841 | 196 | prism_pbuild_exec_kmeans( | |
| 842 | worker_id, | ||
| 843 | shared, | ||
| 844 | dsm_samples, | ||
| 845 | centroids_base, | ||
| 846 | km_workers_base, | ||
| 847 | barrier); | ||
| 848 | 196 | prism_pbuild_exec_root_assign( | |
| 849 | worker_id, shared, dsm_samples, dsm_ra, centroids_base, barrier); | ||
| 850 | |||
| 851 | /* ---- Phase 2c: batched streaming subtree build (page-backed) ---- | ||
| 852 | * | ||
| 853 | * For a hierarchical tree (>= 2 levels) the workers build per-root-child | ||
| 854 | * subtrees into a bounded ring of slots (subtree DSM is nparticipants | ||
| 855 | * slots, independent of the partition count); the leader consumes each | ||
| 856 | * batch, recording layout counts and keeping the blob in a spillable | ||
| 857 | * store, and later streams every subtree's pages by itself from that | ||
| 858 | * store. Workers pass no callback (leader-only consumes each batch). The | ||
| 859 | * barrier sequence is inside prism_pbuild_stream_subtrees, identical for | ||
| 860 | * leader and workers. A flat (1-level) build has no subtrees — the leader | ||
| 861 | * writes the single level directly, and neither side runs the subtree | ||
| 862 | * barriers. */ | ||
| 863 |
2/2✓ Branch 1 taken 94 times.
✓ Branch 2 taken 102 times.
|
196 | if (vs_hkmeans_nlevels(shared->nlist, shared->fan_out) >= 2) |
| 864 | { | ||
| 865 | /* Ring barrier: the leader creates the subtree ring (sized from the | ||
| 866 | * per-child sample counts) and publishes its handle + slot size | ||
| 867 | * before arriving; attach only after. */ | ||
| 868 | 94 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 869 | 94 | void *ring_seg = NULL; | |
| 870 | 38 | char *subtrees_base = | |
| 871 | 94 | prism_pbuild_subtree_ring_attach(shared, &ring_seg); | |
| 872 | 94 | prism_pbuild_stream_subtrees( | |
| 873 | worker_id, | ||
| 874 | shared->nparticipants, | ||
| 875 | dsm_samples, | ||
| 876 | dsm_ra, | ||
| 877 | cents, | ||
| 878 | shared->km_k, | ||
| 879 | shared->nlist, | ||
| 880 | shared->fan_out, | ||
| 881 | dim, | ||
| 882 | shared->metric, | ||
| 883 | shared->km_max_iterations, | ||
| 884 | subtrees_base, | ||
| 885 | shared->subtree_slot_size, | ||
| 886 | barrier, | ||
| 887 | NULL, | ||
| 888 | NULL, | ||
| 889 | NULL); | ||
| 890 | 92 | prism_pbuild_subtree_ring_release(ring_seg); | |
| 891 | } | ||
| 892 | |||
| 893 | /* Barrier: leader finished streaming the centroid tree + published the | ||
| 894 | * routing state + initialized the sorter; workers build their page-backed | ||
| 895 | * router next (from the just-published shared state). */ | ||
| 896 | 194 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 897 | |||
| 898 | /* ---- Phase 3 setup: page-backed router shared by refine + posting scan | ||
| 899 | * --- */ | ||
| 900 | 194 | void *sortshared = shm_toc_lookup(toc, PRISM_DSM_KEY_SORTSHARED, false); | |
| 901 | 194 | BlockNumber first_posting = shared->first_posting; | |
| 902 | 84 | const float *global_mean = | |
| 903 | 194 | shm_toc_lookup(toc, PRISM_DSM_KEY_GLOBAL_MEAN, false); | |
| 904 | |||
| 905 | /* Exact internal-node centroids, published by the leader before the | ||
| 906 | * tree-ready barrier above: the build descent scores the internal | ||
| 907 | * tree levels exactly (an empty collection leaves the hook inert). */ | ||
| 908 | 194 | void *exact_seg = NULL; | |
| 909 | 84 | PrismExactInternalCentroids exact_centroids; | |
| 910 | 194 | prism_exact_centroid_collection_view( | |
| 911 | 194 | prism_pbuild_exact_centroids_attach(shared, &exact_seg), | |
| 912 | &exact_centroids); | ||
| 913 | |||
| 914 | 194 | uint32_t entry_size = (uint32_t)prism_posting_entry_size(dim); | |
| 915 | 194 | RaBitQParams *rq_params = vs_rabitq_create(dim, shared->rabitq_seed); | |
| 916 | |||
| 917 | /* Per-worker storage over the index for page-backed head/centroid reads | ||
| 918 | * (PG opens one on the worker's indexRel; standalone shares the leader's). | ||
| 919 | */ | ||
| 920 | 194 | VsStorage *storage = prism_pbuild_worker_storage(&w); | |
| 921 | |||
| 922 | /* Routing base — the same PrismIndexBase the query/insert build, so the | ||
| 923 | * worker routes each row identically. nlevels + first_centroid (the | ||
| 924 | * streamed tree's root block) come from the shared state the leader | ||
| 925 | * published; the scales + global mean + fastscan bits also from shared. */ | ||
| 926 | 84 | PrismIndexBase base; | |
| 927 | 278 | prism_build_router_base_init( | |
| 928 | &base, | ||
| 929 | rq_params, | ||
| 930 | storage, | ||
| 931 | dim, | ||
| 932 | 194 | shared->nlevels, | |
| 933 | shared->first_centroid, | ||
| 934 | shared->metric, | ||
| 935 | shared->centroid_format, | ||
| 936 | shared->fastscan_bits, | ||
| 937 | 194 | shared->centroid_error_scale, | |
| 938 | 194 | shared->centroid_beam_scale, | |
| 939 | shared->fan_out, | ||
| 940 | shared->nlist, | ||
| 941 | shared->rabitq_seed, | ||
| 942 | global_mean, | ||
| 943 | 194 | vs_alloc((size_t)dim * sizeof(float))); | |
| 944 | /* Build-only accuracy hook: exact scoring of the internal tree levels | ||
| 945 | * (the query and insert paths never set this). */ | ||
| 946 | 194 | base.exact_internal = &exact_centroids; | |
| 947 | |||
| 948 | 84 | PrismQueryState qs; | |
| 949 | 194 | prism_query_state_init(&qs, &base, 1, PRISM_SECONDARY_TOPK); | |
| 950 | |||
| 951 | /* ---- Phase 2.5: page-backed full-table refine (only when subsampled) | ||
| 952 | * ---- Workers route + accumulate; the leader clears/divides and rewrites | ||
| 953 | * heads. Gated on shared->refine (identical on both sides) so the | ||
| 954 | * barrier sequence stays in lockstep. */ | ||
| 955 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 190 times.
|
194 | if (shared->refine) |
| 956 | { | ||
| 957 | /* The accumulator overlays the sample region (dead since the subtree | ||
| 958 | * phase); the leader initialized its header before the tree-ready | ||
| 959 | * barrier above. */ | ||
| 960 | 4 | PrismDsmRefineAccum *accum = prism_pbuild_refine_overlay(dsm_samples); | |
| 961 | 4 | prism_pbuild_exec_refine_paged( | |
| 962 | worker_id, | ||
| 963 | heapRel, | ||
| 964 | indexRel, | ||
| 965 | indexInfo, | ||
| 966 | shared, | ||
| 967 | &qs, | ||
| 968 | first_posting, | ||
| 969 | accum, | ||
| 970 | barrier, | ||
| 971 | NULL, | ||
| 972 | NULL); | ||
| 973 | } | ||
| 974 | |||
| 975 | /* The samples (and the refine overlay riding in them) are dead; hand the | ||
| 976 | * segment back before the posting sort claims its own memory budget. */ | ||
| 977 | 194 | prism_pbuild_samples_release(dsm_samples, sample_seg); | |
| 978 | 194 | dsm_samples = NULL; | |
| 979 | |||
| 980 | /* ---- Phase 3: posting scan -> cluster-keyed sort (page-backed) ---- */ | ||
| 981 | |||
| 982 | /* worker_id is 1..N for launched workers; the sorter's 0-based worker | ||
| 983 | * index is worker_id - 1. Only the workers sort concurrently -- the | ||
| 984 | * leader's merge-only sorter opens after every worker sort finishes and | ||
| 985 | * uses the full budget -- so the budget splits across the workers | ||
| 986 | * alone; splitting across all participants would strand one share for | ||
| 987 | * the whole sort phase. */ | ||
| 988 | 194 | int nsorters = shared->nparticipants > 1 ? shared->nparticipants - 1 : 1; | |
| 989 | 194 | int worker_wm = shared->work_mem_kb / nsorters; | |
| 990 |
1/2✓ Branch 0 taken 110 times.
✗ Branch 1 not taken.
|
194 | if (worker_wm < 64) |
| 991 | 110 | worker_wm = 64; | |
| 992 | 194 | PrismSorter *sorter = prism_pbuild_sort_begin( | |
| 993 | sortshared, seg, worker_id - 1, 0, false, entry_size, worker_wm); | ||
| 994 | |||
| 995 | 84 | PrismBuildRouteCtx route; | |
| 996 | 194 | prism_build_route_ctx_init( | |
| 997 | &route, | ||
| 998 | &qs, | ||
| 999 | sorter, | ||
| 1000 | rq_params, | ||
| 1001 | storage, | ||
| 1002 | first_posting, | ||
| 1003 | dim, | ||
| 1004 | shared->soar_lambda, | ||
| 1005 | shared->boundary_epsilon); | ||
| 1006 | |||
| 1007 | /* Barrier: the leader reset the scan for the posting phase (after refine | ||
| 1008 | * consumed it); workers may now scan. */ | ||
| 1009 | 194 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 1010 | |||
| 1011 | 194 | double heap_tuples = prism_build_scan( | |
| 1012 | heapRel, | ||
| 1013 | indexRel, | ||
| 1014 | indexInfo, | ||
| 1015 | shared, | ||
| 1016 | true, | ||
| 1017 | false, | ||
| 1018 | route_scan_cb, | ||
| 1019 | &route); | ||
| 1020 | |||
| 1021 | /* Sub-batch progress remainder, as in the sampling scan. */ | ||
| 1022 | 194 | prism_build_progress_incr_tuples( | |
| 1023 | 194 | (int64_t)((uint64_t)route.indtuples & 1023)); | |
| 1024 | |||
| 1025 | 194 | prism_pbuild_sort_performsort(sorter); | |
| 1026 | |||
| 1027 | /* Barrier: every worker has finished sorting; the leader merges next. */ | ||
| 1028 | 194 | BarrierArriveAndWait(barrier, WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN); | |
| 1029 | 194 | BarrierDetach(barrier); | |
| 1030 | |||
| 1031 | 194 | prism_pbuild_worker_add_counts( | |
| 1032 | shared, route.indtuples, route.soar_dupes, heap_tuples); | ||
| 1033 | |||
| 1034 | 194 | prism_pbuild_sort_end(sorter); | |
| 1035 | 194 | prism_build_route_ctx_cleanup(&route); | |
| 1036 | 194 | prism_query_state_cleanup(&qs); | |
| 1037 | 194 | vs_free(base.pt_global_mean); | |
| 1038 | 194 | prism_pbuild_exact_centroids_release(exact_seg); | |
| 1039 | 194 | prism_pbuild_worker_storage_release(storage); | |
| 1040 | |||
| 1041 | 194 | prism_pbuild_worker_detach(toc, &w); | |
| 1042 | 194 | } | |
| 1043 |