shyjsarah commented on code in PR #82:
URL:
https://github.com/apache/paimon-vector-index/pull/82#discussion_r3903222561
##########
core/src/pq.rs:
##########
@@ -436,7 +535,29 @@ impl ProductQuantizer {
let c_off = c_base + j * chunk_dim;
fvec_norm_l2sqr(&self.centroids[c_off..c_off +
chunk_dim])
};
- table[t_base + j] = (q_norm + c_norm - 2.0 *
table[t_base + j]).max(0.0);
+ let approximate = q_norm + c_norm - 2.0 * table[t_base
+ j];
+ // Cover the three reductions and final subtraction
before trusting SGEMM.
+ let error_bound =
+ 16.0 * chunk_dim as f32 * f32::EPSILON * (q_norm +
c_norm);
+ table[t_base + j] =
+ if !approximate.is_finite() || approximate <=
error_bound {
Review Comment:
**[Major correctness]** This fallback only asks whether each individual
expanded distance is near its own cancellation bound. It does not protect the
ordering when two large, nonzero candidate distances differ by less than their
numerical error.
I reproduced this on `7296eaa`: the direct/FMA add path selects code 1
(`88314130 < 88314136`), but ordinary in-memory search, unoptimized reader
search, ephemeral reuse, resident precomputation, and optimized
serialized-reader search all rank code 0 first. The new per-sub exact fallback
is not triggered for this subspace.
Please make the direct-encoded 8-bit table contract canonical across all
search paths, or use a pairwise error/margin rule that recomputes candidates
whose ordering is not provably stable. The ephemeral f64-to-f32 path also needs
to preserve the same f32/FMA ordering used for persisted code assignment.
##########
core/src/pq.rs:
##########
@@ -370,6 +386,87 @@ 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) {
+ self.encode_batch_8bit_transposed(data, n, codes);
Review Comment:
**[Major x86_64 fallback regression]** This now transposes the full codebook
for every eligible add, even when runtime dispatch selects the scalar scorer.
For `d=768,m=192`, that is about 0.75 MiB of allocation/writes before
processing a single row.
On a supported x86_64 runtime without AVX2/FMA, versus `314c6bf`, I measured
approximately +39%/+46%/+37%/+45% for rows 1/7/8/31; rows=32 was effectively
unchanged, isolating the fixed transpose cost.
Please restore a small-batch direct path when `score_argmin_kernel` resolves
to scalar, use a backend-specific threshold, or cache the transposed codebook
after training.
--
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]
