shyjsarah commented on code in PR #82: URL: https://github.com/apache/paimon-vector-index/pull/82#discussion_r3877894042
########## core/src/projected_assign.rs: ########## @@ -0,0 +1,1259 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Exact coarse assignment through a low-dimensional projection. +//! +//! `IVFPQIndex::add` must find the nearest coarse centroid of every row. The +//! exact scan is one `n × nlist × d` GEMM, which dominates `add` once PQ +//! encoding is vectorized. This module keeps the result exact but does most of +//! the work in a `d'`-dimensional PCA subspace of the centroids: +//! +//! 1. project the row and centroids with a contractive PCA projection `P` +//! (`d' × d`); +//! 2. for every centroid compute a conservative lower bound on the true +//! squared distance: the projected distance plus the reverse triangle +//! inequality on the components outside the subspace +//! (`|x-c|² >= |P₀(x-c)|² + (|P₀⊥x| - |P₀⊥c|)²`), with every f32 +//! projection, GEMM and f64 rounding error accounted for explicitly; +//! 3. evaluate exact distances in ascending bound order and stop as soon as +//! the next bound cannot beat the best exact distance (branch-and-bound). +//! +//! Step 2 is a valid lower bound for any contractive `P`, so the assignment +//! never depends on how well the PCA converged; `P` only decides how many +//! exact evaluations step 3 needs. Rows without low-dimensional structure +//! degrade to checking every centroid, i.e. the exact scan plus one small +//! GEMM. At `train`, candidate widths are therefore timed on a sample of +//! training vectors, including projection, bound construction, candidate +//! collection, and exact checks. Automatic mode keeps the fastest projection +//! only when it is clearly faster than an exact scan of the same sample. +//! Without a sample, forced mode takes the widest candidate and automatic +//! mode leaves the exact scan in place. +//! +//! The projection is a training artifact: it is derived from the centroids, +//! shared by `IVFPQIndex::from_trained`, and never serialized. + +use crate::blas::{dgemm_a_bt, sgemm_a_bt}; +use crate::distance::{fvec_l2sqr, fvec_l2sqr_four}; +use nalgebra::{DMatrix, SymmetricEigen}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; +use std::time::{Duration, Instant}; + +/// Automatic mode only times widths with `d' * MAX_AUTO_DP_DIVISOR <= d`. +const MAX_AUTO_DP_DIVISOR: usize = 3; +/// Calibration samples below this count as absent. +const MIN_CALIBRATION_ROWS: usize = 256; +/// Rows of the training data used to calibrate `d'`. +pub(crate) const CALIBRATION_ROWS: usize = 2048; +/// With a calibration sample, automatic mode keeps the projection only when +/// its measured time is below this fraction of the exact scan. +const MAX_AUTO_COST_FRACTION: f64 = 0.85; +/// Block subspace iterations; the bound stays valid regardless of convergence. +const SUBSPACE_ITERATIONS: usize = 8; +const SUBSPACE_SEED: u64 = 0x7a5e_c7ed; +/// Keep automatic PCA fitting bounded before allocating its work matrices. +const MAX_AUTO_PCA_BYTES: u128 = 256 * 1024 * 1024; +const MAX_AUTO_PCA_FLOPS: u128 = 100_000_000_000; +/// Smallest projection worth a GEMM; also the rounding unit for `d'`. +const MIN_DP: usize = 8; +/// Rows per parallel block; the projected score matrix is `rows × nlist`. +const MAX_BLOCK_ROWS: usize = 1024; +/// The projection is re-orthonormalized in f64 and scaled by +/// `1 - ORTHO_MARGIN`, so its Gram matrix is `(1 - ORTHO_MARGIN)² · I` up to +/// rounding. `orthonormal_f64` verifies (Gershgorin, including the f64 +/// rounding of the Gram entries) that every eigenvalue of `P Pᵀ` lies within +/// `GRAM_MARGIN` of that diagonal, and declines the projection otherwise, so +/// the constant singular-value bounds below hold for every stored +/// projection. Gram-Schmidt typically lands within ~1e-15. +const ORTHO_MARGIN: f64 = 1e-9; +const GRAM_MARGIN: f64 = 1e-9; +/// Bounds on the squared singular values of the stored projection (see +/// `ORTHO_MARGIN`): `|P v|² / SIGMA_MAX_SQ <= |P₀ v|² <= |P v|² / SIGMA_MIN_SQ` +/// for the orthonormal basis `P₀` of its row space. +const SIGMA_MAX_SQ: f64 = 1.0; +const SIGMA_MIN_SQ: f64 = (1.0 - ORTHO_MARGIN) * (1.0 - ORTHO_MARGIN) * (1.0 - GRAM_MARGIN); +/// Candidates evaluated per branch-and-bound stage before re-pruning. +const STAGE: usize = 32; + +/// How `IVFPQIndex::add` chooses between the exact centroid scan and the +/// projected branch-and-bound. Both produce the exact nearest centroid. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum ProjectedAssignment { + /// Use the projection when the centroids are compressible enough for it + /// to pay off; otherwise scan. + #[default] + Auto, + /// Always build and use the projection (still exact). + Enabled, + /// Always scan. + Disabled, +} + +/// Contractive projection of the coarse centroids plus the per-centroid data +/// the lower bound needs. +#[derive(Clone, Debug)] +pub struct CoarseProjection { + d: usize, + dp: usize, + explained_variance: f64, + /// Centroid mean (f64), subtracted before projecting rows and centroids. + /// Translation cancels in the projected term but not in the residual + /// term, whose norms must not be dominated by a common offset. + mean: Vec<f64>, + /// `dp × d`, scaled so its operator norm does not exceed one. Stored in + /// f64 (exactly the f32 values used for the singular-value bounds) so the + /// row projection carries only f64 rounding. + proj: Vec<f64>, + /// `nlist × dp`: projected centroids. + cents_p: Vec<f32>, + /// `|cents_p[c]|²`. + cents_p_norms: Vec<f64>, + /// `|cents_p[c]|`, cached to keep square roots out of assignment's hot loop. + cents_p_norms_sqrt: Vec<f64>, + /// Upper bound on the f32 GEMM error in each projected centroid. + cents_p_errors: Vec<f64>, + /// Relative error factors of the `rows × nlist` f32 GEMM and of the f64 + /// arithmetic in the bound, fixed by `dp`. + gemm_error_factor: f64, + f64_error_factor: f64, + /// Interval for `|P₀⊥ c|`, the norm of each centroid's component outside + /// the projected subspace, including every rounding error above. + cents_res_lo: Vec<f64>, + cents_res_hi: Vec<f64>, +} + +impl CoarseProjection { + /// Fit a projection to `nlist` centroids of dimension `d`. + /// + /// `calibration` holds `calibration_rows` sample vectors in the centroid + /// space (typically training vectors). Candidate widths up to `d / 3` + /// (`d / 2` when forced) are timed on that sample and the fastest wins; + /// automatic mode keeps the projection only when it also beats the exact + /// scan on the sample by `MAX_AUTO_COST_FRACTION`. Without a sample, + /// forced mode takes the widest width and automatic mode builds nothing. + /// + /// Returns `None` when the projection is not worth it (`force == false`) + /// or cannot be built (too few centroids, zero variance). + pub(crate) fn train( + cents: &[f32], + nlist: usize, + d: usize, + force: bool, + calibration: &[f32], + calibration_rows: usize, + ) -> Option<Self> { + if nlist == 0 || d == 0 { + return None; + } + let calibration_rows = calibration_rows.min(calibration.len() / d); + if !force && calibration_rows < MIN_CALIBRATION_ROWS { + return None; + } + // Largest d' auto mode may pick, rounded down so the gate below holds; + // forced mode searches up to d / 2 (beyond that the projected GEMM + // costs about as much as the scan it is meant to replace). + let auto_limit = (d / MAX_AUTO_DP_DIVISOR) / MIN_DP * MIN_DP; + let block = if force { + (d / 2).max(MIN_DP).min(d) + } else { + auto_limit + }; + if !force && (block < MIN_DP || nlist < 2 * block) { + return None; + } + let block = block.min(nlist); + if !force && !auto_pca_within_budget(nlist, d, block) { + return None; + } + + let mean = column_mean(cents, nlist, d); + let centered: Vec<f32> = cents + .iter() + .enumerate() + .map(|(i, v)| v - mean[i % d] as f32) + .collect(); + let total_variance = centered + .iter() + .map(|v| (*v as f64) * (*v as f64)) + .sum::<f64>() + / nlist as f64; + if total_variance.is_nan() || total_variance <= 0.0 { + return None; + } + + let (basis, eigenvalues) = top_subspace(¢ered, nlist, d, block); + let mut projection = if calibration_rows >= MIN_CALIBRATION_ROWS { + Self::select_width_by_time( + cents, + &mean, + &basis, + nlist, + d, + block, + force, + &calibration[..calibration_rows * d], + calibration_rows, + )? + } else if force { + // No sample to time against: take the widest projection. + Self::from_basis(cents, &mean, &basis, nlist, d, block)? + } else { + return None; + }; + projection.explained_variance = + eigenvalues[..projection.dp].iter().sum::<f64>() / total_variance; + Some(projection) + } + + /// Width minimizing elapsed assignment time on the calibration sample. + #[allow(clippy::too_many_arguments)] + fn select_width_by_time( + cents: &[f32], + mean: &[f64], + basis: &[f32], + nlist: usize, + d: usize, + block: usize, + force: bool, + calibration: &[f32], + calibration_rows: usize, + ) -> Option<Self> { + let exact_elapsed = if force { + None + } else { + let started = Instant::now(); + let _ = + crate::kmeans::find_nearest_batch(calibration, calibration_rows, cents, nlist, d); + Some(started.elapsed()) + }; + let mut best: Option<(Duration, Self)> = None; + let mut widths: Vec<usize> = [8usize, 4, 2] + .iter() + .map(|div| (block / div).div_ceil(MIN_DP) * MIN_DP) + .chain([(block * 3 / 4).div_ceil(MIN_DP) * MIN_DP, block]) + .filter(|&w| w >= MIN_DP && w <= block) + .collect(); + if force && widths.is_empty() { + widths.push(block); + } + widths.sort_unstable(); + widths.dedup(); + for dp in widths.into_iter().rev() { + let Some(candidate) = Self::from_basis(cents, mean, basis, nlist, d, dp) else { + continue; + }; + let started = Instant::now(); + let _ = candidate.assign_with_stats(calibration, calibration_rows, cents, nlist); + let elapsed = started.elapsed(); + if best.as_ref().is_none_or(|(time, _)| elapsed <= *time) { + best = Some((elapsed, candidate)); + } + } + let (elapsed, candidate) = best?; + if let Some(exact_elapsed) = + exact_elapsed.filter(|exact_elapsed| !is_fast_enough(elapsed, *exact_elapsed)) + { + crate::logging::emit_log( + crate::logging::LogLevel::Info, + &format!( + "IVF-PQ projected assignment not used: best width d'={} measured {:.0}% of exact scan time", + candidate.dp, + elapsed.as_secs_f64() / exact_elapsed.as_secs_f64() * 100.0 + ), + ); + return None; + } + Some(candidate) + } + + fn from_basis( + cents: &[f32], + mean: &[f64], + basis: &[f32], + nlist: usize, + d: usize, + dp: usize, + ) -> Option<Self> { + let proj = orthonormal_f64(&basis[..dp * d], dp, d)?; + let mean = mean.to_vec(); + // Center and project in f64, then round to f32 for the row×centroid + // GEMM. `cent_norms` are `|c - mean|²` in f64. + let (cents_p, cent_norms) = project_rows(cents, nlist, d, &proj, dp, &mean); + let cents_p_norms: Vec<f64> = (0..nlist) + .map(|c| norm_l2sqr_f64(¢s_p[c * dp..(c + 1) * dp])) + .collect(); + let cents_p_norms_sqrt: Vec<f64> = cents_p_norms.iter().map(|norm| norm.sqrt()).collect(); + let cents_p_errors: Vec<f64> = (0..nlist) + .map(|c| { + projection_error( + centered_norm_upper(cent_norms[c], d), + cents_p_norms_sqrt[c], + d, + dp, + ) + }) + .collect(); + let (cents_res_lo, cents_res_hi): (Vec<f64>, Vec<f64>) = (0..nlist) + .map(|c| residual_interval(cent_norms[c], d, cents_p_norms[c], cents_p_errors[c])) + .unzip(); + Some(Self { + d, + dp, + explained_variance: 0.0, + mean, + proj, + cents_p, + cents_p_norms, + cents_p_norms_sqrt, + cents_p_errors, + gemm_error_factor: 2.0 * gamma_f32(dp.saturating_mul(2).saturating_add(1)), + f64_error_factor: gamma_f64(dp.saturating_mul(2).saturating_add(8)), + cents_res_lo, + cents_res_hi, + }) + } + + pub fn dimension(&self) -> usize { + self.dp + } + + pub fn explained_variance(&self) -> f64 { + self.explained_variance + } + + /// Exact nearest centroid of every row (ties: smallest centroid index). + /// `data` must already be in the centroid space (normalized / rotated). + pub(crate) fn assign(&self, data: &[f32], n: usize, cents: &[f32], nlist: usize) -> Vec<usize> { + self.assign_with_stats(data, n, cents, nlist).0 + } + + /// `assign` plus the total number of exact distance evaluations, for + /// tests and benchmarks. + fn assign_with_stats( + &self, + data: &[f32], + n: usize, + cents: &[f32], + nlist: usize, + ) -> (Vec<usize>, usize) { + let d = self.d; + let dp = self.dp; + debug_assert_eq!(self.cents_p.len(), nlist * dp); + let mut out = vec![0usize; n]; + // Per worker: bounds (f64) and candidates (f64 + u32), plus projected + // rows and their f64 norms. The planner already counts the score matrix. + let fixed_scratch = nlist.saturating_mul( + (std::mem::size_of::<f64>() + std::mem::size_of::<(f64, u32)>()) + .div_ceil(std::mem::size_of::<f32>()), + ); + let row_scratch = dp.saturating_add(2); Review Comment: **[Minor — incomplete scratch accounting]** Honoring `parallel=false` fixes the original worker-amplification case, but `row_scratch = dp + 2` omits transient buffers held by `project_rows`: `rows*d` f64 centered data, `rows*dp` f64 output, and the f32 projection allocated while those buffers are still live. The actual row cost is approximately `2*d + 3*dp + 2` f32 units. For `d=8192, nlist=dp=512, threads=32`, the planner models about 16 MiB but can select a parallel plan whose projection phase reaches roughly 271 MiB. Please account for the projection and score phases explicitly (or move to reusable caller-owned buffers), and add an end-to-end planner regression rather than testing only the supplied estimates. ########## 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. ########## 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/ivfpq.rs: ########## @@ -252,7 +252,7 @@ impl IVFPQIndex { let code_size = self.pq.code_size(); let mut codes = vec![0u8; n * code_size]; - self.pq.encode_batch(&to_encode, n, &mut codes); + self.pq.encode_batch_blocked(&to_encode, n, &mut codes); Review Comment: **[Major correctness]** This switches persisted code assignment to direct squared differences, but IVF-PQ search still scores those codes with the norm/dot-product distance table. Those arithmetic definitions can disagree materially, not just on ulp-level ties. I reproduced an end-to-end public API case with finite values where exact query rows are assigned code 1, but search scores code 1 as `8388608` and a farther code-0 row as `0`, returning the farther row first. The merge base returns an exact row first. Please make assignment and all IVF-PQ search-table paths use the same numerical distance definition, and add the train/add/write/open/search reproduction as a regression test. ########## core/src/pq.rs: ########## @@ -370,6 +376,113 @@ impl ProductQuantizer { ); } + /// Blocked batch encode for the IVF-PQ add path. + /// + /// The nbits=8 path uses a transposed-codebook kernel: per sub-quantizer + /// the centroids are transposed once to `[dsub][ksub]` so the inner + /// distance loop is stride-1 over `ksub` and runs on SIMD (NEON/AVX2, + /// scalar fallback). This removes the per-vector-per-sub GEMM calls and + /// their distance-table memory traffic. Distances are accumulated directly + /// from coordinate differences, avoiding the cancellation possible in the + /// norm/dot identity used by [`Self::encode_batch`]. The different arithmetic + /// can produce different valid codes, so use this only where codes are freshly + /// produced (index build), not where byte-stable output is pinned. Trained + /// codebooks are expected to contain only finite values. + pub(crate) fn encode_batch_blocked(&self, data: &[f32], n: usize, codes: &mut [u8]) { + // The 4-bit packed path keeps the original per-vector implementation. + if self.nbits == 8 && (0..self.m).all(|sub| self.chunk_dim(sub) >= 4) { + if n < ENCODE_TRANSPOSE_MIN_ROWS { + self.encode_batch_8bit_direct(data, n, codes); Review Comment: **[Major x86_64 performance regression]** This unconditional `<32` branch routes small adds through the generic per-centroid direct loop. On optimized x86_64, its inner `f32::mul_add` lowers to repeated `fmaf` calls. For the included 31-row `d=768,m=192` shape, I measured about `3.63 ms` here versus `1.11 ms` at the merge base (>3x slower). Please retain the previous batch path on x86 unless a target-specific direct kernel is available, or choose the threshold per backend. ########## core/src/pq.rs: ########## @@ -527,6 +640,189 @@ fn argmin_code(distances: &[f32]) -> u8 { best as u8 } +/// Row block for the transposed batch-encode path. 512 rows keeps the +/// per-thread score buffer at 1 KiB and yields ~20 blocks per thread at the +/// production 32,768-row add batch. +const MAX_ENCODE_BLOCK_ROWS: usize = 512; +/// Below this row count the one-time codebook transpose is not worth it. +const ENCODE_TRANSPOSE_MIN_ROWS: usize = 32; + +fn encode_block_rows(rows: usize, workers: usize) -> usize { + rows.div_ceil(workers.max(1)) + .clamp(1, MAX_ENCODE_BLOCK_ROWS) +} + +/// Squared L2 with the same accumulation order as the transposed kernels. +#[inline] +fn fvec_l2sqr_fma(a: &[f32], b: &[f32]) -> f32 { + debug_assert_eq!(a.len(), b.len()); + debug_assert!(!a.is_empty()); + let diff = a[0] - b[0]; + let mut sum = diff * diff; + for i in 1..a.len() { + let diff = a[i] - b[i]; + sum = diff.mul_add(diff, sum); + } + sum +} + +/// Squared-L2 argmin over a transposed sub-codebook. +/// +/// `t` is `[dsub][ksub]` (stride-1 over `j`) and `scores` is a reusable +/// `ksub`-sized scratch buffer. Ties resolve to the smallest index, matching +/// `argmin_code`'s strictly-smaller update rule. +#[inline] +fn score_argmin(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32]) -> u8 { + assert_eq!(q.len().checked_mul(ksub), Some(t.len())); + assert!(scores.len() >= ksub); + #[cfg(target_arch = "aarch64")] + { + if q.len() == 4 && ksub.is_multiple_of(4) { + // SAFETY: NEON is baseline on aarch64; slice bounds checked by caller. + return unsafe { score_argmin_neon_d4(q, t, ksub) }; + } + } + #[cfg(target_arch = "x86_64")] + { + if q.len() == 4 + && ksub.is_multiple_of(8) + && is_x86_feature_detected!("avx2") + && is_x86_feature_detected!("fma") + { + // SAFETY: AVX2 and FMA presence checked above; slice bounds checked by caller. + return unsafe { score_argmin_avx2_d4(q, t, ksub) }; + } + } + score_argmin_scalar(q, t, ksub, scores) Review Comment: **[Major x86_64 performance regression]** The outer gate enables this transposed path for every `dsub>=4`, but only `dsub==4` has an AVX2/NEON specialization. `dsub=8` reaches this scalar fallback, whose `mul_add` also becomes inner-loop `fmaf` calls on optimized x86_64. For the documented 32,768-row `dsub=8` case, I measured about `212 ms` here versus `74 ms` at the merge base (~3x slower). Please keep unsupported shapes on the matrixmultiply path until a wider SIMD kernel is available, or add an arbitrary-`dsub` AVX2 implementation. -- 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]
