shyjsarah commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3879148543
##########
core/src/kmeans.rs:
##########
@@ -728,16 +989,51 @@ pub fn find_topk(
if nprobe == 0 {
return (Vec::new(), Vec::new());
}
- let mut dists: Vec<(f32, usize)> = (0..k)
- .map(|c| (fvec_l2sqr(point, ¢roids[c * d..(c + 1) * d]), c))
- .collect();
+ if nprobe == 1 {
+ let (distance, index) = find_top1(point, centroids, k, d);
+ return (vec![index], vec![distance]);
+ }
+ let mut dists = Vec::with_capacity(k);
+ let four_end = k / 4 * 4;
+ for c in (0..four_end).step_by(4) {
+ let distances = fvec_l2sqr_four(
+ point,
+ ¢roids[c * d..(c + 1) * d],
+ ¢roids[(c + 1) * d..(c + 2) * d],
+ ¢roids[(c + 2) * d..(c + 3) * d],
+ ¢roids[(c + 3) * d..(c + 4) * d],
+ );
+ dists.extend(
+ distances
+ .into_iter()
+ .enumerate()
+ .map(|(offset, distance)| (distance, c + offset)),
+ );
+ }
+ dists.extend((four_end..k).map(|c| (fvec_l2sqr(point, ¢roids[c * d..(c
+ 1) * d]), c)));
select_topk_prefix(&mut dists, nprobe);
let indices: Vec<usize> = dists[..nprobe].iter().map(|&(_, i)|
i).collect();
let distances: Vec<f32> = dists[..nprobe].iter().map(|&(d, _)|
d).collect();
(indices, distances)
}
-/// Batch find top-nprobe nearest centroids for multiple queries using sgemm.
+fn find_topk_batch_direct(
+ queries: &[f32],
+ nq: usize,
+ centroids: &[f32],
+ k: usize,
+ d: usize,
+ nprobe: usize,
+) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
+ let search = |qi| find_topk(&queries[qi * d..(qi + 1) * d], centroids, k,
d, nprobe);
+ if use_parallel_direct_topk(nq, k, nprobe, rayon::current_num_threads()) {
+ (0..nq).into_par_iter().map(search).unzip()
+ } else {
+ (0..nq).map(search).unzip()
Review Comment:
**[Major — performance regression]** Bounding memory by switching this
entire batch to a sequential iterator creates a large latency cliff for small
query batches on large-`nlist` indexes. Matching-checksum 8-thread release runs
at `k=262144, nq=8, nprobe=2` regressed from 3.58 to 10.02 ms for `d=8` and
from 3.05 to 15.13 ms for `d=32` (about 2.8x-5.0x). At `k=1048576, nq=32,
nprobe=2, threads=32`, the same mechanism is about 3x slower. The memory cap is
necessary, but please retain bounded concurrency—for example, O(`nprobe`)
per-worker heaps, centroid tiling with capped reusable buffers, or limiting the
active query count instead of serializing every query.
##########
core/src/kmeans.rs:
##########
@@ -767,39 +1066,204 @@ pub(crate) fn find_topk_batch_with_centroid_norms(
if nprobe == 0 {
return (vec![Vec::new(); nq], vec![Vec::new(); nq]);
}
+ if use_direct_batch(nq, rayon::current_num_threads()) || d == 0 {
+ return find_topk_batch_direct(queries, nq, centroids, k, d, nprobe);
+ }
+ let mut out: Vec<(Vec<usize>, Vec<f32>)> = vec![(Vec::new(), Vec::new());
nq];
+ certified_topk_blocks(
+ queries,
+ nq,
+ centroids,
+ centroid_norms,
+ k,
+ d,
+ nprobe,
+ &mut out,
+ |slot, top| {
+ slot.0.extend(top.iter().map(|&(_, i)| i));
+ slot.1.extend(top.iter().map(|&(dist, _)| dist));
+ },
+ );
+ out.into_iter().unzip()
+}
- if nq == 1 {
- let (indices, distances) = find_topk(&queries[..d], centroids, k, d,
nprobe);
- return (vec![indices], vec![distances]);
+/// Upper bound on the largest centroid norm from the f32 squared norms.
+fn centroid_norm_upper(c_norms: &[f32], d: usize) -> f64 {
+ let max = c_norms.iter().copied().fold(0.0f32, f32::max) as f64;
+ (max / (1.0 - gamma_f32(d))).sqrt()
+}
+
+/// Bound on `|fl(|x|² + |c|² - 2 x·c) - |x - c|²|` for a row of f32 squared
+/// norm `x_norm` and centroids of norm at most `c_max`: the two norms and the
+/// inner product each carry the standard `γ_d` dot-product rounding, the two
+/// combining operations one rounding each.
+fn gemm_error(x_norm: f32, c_max: f64, d: usize) -> f64 {
+ let x_up = (x_norm as f64 / (1.0 - gamma_f32(d))).sqrt();
+ let sum = x_up + c_max;
+ gamma_f32(d + 2) * sum * sum
+}
+
+fn gamma_f32(operations: usize) -> f64 {
+ let error = (f32::EPSILON as f64 / 2.0) * operations as f64;
+ if error < 1.0 {
+ error / (1.0 - error)
+ } else {
+ f64::INFINITY
}
+}
- // Precompute norms
- let q_norms: Vec<f32> = (0..nq)
- .map(|i| fvec_norm_l2sqr(&queries[i * d..(i + 1) * d]))
- .collect();
- // Batch inner products: ip[nq × k] = queries[nq × d] · centroids[k × d]^T
- let mut ip_matrix = vec![0.0f32; nq * k];
- sgemm_a_bt(nq, k, d, 1.0, queries, centroids, 0.0, &mut ip_matrix);
-
- // Extract top-nprobe per query
- let mut all_indices = Vec::with_capacity(nq);
- let mut all_distances = Vec::with_capacity(nq);
-
- for qi in 0..nq {
- let row = qi * k;
- let mut dists: Vec<(f32, usize)> = (0..k)
- .map(|c| {
- let dist = q_norms[qi] + centroid_norms[c] - 2.0 *
ip_matrix[row + c];
- (dist.max(0.0), c)
- })
- .collect();
- select_topk_prefix(&mut dists, nprobe);
+fn certified_error(e_gemm: f64, kth: f64, d: usize) -> f64 {
+ let gamma = gamma_f32(d + 3);
+ let denominator = 1.0 - 2.0 * gamma;
+ if denominator > 0.0 {
+ (e_gemm * (1.0 + gamma) + gamma * kth) / denominator
+ } else {
+ f64::INFINITY
+ }
+}
- all_indices.push(dists[..nprobe].iter().map(|&(_, i)| i).collect());
- all_distances.push(dists[..nprobe].iter().map(|&(d, _)| d).collect());
+/// Direct distances of four centroids at a time (same accumulation order as
+/// `fvec_l2sqr`, so the result equals `find_topk`'s bit for bit).
+fn direct_distances(x: &[f32], centroids: &[f32], d: usize, slots: &mut [(f32,
usize)]) {
+ let (chunks, remainder) = slots.as_chunks_mut::<4>();
+ for group in chunks {
+ let dists = fvec_l2sqr_four(
+ x,
+ ¢roids[group[0].1 * d..(group[0].1 + 1) * d],
+ ¢roids[group[1].1 * d..(group[1].1 + 1) * d],
+ ¢roids[group[2].1 * d..(group[2].1 + 1) * d],
+ ¢roids[group[3].1 * d..(group[3].1 + 1) * d],
+ );
+ for (slot, dist) in group.iter_mut().zip(dists) {
+ slot.0 = dist;
+ }
+ }
+ for slot in remainder {
+ slot.0 = fvec_l2sqr(x, ¢roids[slot.1 * d..(slot.1 + 1) * d]);
}
+}
- (all_indices, all_distances)
+/// Top-`nprobe` centroids of one row from its SGEMM inner products, with the
+/// contract of `find_topk` (direct `fvec_l2sqr` distances, ties by index).
+///
+/// With `E` chosen so `E = gemm_error + γ(d+3)·(D + 2E + gemm_error)`
+/// and `D` the `nprobe`-th SGEMM distance, every candidate through `D + 2E`
+/// is within `E` of its direct distance: a centroid whose SGEMM distance is
+/// below `D - 2E` is in the direct top-`nprobe` set and one above `D + 2E`
+/// is not. When nothing outside the SGEMM top-`nprobe` lies within `D + 2E`
+/// the set is certified as is (the common case); otherwise the direct kernel
+/// decides the centroids inside `[D - 2E, D + 2E]`, typically one or two.
+/// The selected centroids are then re-measured with the direct kernel and
+/// ordered by (distance, index), so the returned values equal `find_topk`.
+#[allow(clippy::too_many_arguments)]
+fn certified_topk_row<'a>(
+ x: &[f32],
+ x_norm: f32,
+ ip: &[f32],
+ centroids: &[f32],
+ c_norms: &[f32],
+ c_max: f64,
+ d: usize,
+ nprobe: usize,
+ dists: &'a mut Vec<(f32, usize)>,
+) -> &'a [(f32, usize)] {
+ let k = c_norms.len();
+ let e_gemm = gemm_error(x_norm, c_max, d);
+ dists.clear();
+ if nprobe == 1 && k > 1 {
+ // Argmin pass tracking the runner-up value (no scratch, no select).
+ let mut best = f32::INFINITY;
+ let mut best_idx = 0;
+ let mut second = f32::INFINITY;
+ for c in 0..k {
+ let approximate = x_norm + c_norms[c] - 2.0 * ip[c];
+ if !approximate.is_finite() {
+ let (ids, distances) = find_topk(x, centroids, k, d, nprobe);
+ dists.extend(distances.into_iter().zip(ids));
+ return dists;
+ }
+ let dist = approximate.max(0.0);
+ if dist < best {
+ second = best;
+ best = dist;
+ best_idx = c;
+ } else if dist < second {
+ second = dist;
+ }
+ }
+ let upper = best as f64 + 2.0 * certified_error(e_gemm, best as f64,
d);
Review Comment:
**[Major — correctness]** The certification error is relative-only and does
not include an additive allowance for gradual underflow. This can make the band
hundreds of times smaller than one f32 subnormal ULP and certify the wrong
centroid. With a finite deterministic `d=1` fixture, scalar `find_topk` selects
centroid 0, while an eight-row one-thread blocked batch certifies centroid 8
for every row; a separate `nprobe=3` fixture also returns a strictly farther
centroid. End-to-end, Disabled and Enabled assign different lists, scalar and
batch `nprobe=1` queries probe different lists, and the disagreement survives
IVFPQ/IVFFlat serialization on both arm64 and x86_64. Please add
architecture-independent absolute underflow/FTZ terms to the norm, GEMM,
combine, and direct-distance certificate—or fail closed to the direct path
whenever terms can enter the subnormal range—and add the exact bit-pattern
regressions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
