shyjsarah commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3900226691
##########
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]
