diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bc4533d..c206376 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -288,6 +288,19 @@ jobs: - name: Build FFI library run: cargo build --release -p paimon-vindex-ffi + - name: Test AArch64 PQ kernels + if: runner.os == 'macOS' + run: cargo test --release -p paimon-vindex-core pq::tests + + - name: Test no-AVX2/FMA PQ fallback under Rosetta + if: runner.os == 'macOS' + env: + PAIMON_EXPECT_PQ_SGEMM_FALLBACK: '1' + run: | + rustup target add x86_64-apple-darwin + cargo test --release -p paimon-vindex-core --target x86_64-apple-darwin config_from_options_pq_encoding + cargo test --release -p paimon-vindex-core --target x86_64-apple-darwin test_encode_batch_8bit_sgemm_is_thread_and_split_invariant + - name: Build wheel working-directory: python run: | diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs index a62566b..0d79609 100644 --- a/core/benches/ann_bench.rs +++ b/core/benches/ann_bench.rs @@ -667,6 +667,7 @@ fn index_specs(config: &Config) -> Vec { metric: MetricType::L2, use_opq: false, use_approximate_coarse_assignment: true, + canonical_pq_encoding: false, }, searches: vec![ivf_search], }, diff --git a/core/benches/ivfpq_add_bench.rs b/core/benches/ivfpq_add_bench.rs index 3893cb8..1cfffff 100644 --- a/core/benches/ivfpq_add_bench.rs +++ b/core/benches/ivfpq_add_bench.rs @@ -106,16 +106,10 @@ const CASES: [Case; 10] = [ }, ]; -fn new_index( - case: Case, - quantizer_centroids: &[f32], - centroids: &[f32], - norms: &[f32], -) -> IVFPQIndex { +fn new_index(case: Case, quantizer_centroids: &[f32], centroids: &[f32]) -> IVFPQIndex { let mut index = IVFPQIndex::new(case.d, case.nlist, case.m, MetricType::L2, false); index.set_quantizer_centroids(quantizer_centroids.to_vec()); - index.pq.centroids = centroids.to_vec(); - index.pq.centroid_norms_cache = norms.to_vec(); + index.pq.set_centroids(centroids.to_vec()); index } @@ -138,11 +132,6 @@ fn bench_ivfpq_add(c: &mut Criterion) { let centroids = (0..case.m * 256 * dsub) .map(|_| rng.gen_range(-1.0f32..1.0)) .collect::>(); - let norms = centroids - .chunks_exact(dsub) - .map(|centroid| centroid.iter().map(|value| value * value).sum()) - .collect::>(); - group.throughput(Throughput::Elements(case.rows as u64)); group.bench_with_input( BenchmarkId::new( @@ -155,7 +144,7 @@ fn bench_ivfpq_add(c: &mut Criterion) { &case, |b, &case| { b.iter_batched( - || new_index(case, &quantizer_centroids, ¢roids, &norms), + || new_index(case, &quantizer_centroids, ¢roids), |mut index| index.add(black_box(&data), black_box(&ids), case.rows), BatchSize::LargeInput, ); diff --git a/core/benches/ivfpq_batch_reuse_bench.rs b/core/benches/ivfpq_batch_reuse_bench.rs index 7a67db8..4713e42 100644 --- a/core/benches/ivfpq_batch_reuse_bench.rs +++ b/core/benches/ivfpq_batch_reuse_bench.rs @@ -103,9 +103,11 @@ fn main() { .map(|_| rng.gen_range(-1.0f32..1.0)) .collect(), ); - index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub) - .map(|_| rng.gen_range(-1.0f32..1.0)) - .collect(); + index.pq.set_centroids( + (0..M * index.pq.ksub() * index.pq.dsub()) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(), + ); for list_id in 0..NLIST { let first_id = list_id * ROWS_PER_LIST; index.ids[list_id] = (first_id..first_id + ROWS_PER_LIST) diff --git a/core/benches/ivfpq_filter_scan_bench.rs b/core/benches/ivfpq_filter_scan_bench.rs index a886c51..9f8b48b 100644 --- a/core/benches/ivfpq_filter_scan_bench.rs +++ b/core/benches/ivfpq_filter_scan_bench.rs @@ -104,9 +104,11 @@ fn main() { .map(|_| rng.gen_range(-1.0f32..1.0)) .collect(), ); - index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub) - .map(|_| rng.gen_range(-1.0f32..1.0)) - .collect(); + index.pq.set_centroids( + (0..M * index.pq.ksub() * index.pq.dsub()) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(), + ); for list_id in 0..NLIST { let first_id = list_id * ROWS_PER_LIST; index.ids[list_id] = (first_id..first_id + ROWS_PER_LIST) diff --git a/core/benches/ivfpq_train_bench.rs b/core/benches/ivfpq_train_bench.rs index 07ce98b..34a79e1 100644 --- a/core/benches/ivfpq_train_bench.rs +++ b/core/benches/ivfpq_train_bench.rs @@ -103,7 +103,7 @@ fn run_scenario(s: &Scenario) { // Keep results observable so nothing is optimized away. let checksum: f32 = - centroids.iter().take(8).sum::() + pq.centroids.iter().take(8).sum::(); + centroids.iter().take(8).sum::() + pq.centroids().iter().take(8).sum::(); println!( "{:<11} {:>8} {:>5} {:>6} {:>5} {:>8.3} {:>9.3} {:>8.3} {:>14.3}", diff --git a/core/src/diskann.rs b/core/src/diskann.rs index 4b2cb95..30ae782 100644 --- a/core/src/diskann.rs +++ b/core/src/diskann.rs @@ -317,8 +317,8 @@ impl DiskAnnIndex { fn training_plan(&self, n: usize) -> io::Result { pq_training_plan_with_sample_buffers( self.d, - self.pq.m, - self.pq.ksub, + self.pq.m(), + self.pq.ksub(), n, self.build_params.memory_budget_bytes, usize::from(self.metric == MetricType::Cosine) + 1, @@ -338,14 +338,15 @@ impl DiskAnnIndex { let row_ids = checked_bytes(n, size_of::(), "row IDs")?; let row_id_encoding_scratch = row_id_encoding_scratch_bytes(n)?; let pq_codes = checked_bytes(n, self.pq.code_size(), "PQ codes")?; - let pq_codebook = checked_bytes(self.pq.centroids.len(), size_of::(), "PQ codebook")?; + let pq_codebook = + checked_bytes(self.pq.centroids().len(), size_of::(), "PQ codebook")?; let pq_build_distances = if self.build_params.build_distance == DiskAnnBuildDistance::ProductQuantized { self.pq - .m - .checked_mul(self.pq.ksub) - .and_then(|value| value.checked_mul(self.pq.ksub)) + .m() + .checked_mul(self.pq.ksub()) + .and_then(|value| value.checked_mul(self.pq.ksub())) .and_then(|value| value.checked_mul(size_of::())) .ok_or_else(|| invalid_input("DiskANN PQ build-distance table size overflows"))? } else { @@ -438,12 +439,17 @@ impl DiskAnnIndex { } pub(crate) fn validate_for_write(&self) -> io::Result<()> { - validate_diskann_format_configuration(self.d, self.pq.m, self.pq.nbits, self.build_params)?; + validate_diskann_format_configuration( + self.d, + self.pq.m(), + self.pq.nbits(), + self.build_params, + )?; validate_diskann_training_budget( self.d, self.metric, - self.pq.m, - self.pq.nbits, + self.pq.m(), + self.pq.nbits(), self.build_params.memory_budget_bytes, )?; if self.build_params.memory_budget_bytes == 0 { @@ -489,22 +495,22 @@ impl DiskAnnIndex { } let expected_ksub = 1usize - .checked_shl(self.pq.nbits as u32) + .checked_shl(self.pq.nbits() as u32) .ok_or_else(|| invalid_input("DiskANN PQ centroid count overflows usize"))?; let expected_centroids = self .d .checked_mul(expected_ksub) .ok_or_else(|| invalid_input("DiskANN PQ codebook shape overflows usize"))?; - if self.pq.d != self.d - || self.pq.ksub != expected_ksub - || self.pq.centroids.len() != expected_centroids + if self.pq.d() != self.d + || self.pq.ksub() != expected_ksub + || self.pq.centroids().len() != expected_centroids || !self.pq.has_valid_layout() { return Err(invalid_input("DiskANN PQ codebook shape is invalid")); } if let Some(offset) = self .pq - .centroids + .centroids() .iter() .position(|value| !value.is_finite()) { diff --git a/core/src/diskann_io.rs b/core/src/diskann_io.rs index e2106cf..b22a2e8 100644 --- a/core/src/diskann_io.rs +++ b/core/src/diskann_io.rs @@ -1433,10 +1433,8 @@ impl DiskAnnIndexReader { ))); } - let (mut pq, row_ids, pq_codes, adjacency_index) = + let (pq, row_ids, pq_codes, adjacency_index) = read_resident_sections(&mut self.reader, &self.header)?; - pq.try_rebuild_norms_cache() - .map_err(|_| invalid_data("DiskANN PQ norms allocation failed"))?; validate_pq_code_padding(&self.header, &pq_codes)?; let adjacency_validation = AdjacencyValidationCache::new(adjacency_page_count(&self.header)?)?; @@ -2312,8 +2310,8 @@ pub fn write_diskann_index_with_stats( index.d, index.ids.len(), prepared.graph.entry_node, - index.pq.m, - index.pq.nbits, + index.pq.m(), + index.pq.nbits(), index.metric, index.build_params, row_ids_len, @@ -2588,38 +2586,38 @@ fn write_pq_codebook( put_u32( &mut header, 8, - u32::try_from(pq.d).map_err(|_| invalid_input("DiskANN PQ dimension exceeds u32"))?, + u32::try_from(pq.d()).map_err(|_| invalid_input("DiskANN PQ dimension exceeds u32"))?, ); put_u32( &mut header, 12, - u32::try_from(pq.m).map_err(|_| invalid_input("DiskANN PQ m exceeds u32"))?, + u32::try_from(pq.m()).map_err(|_| invalid_input("DiskANN PQ m exceeds u32"))?, ); put_u32( &mut header, 16, - u32::try_from(pq.nbits).map_err(|_| invalid_input("DiskANN PQ bits exceeds u32"))?, + u32::try_from(pq.nbits()).map_err(|_| invalid_input("DiskANN PQ bits exceeds u32"))?, ); put_u32( &mut header, 20, - u32::try_from(pq.ksub).map_err(|_| invalid_input("DiskANN PQ ksub exceeds u32"))?, + u32::try_from(pq.ksub()).map_err(|_| invalid_input("DiskANN PQ ksub exceeds u32"))?, ); put_u32( &mut header, 24, - u32::try_from(pq.chunk_offsets.len()) + u32::try_from(pq.chunk_offsets().len()) .map_err(|_| invalid_input("DiskANN PQ chunk-offset count exceeds u32"))?, ); writer.write_bytes(&header)?; - for &offset in &pq.chunk_offsets { + for &offset in pq.chunk_offsets() { writer.write_bytes( &u32::try_from(offset) .map_err(|_| invalid_input("DiskANN PQ chunk offset exceeds u32"))? .to_le_bytes(), )?; } - for &value in &pq.centroids { + for &value in pq.centroids() { writer.write_bytes(&value.to_le_bytes())?; } Ok(()) @@ -3050,7 +3048,7 @@ fn decode_pq_codebook(bytes: &[u8], header: &DiskAnnHeader) -> io::Result io::Result crate::diskann_io::DiskAnnIndexReader { let query_count = queries.len() / dimension; let pq_m = self.header.pq_m as usize; let pq = self.pq()?; - let pq_ksub = pq.ksub; + let pq_ksub = pq.ksub(); let pq_code_size = pq.code_size(); let pq_codes = self.pq_codes()?; let metric = self.header.metric_type(); @@ -1881,7 +1881,7 @@ impl crate::diskann_io::DiskAnnIndexReader { self.header.max_degree as usize, )?; let pq = self.pq()?; - let distance_table_len = pq.m * pq.ksub; + let distance_table_len = pq.m() * pq.ksub(); pq.compute_distance_table( query, self.header.metric_type(), @@ -2024,7 +2024,7 @@ impl crate::diskann_io::DiskAnnIndexReader { let result = (|| { scratch.begin_rerank(); let pq = self.pq()?; - let distance_table_len = pq.m * pq.ksub; + let distance_table_len = pq.m() * pq.ksub(); pq.compute_distance_table( query, self.header.metric_type(), @@ -5432,7 +5432,7 @@ mod tests { reader.ensure_resident().unwrap(); assert_eq!(reader.header.pq_bits, 4); - assert_eq!(reader.pq().unwrap().ksub, 16); + assert_eq!(reader.pq().unwrap().ksub(), 16); assert_eq!(reader.pq_codes().unwrap().len(), indexed_count); let (result_ids, distances) = reader.search(query, 5, 100).unwrap(); diff --git a/core/src/index.rs b/core/src/index.rs index 773984f..9acdee5 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -155,6 +155,9 @@ pub enum VectorIndexConfig { metric: MetricType, use_opq: bool, use_approximate_coarse_assignment: bool, + /// Use canonical expanded-form PQ encoding instead of the default + /// transposed direct-L2 encoder. + canonical_pq_encoding: bool, }, IvfRq { dimension: usize, @@ -207,6 +210,7 @@ impl VectorIndexConfig { metric, use_opq, use_approximate_coarse_assignment: true, + canonical_pq_encoding: false, }; validate_config(&config)?; Ok(config) @@ -283,6 +287,8 @@ pub struct ResolvedVectorIndexConfig { pub rq_bits: Option, pub use_opq: bool, pub use_approximate_coarse_assignment: bool, + /// Only meaningful for IVF-PQ. + pub canonical_pq_encoding: bool, pub diskann_build: Option, } @@ -310,6 +316,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { rq_bits: None, use_opq: false, use_approximate_coarse_assignment: *use_approximate_coarse_assignment, + canonical_pq_encoding: false, diskann_build: None, }, VectorIndexConfig::IvfPq { @@ -319,6 +326,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { metric, use_opq, use_approximate_coarse_assignment, + canonical_pq_encoding, } => Self { index_type: IndexType::IvfPq, dimension: *dimension, @@ -329,6 +337,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { rq_bits: None, use_opq: *use_opq, use_approximate_coarse_assignment: *use_approximate_coarse_assignment, + canonical_pq_encoding: *canonical_pq_encoding, diskann_build: None, }, VectorIndexConfig::IvfRq { @@ -347,6 +356,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { rq_bits: Some(*bits), use_opq: false, use_approximate_coarse_assignment: *use_approximate_coarse_assignment, + canonical_pq_encoding: false, diskann_build: None, }, VectorIndexConfig::DiskAnn { @@ -365,6 +375,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { rq_bits: None, use_opq: false, use_approximate_coarse_assignment: false, + canonical_pq_encoding: false, diskann_build: Some(*build), }, } @@ -431,6 +442,7 @@ impl VectorIndexBuildPlan { } IndexType::DiskAnn => true, }; + let canonical_pq_encoding = parse_ivf_pq_encoding_option(&mut options, index_type)?; let config = match index_type { IndexType::IvfFlat => VectorIndexConfig::IvfFlat { @@ -460,6 +472,7 @@ impl VectorIndexBuildPlan { None => target_recall.is_some_and(|recall| recall >= 0.9), }, use_approximate_coarse_assignment, + canonical_pq_encoding, }, IndexType::IvfRq => { let explicit_bits = options @@ -639,6 +652,28 @@ fn parse_ivf_coarse_assignment_option(options: &mut ConfigOptions) -> io::Result } } +fn parse_ivf_pq_encoding_option( + options: &mut ConfigOptions, + index_type: IndexType, +) -> io::Result { + let value = options.optional("ivf.pq-encoding"); + if index_type != IndexType::IvfPq { + return match value { + None => Ok(false), + Some(_) => Err(invalid_input( + "option 'ivf.pq-encoding' is only valid for IVF-PQ", + )), + }; + } + match value.as_deref().map(str::trim) { + None | Some("auto") => Ok(false), + Some("canonical") => Ok(true), + Some(_) => Err(invalid_input( + "option 'ivf.pq-encoding' must be auto or canonical", + )), + } +} + fn parse_deployment_profile_option(name: &str, value: &str) -> io::Result { match value.trim() { "auto" => Ok(DeploymentProfile::Auto), @@ -1339,9 +1374,11 @@ impl VectorIndexWriter { metric, use_opq, use_approximate_coarse_assignment, + canonical_pq_encoding, } => { let mut index = IVFPQIndex::new(dimension, nlist, m, metric, use_opq); index.set_approximate_coarse_assignment(use_approximate_coarse_assignment); + index.set_canonical_pq_encoding(canonical_pq_encoding); Self::IvfPq(index) } VectorIndexConfig::IvfRq { @@ -1524,7 +1561,7 @@ impl VectorIndexReader { metric: reader.metric, total_vectors: reader.total_vectors, pq_m: Some(reader.m), - pq_bits: Some(reader.pq.nbits), + pq_bits: Some(reader.pq.nbits()), rq_bits: None, diskann: None, }, @@ -2946,6 +2983,7 @@ mod tests { metric: MetricType::L2, use_opq: false, use_approximate_coarse_assignment: true, + canonical_pq_encoding: false, }, VectorIndexConfig::IvfRq { dimension: 8, @@ -3046,6 +3084,7 @@ mod tests { metric: MetricType::L2, use_opq: false, use_approximate_coarse_assignment: true, + canonical_pq_encoding: false, }) { Ok(_) => panic!("invalid PQ config should be rejected"), Err(err) => err, @@ -4295,6 +4334,68 @@ mod tests { assert!(error.to_string().contains("must be auto or exact")); } + #[test] + fn config_from_options_pq_encoding() { + let supports_transposed = crate::pq::supports_transposed_pq_encoding(); + if std::env::var("PAIMON_EXPECT_PQ_SGEMM_FALLBACK").as_deref() == Ok("1") { + assert!(!supports_transposed, "expected the SGEMM fallback backend"); + } + + let base = [ + ("index.type", "ivf_pq"), + ("dimension", "4"), + ("nlist", "1"), + ("pq.m", "1"), + ("metric", "l2"), + ]; + let with = |extra: &[(&str, &str)]| { + let mut all = base.to_vec(); + all.extend_from_slice(extra); + VectorIndexConfig::from_options(&options(&all)) + }; + + let auto = with(&[]).unwrap(); + assert!(!auto.resolved().canonical_pq_encoding); + assert!( + !with(&[("ivf.pq-encoding", "auto")]) + .unwrap() + .resolved() + .canonical_pq_encoding + ); + let canonical = with(&[("ivf.pq-encoding", "canonical")]).unwrap(); + assert!(canonical.resolved().canonical_pq_encoding); + + let encoded_code = |config| { + let VectorIndexWriter::IvfPq(mut index) = + VectorIndexWriter::from_config(config).unwrap() + else { + unreachable!() + }; + index.set_quantizer_centroids(vec![0.0; 4]); + let mut centroids = vec![100_000_016.0; 4 * 256]; + centroids[0..4].fill(100_000_008.0); + centroids[4..8].fill(100_000_000.0); + index.pq.set_centroids(centroids); + index.add(&[100_000_000.0; 4], &[0], 1); + index.codes[0][0] + }; + assert_eq!(encoded_code(auto), u8::from(supports_transposed)); + assert_eq!(encoded_code(canonical), 0); + + let error = with(&[("ivf.pq-encoding", "sgemm")]).unwrap_err(); + assert!(error.to_string().contains("must be auto or canonical")); + + let error = VectorIndexConfig::from_options(&options(&[ + ("index.type", "ivf_flat"), + ("dimension", "8"), + ("nlist", "4"), + ("metric", "l2"), + ("ivf.pq-encoding", "canonical"), + ])) + .unwrap_err(); + assert!(error.to_string().contains("only valid for IVF-PQ")); + } + #[test] fn config_from_options_rejects_unknown_options() { let err = VectorIndexConfig::from_options(&options(&[ diff --git a/core/src/io.rs b/core/src/io.rs index 3a55823..3067be1 100644 --- a/core/src/io.rs +++ b/core/src/io.rs @@ -280,9 +280,9 @@ fn checked_list_bytes(count: usize, bytes_per_entry: usize) -> io::Result pub fn write_index(index: &IVFPQIndex, out: &mut dyn SeekWrite) -> io::Result<()> { let d = index.d; let nlist = index.nlist; - let m = index.pq.m; - let ksub = index.pq.ksub; - let dsub = index.pq.dsub; + let m = index.pq.m(); + let ksub = index.pq.ksub(); + let dsub = index.pq.dsub(); let code_size = index.pq.code_size(); if ksub == 16 && !m.is_multiple_of(2) { return Err(io::Error::new( @@ -364,7 +364,7 @@ pub fn write_index(index: &IVFPQIndex, out: &mut dyn SeekWrite) -> io::Result<() } write_f32_slice(out, index.quantizer_centroids())?; - write_f32_slice(out, &index.pq.centroids)?; + write_f32_slice(out, index.pq.centroids())?; // Compute offsets for inverted lists // Delta-varint format per list: [base_id: i64][id_bytes_len: u32][id_bytes][codes] @@ -640,16 +640,7 @@ impl IVFPQIndexReader { total_vectors, opq: None, quantizer_centroids: Vec::new(), - pq: ProductQuantizer { - d, - m, - nbits: ksub.trailing_zeros() as usize, - dsub, - ksub, - chunk_offsets: (0..=m).map(|chunk| chunk * dsub).collect(), - centroids: Vec::new(), - centroid_norms_cache: Vec::new(), - }, + pq: ProductQuantizer::with_nbits(d, m, ksub.trailing_zeros() as usize), list_offsets: Vec::new(), list_counts: Vec::new(), list_id_bytes_lens: Vec::new(), @@ -743,17 +734,7 @@ impl IVFPQIndexReader { let pq_centroids = bytes_to_f32_vec(&metadata[position..position + pq_centroid_bytes])?; position += pq_centroid_bytes; - self.pq = ProductQuantizer { - d, - m, - nbits: ksub.trailing_zeros() as usize, - dsub, - ksub, - chunk_offsets: (0..=m).map(|chunk| chunk * dsub).collect(), - centroids: pq_centroids, - centroid_norms_cache: Vec::new(), - }; - self.pq.rebuild_norms_cache(); + self.pq.set_centroids(pq_centroids); self.list_offsets = vec![0i64; nlist]; self.list_counts = vec![0i32; nlist]; @@ -1226,9 +1207,9 @@ fn compute_precomputed_table( nlist: usize, d: usize, ) -> Vec { - let m = pq.m; - let ksub = pq.ksub; - let dsub = pq.dsub; + let m = pq.m(); + let ksub = pq.ksub(); + let dsub = pq.dsub(); let table_size = nlist * m * ksub; let mut table = vec![0.0f32; table_size]; @@ -1247,7 +1228,7 @@ fn compute_precomputed_table( let pq_off = pq_base + j * dsub; let mut ip = 0.0f32; for dd in 0..dsub { - ip += sub_centroid[dd] * pq.centroids[pq_off + dd]; + ip += sub_centroid[dd] * pq.centroids()[pq_off + dd]; } list_table[sub * ksub + j] = pq_norms[sub * ksub + j] + 2.0 * ip; } @@ -1547,7 +1528,7 @@ mod tests { let mut cursor = Cursor::new(&buf); let mut reader = IVFPQIndexReader::open(&mut cursor).unwrap(); - assert_eq!(reader.pq.nbits, 4); + assert_eq!(reader.pq.nbits(), 4); assert_eq!(reader.pq.code_size(), m / 2); let (result_ids, result_dists) = reader.search(&data[0..d], 5, 4).unwrap(); diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 15e530c..a7b1082 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -80,6 +80,7 @@ pub struct IVFPQIndex { /// Block-layout packed codes for 4-bit FastScan. One per list. fastscan_codes: Vec>, coarse_assignment: CoarseAssignment, + canonical_pq_encoding: bool, } impl IVFPQIndex { @@ -117,6 +118,7 @@ impl IVFPQIndex { precomputed_table: Vec::new(), fastscan_codes: Vec::new(), coarse_assignment: CoarseAssignment::default(), + canonical_pq_encoding: false, } } @@ -124,6 +126,16 @@ impl IVFPQIndex { &self.quantizer_centroids } + /// Uses the canonical per-vector PQ encoder instead of the default + /// automatic encoder. This build-time setting is not serialized. + pub(crate) fn set_canonical_pq_encoding(&mut self, canonical: bool) { + assert!( + self.ids.iter().all(Vec::is_empty), + "cannot change PQ encoding after vectors have been added" + ); + self.canonical_pq_encoding = canonical; + } + /// Enables automatic Vamana coarse assignment for large centroid matrices. /// Disable it to keep vector assignment exact. pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) { @@ -174,16 +186,7 @@ impl IVFPQIndex { metric: trained.metric, by_residual: trained.by_residual, quantizer_centroids: trained.quantizer_centroids.clone(), - pq: ProductQuantizer { - d: trained.pq.d, - m: trained.pq.m, - nbits: trained.pq.nbits, - dsub: trained.pq.dsub, - ksub: trained.pq.ksub, - chunk_offsets: trained.pq.chunk_offsets.clone(), - centroids: trained.pq.centroids.clone(), - centroid_norms_cache: trained.pq.centroid_norms_cache.clone(), - }, + pq: trained.pq.clone_without_transposed_cache(), opq: trained.opq.as_ref().map(|o| OPQMatrix { d: o.d, m: o.m, @@ -199,6 +202,7 @@ impl IVFPQIndex { precomputed_table: Vec::new(), fastscan_codes: Vec::new(), coarse_assignment: CoarseAssignment::default(), + canonical_pq_encoding: trained.canonical_pq_encoding, }; index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled()); index @@ -303,7 +307,11 @@ impl IVFPQIndex { let code_size = self.pq.code_size(); let mut codes = vec![0u8; n * code_size]; - self.pq.encode_batch_blocked(&to_encode, n, &mut codes); + if self.canonical_pq_encoding { + self.pq.encode_batch(&to_encode, n, &mut codes); + } else { + self.pq.encode_batch_blocked(&to_encode, n, &mut codes); + } for i in 0..n { let list_id = assignments[i]; @@ -322,7 +330,7 @@ impl IVFPQIndex { /// Build fastscan block codes for 4-bit search acceleration. /// Call after all vectors are added. Lightweight — only reorganizes existing codes. pub fn build_search_structures(&mut self) { - if self.pq.nbits == 4 { + if self.pq.nbits() == 4 { let cs = self.pq.code_size(); self.fastscan_codes = self .codes @@ -345,8 +353,8 @@ impl IVFPQIndex { /// Costs ~10ms to build and uses nlist * M * ksub * 4 bytes of memory. pub fn build_precomputed_table(&mut self) { let d = self.d; - let m = self.pq.m; - let ksub = self.pq.ksub; + let m = self.pq.m(); + let ksub = self.pq.ksub(); let nlist = self.nlist; if self.metric != MetricType::L2 || !self.by_residual { @@ -362,14 +370,15 @@ impl IVFPQIndex { .for_each(|(i, list_table)| { let centroid = &self.quantizer_centroids[i * d..(i + 1) * d]; for sub in 0..m { - let sub_centroid = ¢roid[sub * self.pq.dsub..(sub + 1) * self.pq.dsub]; - let pq_base = sub * ksub * self.pq.dsub; + let sub_centroid = + ¢roid[sub * self.pq.dsub()..(sub + 1) * self.pq.dsub()]; + let pq_base = sub * ksub * self.pq.dsub(); for j in 0..ksub { - let pq_off = pq_base + j * self.pq.dsub; + let pq_off = pq_base + j * self.pq.dsub(); let ip = fvec_inner_product( sub_centroid, - &self.pq.centroids[pq_off..pq_off + self.pq.dsub], + &self.pq.centroids()[pq_off..pq_off + self.pq.dsub()], ); list_table[sub * ksub + j] = pq_norms[sub * ksub + j] + 2.0 * ip; } @@ -413,8 +422,8 @@ impl IVFPQIndex { result_labels: &mut [i64], ) { let d = self.d; - let m = self.pq.m; - let ksub = self.pq.ksub; + let m = self.pq.m(); + let ksub = self.pq.ksub(); let processed_queries = self.preprocess_queries(queries, nq); @@ -428,7 +437,7 @@ impl IVFPQIndex { ); let use_precomputed = !self.precomputed_table.is_empty(); - let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits == 4; + let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits() == 4; let matching_rows_by_list = filter.map(|filter| { let mut probed_lists = vec![false; self.nlist]; for probe_indices in &all_probe_indices { @@ -524,7 +533,7 @@ impl IVFPQIndex { heap.push(dis0 + dists[i], self.ids[list_id][i]); } } - } else if self.pq.nbits == 4 { + } else if self.pq.nbits() == 4 { scan_codes_4bit( &sim_table, &self.codes[list_id], @@ -621,8 +630,8 @@ impl IVFPQIndex { result_labels: &mut [i64], ) { let d = self.d; - let m = self.pq.m; - let ksub = self.pq.ksub; + let m = self.pq.m(); + let ksub = self.pq.ksub(); let processed_queries = self.preprocess_queries(queries, nq); let (all_probe_indices, all_coarse_dists) = kmeans::find_topk_batch( @@ -635,7 +644,7 @@ impl IVFPQIndex { ); let use_precomputed = !self.precomputed_table.is_empty(); - let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits == 4; + let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits() == 4; let results: Vec> = (0..nq) .into_par_iter() @@ -697,7 +706,7 @@ impl IVFPQIndex { for i in 0..scan_count { heap.push(dis0 + dists[i], self.ids[list_id][i]); } - } else if self.pq.nbits == 4 { + } else if self.pq.nbits() == 4 { scan_codes_4bit( &sim_table, &self.codes[list_id], @@ -786,24 +795,24 @@ impl IVFPQIndex { self.by_residual, other.by_residual ))); } - if self.pq.d != other.pq.d - || self.pq.m != other.pq.m - || self.pq.nbits != other.pq.nbits - || self.pq.dsub != other.pq.dsub - || self.pq.ksub != other.pq.ksub + if self.pq.d() != other.pq.d() + || self.pq.m() != other.pq.m() + || self.pq.nbits() != other.pq.nbits() + || self.pq.dsub() != other.pq.dsub() + || self.pq.ksub() != other.pq.ksub() { return Err(invalid_merge_input(format!( "PQ layout mismatch: self=(d={}, m={}, nbits={}, dsub={}, ksub={}), other=(d={}, m={}, nbits={}, dsub={}, ksub={})", - self.pq.d, - self.pq.m, - self.pq.nbits, - self.pq.dsub, - self.pq.ksub, - other.pq.d, - other.pq.m, - other.pq.nbits, - other.pq.dsub, - other.pq.ksub + self.pq.d(), + self.pq.m(), + self.pq.nbits(), + self.pq.dsub(), + self.pq.ksub(), + other.pq.d(), + other.pq.m(), + other.pq.nbits(), + other.pq.dsub(), + other.pq.ksub() ))); } if self.opq.is_some() != other.opq.is_some() { @@ -823,7 +832,7 @@ impl IVFPQIndex { if self.quantizer_centroids != other.quantizer_centroids { return Err(invalid_merge_input("coarse centroids mismatch")); } - if self.pq.centroids != other.pq.centroids { + if self.pq.centroids() != other.pq.centroids() { return Err(invalid_merge_input("PQ codebooks mismatch")); } @@ -1034,37 +1043,37 @@ fn fill_list_precomputed_table( pq_norms: &[f32], table: &mut Vec, ) { - debug_assert_eq!(coarse_centroid.len(), pq.d); - debug_assert_eq!(pq_norms.len(), pq.m * pq.ksub); - table.resize(pq.m * pq.ksub, 0.0); - for sub in 0..pq.m { + debug_assert_eq!(coarse_centroid.len(), pq.d()); + debug_assert_eq!(pq_norms.len(), pq.m() * pq.ksub()); + table.resize(pq.m() * pq.ksub(), 0.0); + for sub in 0..pq.m() { let range = pq.chunk_range(sub); let chunk_dim = range.len(); - let pq_base = range.start * pq.ksub; - for code in 0..pq.ksub { + let pq_base = range.start * pq.ksub(); + for code in 0..pq.ksub() { let pq_offset = pq_base + code * chunk_dim; let mut inner_product = 0.0f32; for dimension in 0..chunk_dim { - inner_product += - coarse_centroid[range.start + dimension] * pq.centroids[pq_offset + dimension]; + inner_product += coarse_centroid[range.start + dimension] + * pq.centroids()[pq_offset + dimension]; } - let table_offset = sub * pq.ksub + code; + let table_offset = sub * pq.ksub() + code; table[table_offset] = pq_norms[table_offset] + 2.0 * inner_product; } } } fn compute_stable_ephemeral_pq_norms(pq: &ProductQuantizer) -> Vec { - let mut norms = vec![0.0f64; pq.m * pq.ksub]; - for sub in 0..pq.m { + let mut norms = vec![0.0f64; pq.m() * pq.ksub()]; + for sub in 0..pq.m() { let range = pq.chunk_range(sub); let chunk_dim = range.len(); - let pq_base = range.start * pq.ksub; - for code in 0..pq.ksub { + let pq_base = range.start * pq.ksub(); + for code in 0..pq.ksub() { let pq_offset = pq_base + code * chunk_dim; - norms[sub * pq.ksub + code] = (0..chunk_dim) + norms[sub * pq.ksub() + code] = (0..chunk_dim) .map(|dimension| { - let value = f64::from(pq.centroids[pq_offset + dimension]); + let value = f64::from(pq.centroids()[pq_offset + dimension]); value * value }) .sum(); @@ -1079,38 +1088,38 @@ fn fill_stable_ephemeral_list_table( pq_norms: &[f64], table: &mut Vec, ) { - table.resize(pq.m * pq.ksub, 0.0); - for sub in 0..pq.m { + table.resize(pq.m() * pq.ksub(), 0.0); + for sub in 0..pq.m() { let range = pq.chunk_range(sub); let chunk_dim = range.len(); - let pq_base = range.start * pq.ksub; - for code in 0..pq.ksub { + let pq_base = range.start * pq.ksub(); + for code in 0..pq.ksub() { let pq_offset = pq_base + code * chunk_dim; let mut inner_product = 0.0f64; for dimension in 0..chunk_dim { - let pq_value = f64::from(pq.centroids[pq_offset + dimension]); + let pq_value = f64::from(pq.centroids()[pq_offset + dimension]); inner_product += f64::from(coarse_centroid[range.start + dimension]) * pq_value; } - let offset = sub * pq.ksub + code; + let offset = sub * pq.ksub() + code; table[offset] = pq_norms[offset] + 2.0 * inner_product; } } } fn fill_stable_ephemeral_query_table(query: &[f32], pq: &ProductQuantizer, table: &mut Vec) { - table.resize(pq.m * pq.ksub, 0.0); - for sub in 0..pq.m { + table.resize(pq.m() * pq.ksub(), 0.0); + for sub in 0..pq.m() { let range = pq.chunk_range(sub); let chunk_dim = range.len(); - let pq_base = range.start * pq.ksub; - for code in 0..pq.ksub { + let pq_base = range.start * pq.ksub(); + for code in 0..pq.ksub() { let pq_offset = pq_base + code * chunk_dim; let mut inner_product = 0.0f64; for dimension in 0..chunk_dim { inner_product += f64::from(query[range.start + dimension]) - * f64::from(pq.centroids[pq_offset + dimension]); + * f64::from(pq.centroids()[pq_offset + dimension]); } - table[sub * pq.ksub + code] = inner_product; + table[sub * pq.ksub() + code] = inner_product; } } } @@ -1123,16 +1132,16 @@ fn combine_stable_ephemeral_tables( pq: &ProductQuantizer, sim_table: &mut Vec, ) { - sim_table.resize(pq.m * pq.ksub, 0.0); - for sub in 0..pq.m { + sim_table.resize(pq.m() * pq.ksub(), 0.0); + for sub in 0..pq.m() { let range = pq.chunk_range(sub); let mut residual_norm = 0.0f64; for dimension in range { let residual = f64::from(query[dimension]) - f64::from(coarse_centroid[dimension]); residual_norm += residual * residual; } - let table_base = sub * pq.ksub; - for code in 0..pq.ksub { + let table_base = sub * pq.ksub(); + for code in 0..pq.ksub() { let offset = table_base + code; sim_table[offset] = (residual_norm + list_table[offset] - 2.0 * query_table[offset]).max(0.0) as f32; @@ -1531,7 +1540,7 @@ pub fn search_with_reader_filter( let (_, _, dis0) = lists_to_read[batch_start]; let sim_table = by_residual .then(|| reader_sim_table(reader, first_list, &q, &ip_table, use_precomputed)); - let pq_nbits = reader.pq.nbits; + let pq_nbits = reader.pq.nbits(); let transposed_codes = reader.transposed_codes; let mut scratch = ReaderScanScratch::default(); reader.for_each_streamed_list_chunk(first_list, |pq, ids, codes| { @@ -1674,7 +1683,7 @@ fn scan_reader_list( &entry.ids, ctx.m, ctx.ksub, - ctx.pq.nbits, + ctx.pq.nbits(), ctx.transposed_codes, dis0, matching_rows, @@ -2274,7 +2283,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( && by_residual && !reader.precomputed_table.is_empty() && reused_query_tables_fit_budget; - let allow_ephemeral_precomputed = reader.pq.nbits == 8 + let allow_ephemeral_precomputed = reader.pq.nbits() == 8 && metric == MetricType::L2 && by_residual && !use_precomputed @@ -2299,7 +2308,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( }; // Non-residual tables depend only on the query and PQ codebook, so every // probed list for that query can share one table. - let reuse_non_residual_tables = reader.pq.nbits == 8 + let reuse_non_residual_tables = reader.pq.nbits() == 8 && !by_residual && probe_end - probe_start > 1 && match reuse_mode { @@ -2375,7 +2384,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( }) }) .collect::>(); - let pq_nbits = reader.pq.nbits; + let pq_nbits = reader.pq.nbits(); let transposed_codes = reader.transposed_codes; // The loop is sequential across queries. Reuse one chunk-sized // distance buffer instead of retaining one per query. @@ -2471,7 +2480,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( .get(list.list_id as u32) .copied() .unwrap_or_default(), - reader.pq.nbits, + reader.pq.nbits(), reader.transposed_codes, ); } @@ -2615,7 +2624,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( &loaded_lists[position].ids, m, ksub, - reader.pq.nbits, + reader.pq.nbits(), reader.transposed_codes, dis0, matching_rows_by_list[position].as_ref(), @@ -2680,7 +2689,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( elapsed_since(total_started), nq, scanned_nprobe, - reader.pq.nbits, + reader.pq.nbits(), k, unique_lists.len(), filter.is_some(), @@ -2700,7 +2709,7 @@ fn search_batch_reader_filter_with_reuse_mode_and_observer( budget_bytes={reuse_max_bytes} tables_built={tables_built}", tables_built > 0, metric.as_str(), - reader.pq.nbits, + reader.pq.nbits(), unique_lists.len(), filter.is_some(), reuse_required_bytes, @@ -3303,7 +3312,7 @@ mod tests { let ids: Vec = (0..n as i64).collect(); let mut index = IVFPQIndex::with_nbits(d, nlist, m, 4, MetricType::L2, false); - assert_eq!(index.pq.ksub, 16); + assert_eq!(index.pq.ksub(), 16); assert_eq!(index.pq.code_size(), 4); index.train(&data, n); @@ -3365,6 +3374,34 @@ mod tests { assert!(dists_full[0] <= dists_limited[0] + 1e-6); } + #[test] + fn canonical_pq_encoding_is_preserved_by_from_trained() { + let mut trainer = IVFPQIndex::new(4, 1, 1, MetricType::L2, false); + assert!(!trainer.canonical_pq_encoding); + trainer.set_quantizer_centroids(vec![0.0; 4]); + let mut centroids = vec![100_000_016.0; 4 * 256]; + centroids[0..4].fill(100_000_008.0); + centroids[4..8].fill(100_000_000.0); + trainer.pq.set_centroids(centroids); + trainer.set_canonical_pq_encoding(true); + + let n = 7; + let data = vec![100_000_000.0; n * 4]; + let mut expected = vec![0; n]; + trainer.pq.encode_batch(&data, n, &mut expected); + assert_eq!(expected, vec![0; n]); + + let mut worker = IVFPQIndex::from_trained(&trainer); + assert!(worker.canonical_pq_encoding); + worker.add(&data, &(0..n as i64).collect::>(), n); + assert_eq!(worker.codes[0], expected); + + assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + worker.set_canonical_pq_encoding(false); + })) + .is_err()); + } + #[test] fn test_from_trained_and_merge() { let d = 16; @@ -3435,7 +3472,9 @@ mod tests { assert_invalid_merge(&base, &mismatched_centroids, "coarse centroids mismatch"); let mut mismatched_codebooks = IVFPQIndex::from_trained(&trainer); - mismatched_codebooks.pq.centroids[0] += 1.0; + let mut centroids = mismatched_codebooks.pq.centroids().to_vec(); + centroids[0] += 1.0; + mismatched_codebooks.pq.set_centroids(centroids); assert_invalid_merge(&base, &mismatched_codebooks, "PQ codebooks mismatch"); let mismatched_opq = IVFPQIndex::new(d, nlist, m, MetricType::L2, true); @@ -3837,7 +3876,7 @@ mod tests { &pq_norms, &mut actual, ); - let table_size = m * index.pq.ksub; + let table_size = m * index.pq.ksub(); assert_eq!( actual, index.precomputed_table[list_id * table_size..(list_id + 1) * table_size] diff --git a/core/src/pq.rs b/core/src/pq.rs index edc6a3e..96fd06f 100644 --- a/core/src/pq.rs +++ b/core/src/pq.rs @@ -21,6 +21,7 @@ use crate::distance::{ }; use crate::kmeans::{self, KMeansConfig}; use rayon::prelude::*; +use std::sync::OnceLock; /// Product Quantizer aligned with Faiss's ProductQuantizer. /// @@ -31,23 +32,63 @@ use rayon::prelude::*; /// /// Centroids are chunk-major. Chunk `m` starts at /// `chunk_offsets[m] * ksub`, and each of its `ksub` centroids contains -/// `chunk_offsets[m + 1] - chunk_offsets[m]` contiguous components. +/// `chunk_offsets[m + 1] - chunk_offsets[m]` contiguous components. Use +/// [`Self::set_centroids`] to replace them so derived caches stay synchronized. +/// Layout is immutable after construction so derived caches stay valid. +/// +/// ```compile_fail,E0616 +/// use paimon_vindex_core::pq::ProductQuantizer; +/// let mut pq = ProductQuantizer::new(12, 2); +/// pq.chunk_offsets = vec![0, 4, 12]; +/// ``` pub struct ProductQuantizer { - pub d: usize, - pub m: usize, - pub nbits: usize, + d: usize, + m: usize, + nbits: usize, /// Uniform chunk width for legacy formats, or the largest chunk width for /// a balanced non-uniform layout. - pub dsub: usize, - pub ksub: usize, - pub chunk_offsets: Vec, - pub centroids: Vec, + dsub: usize, + ksub: usize, + chunk_offsets: Vec, + centroids: Vec, + centroids_are_finite: bool, /// Pre-computed squared norms of each centroid: [M * ksub]. /// Avoids recomputing per query for L2 distance table. - pub centroid_norms_cache: Vec, + centroid_norms_cache: Vec, + transposed_codebook_cache: OnceLock<(Vec, usize)>, } impl ProductQuantizer { + /// Vector dimension. + pub fn d(&self) -> usize { + self.d + } + + /// Number of subquantizers. + pub fn m(&self) -> usize { + self.m + } + + /// Bits per PQ code. + pub fn nbits(&self) -> usize { + self.nbits + } + + /// Largest subvector dimension. + pub fn dsub(&self) -> usize { + self.dsub + } + + /// Centroids per subquantizer. + pub fn ksub(&self) -> usize { + self.ksub + } + + /// Contiguous subvector boundaries. + pub fn chunk_offsets(&self) -> &[usize] { + &self.chunk_offsets + } + pub fn new(d: usize, m: usize) -> Self { Self::with_nbits(d, m, 8) } @@ -125,7 +166,9 @@ impl ProductQuantizer { ksub, chunk_offsets, centroids: Vec::new(), + centroids_are_finite: true, centroid_norms_cache: Vec::new(), + transposed_codebook_cache: OnceLock::new(), } } @@ -155,6 +198,50 @@ impl ProductQuantizer { ) } + /// Return the chunk-major PQ codebook. + pub fn centroids(&self) -> &[f32] { + &self.centroids + } + + /// Replace the PQ codebook and refresh its derived norm and transpose caches. + pub fn set_centroids(&mut self, centroids: Vec) { + self.try_set_centroids(centroids) + .expect("PQ centroid norms allocation failed"); + } + + pub(crate) fn try_set_centroids( + &mut self, + centroids: Vec, + ) -> Result<(), std::collections::TryReserveError> { + assert_eq!( + centroids.len(), + self.d * self.ksub, + "PQ centroids must hold d * ksub values" + ); + let norms = self.try_compute_centroid_norms(¢roids)?; + self.centroids_are_finite = centroids.iter().all(|value| value.is_finite()); + self.centroids = centroids; + self.centroid_norms_cache = norms; + self.transposed_codebook_cache.take(); + Ok(()) + } + + /// Copy trained state while leaving the writer-local transpose lazy. + pub(crate) fn clone_without_transposed_cache(&self) -> Self { + Self { + d: self.d, + m: self.m, + nbits: self.nbits, + dsub: self.dsub, + ksub: self.ksub, + chunk_offsets: self.chunk_offsets.clone(), + centroids: self.centroids.clone(), + centroids_are_finite: self.centroids_are_finite, + centroid_norms_cache: self.centroid_norms_cache.clone(), + transposed_codebook_cache: OnceLock::new(), + } + } + /// Train the codebooks from training data. /// data: flat [n * d], n training vectors. pub fn train(&mut self, data: &[f32], n: usize) { @@ -243,23 +330,19 @@ impl ProductQuantizer { .install(train_subquantizers) }; - self.centroids = vec![0.0f32; d * ksub]; + let mut centroids = vec![0.0f32; d * ksub]; for (sub, sub_centroids) in sub_results.into_iter().enumerate() { let chunk_dim = self.chunk_dim(sub); let dst_offset = self.centroid_chunk_base(sub); - self.centroids[dst_offset..dst_offset + ksub * chunk_dim] - .copy_from_slice(&sub_centroids); + centroids[dst_offset..dst_offset + ksub * chunk_dim].copy_from_slice(&sub_centroids); } - self.rebuild_norms_cache(); + self.set_centroids(centroids); } - /// Rebuild the centroid norms cache. Called after training or loading centroids. - pub fn rebuild_norms_cache(&mut self) { - self.try_rebuild_norms_cache() - .expect("PQ centroid norms allocation failed"); - } - - pub fn try_rebuild_norms_cache(&mut self) -> Result<(), std::collections::TryReserveError> { + fn try_compute_centroid_norms( + &self, + centroids: &[f32], + ) -> Result, std::collections::TryReserveError> { let mut norms = Vec::new(); norms.try_reserve_exact(self.m * self.ksub)?; norms.resize(self.m * self.ksub, 0.0f32); @@ -268,12 +351,10 @@ impl ProductQuantizer { let c_base = self.centroid_chunk_base(sub); for j in 0..self.ksub { let c_off = c_base + j * chunk_dim; - norms[sub * self.ksub + j] = - fvec_norm_l2sqr(&self.centroids[c_off..c_off + chunk_dim]); + norms[sub * self.ksub + j] = fvec_norm_l2sqr(¢roids[c_off..c_off + chunk_dim]); } } - self.centroid_norms_cache = norms; - Ok(()) + Ok(norms) } /// Bytes per encoded vector. @@ -357,11 +438,10 @@ impl ProductQuantizer { /// Encode multiple vectors in parallel. /// - /// This is the byte-stable path: results are bit-identical to the - /// per-vector [`Self::encode`], which golden storage fixtures rely on - /// (DiskANN serializes these codes). IVF-PQ's add path uses - /// [`Self::encode_batch_blocked`] instead, which batches the same distance - /// calculation into larger SGEMM calls. + /// This is the canonical path: results are bit-identical to the + /// per-vector [`Self::encode`] on the same runtime backend, which golden + /// storage fixtures rely on (DiskANN serializes these codes). IVF-PQ's add + /// path uses [`Self::encode_batch_blocked`] instead. pub fn encode_batch(&self, data: &[f32], n: usize, codes: &mut [u8]) { let d = self.d; let cs = self.code_size(); @@ -376,19 +456,41 @@ impl ProductQuantizer { ); } - /// Blocked batch encode for the IVF-PQ add path. + /// Automatic batch encode for the IVF-PQ add path. + /// + /// The backend depends only on shape and CPU features, never on `n`: + /// - 8-bit shapes with `dsub >= 4` use the transposed direct-L2 encoder + /// when this CPU has a fast fused kernel. + /// - The same finite shape uses blocked SGEMM otherwise, avoiding billions + /// of software `fmaf` calls on x86 CPUs without FMA. + /// - Other shapes, and non-finite codebooks on the SGEMM fallback, use the + /// canonical encoder. + /// + /// Each backend is deterministic across thread counts and batch splits, + /// but their floating-point contracts differ. The transposed path treats + /// NaN distances as losing, lets infinite distances lose, returns code 0 + /// when no distance is below `f32::MAX`, and computes direct squared L2. + /// The SGEMM and canonical paths use the expanded form + /// `|q|² + |c|² - 2q·c`, including its existing NaN and large-offset + /// behavior. Canonical mode reproduces the per-vector encoder on the same + /// CPU and runtime backend; neither mode promises cross-CPU code identity. pub(crate) fn encode_batch_blocked(&self, data: &[f32], n: usize, codes: &mut [u8]) { - if self.nbits == 8 - && n >= encode_sgemm_min_rows(rayon::current_num_threads()) - && (0..self.m).all(|sub| self.chunk_dim(sub) >= 4) - && self.centroids.iter().all(|value| value.is_finite()) - { - self.encode_batch_8bit_sgemm(data, n, codes); + if self.nbits == 8 && self.ksub == 256 && (0..self.m).all(|sub| self.chunk_dim(sub) >= 4) { + if supports_transposed_pq_encoding() { + self.encode_batch_8bit_transposed(data, n, codes); + } else if self.centroids_are_finite { + self.encode_batch_8bit_sgemm(data, n, codes); + } else { + self.encode_batch(data, n, codes); + } return; } self.encode_batch(data, n, codes); } + /// Blocked SGEMM fallback for CPUs without a fast transposed kernel. + /// Unlike the former main path, this runs for every `n`; block sizes only + /// divide work and never select another encoder. fn encode_batch_8bit_sgemm(&self, data: &[f32], n: usize, codes: &mut [u8]) { let d = self.d; let m = self.m; @@ -396,6 +498,8 @@ impl ProductQuantizer { let cs = self.code_size(); debug_assert_eq!(cs, m); debug_assert_eq!(ksub, 256); + debug_assert!((0..m).all(|sub| self.chunk_dim(sub) >= 4)); + debug_assert!(self.centroids_are_finite); let max_dsub = (0..m).map(|sub| self.chunk_dim(sub)).max().unwrap_or(0); let computed_norms = self @@ -455,6 +559,73 @@ impl ProductQuantizer { ); } + /// Transposed codebook: per sub a `[dsub][ksub]` block at a uniform + /// stride of `max_dsub * ksub`, so sub lookup stays O(1) for balanced + /// non-uniform chunks. Returns the table and the per-sub stride. + fn build_transposed_codebook(&self) -> (Vec, usize) { + let m = self.m; + let ksub = self.ksub; + let max_dsub = (0..m).map(|sub| self.chunk_dim(sub)).max().unwrap_or(0); + let sub_stride = max_dsub + .checked_mul(ksub) + .expect("transposed codebook stride overflows usize"); + let mut transposed = vec![0.0f32; m * sub_stride]; + for sub in 0..m { + let dsub = self.chunk_dim(sub); + let c_base = self.centroid_chunk_base(sub); + let dst = &mut transposed[sub * sub_stride..sub * sub_stride + dsub * ksub]; + for j in 0..ksub { + for k in 0..dsub { + dst[k * ksub + j] = self.centroids[c_base + j * dsub + k]; + } + } + } + (transposed, sub_stride) + } + + /// Transposed-codebook encode for nbits=8. Rows are split into blocks + /// only for parallelism; each row is encoded independently. + fn encode_batch_8bit_transposed(&self, data: &[f32], n: usize, codes: &mut [u8]) { + let d = self.d; + let m = self.m; + let ksub = self.ksub; + let cs = self.code_size(); + debug_assert_eq!(cs, m); + + let (transposed, sub_stride) = self + .transposed_codebook_cache + .get_or_init(|| self.build_transposed_codebook()); + let sub_stride = *sub_stride; + let kernels: Vec = (0..m) + .map(|sub| score_argmin_kernel(self.chunk_dim(sub), ksub)) + .collect(); + + let block_rows = encode_block_rows(n, rayon::current_num_threads()); + codes[..n * cs] + .par_chunks_mut(block_rows * cs) + .enumerate() + .for_each_init( + || vec![0.0f32; ksub], + |scores, (block_idx, block_codes)| { + let row0 = block_idx * block_rows; + let rows = block_rows.min(n - row0); + let block_data = &data[row0 * d..(row0 + rows) * d]; + + for r in 0..rows { + let row = &block_data[r * d..(r + 1) * d]; + for sub in 0..m { + let range = self.chunk_range(sub); + let dsub = range.len(); + let q = &row[range]; + let t = &transposed[sub * sub_stride..sub * sub_stride + dsub * ksub]; + block_codes[r * cs + sub] = + score_argmin(kernels[sub], q, t, ksub, scores); + } + } + }, + ); + } + /// Decode PQ codes back to an approximate vector. pub fn decode(&self, codes: &[u8], x: &mut [f32]) { for sub in 0..self.m { @@ -612,20 +783,316 @@ fn argmin_code(distances: &[f32]) -> u8 { best as u8 } -/// Row block for the batched SGEMM encode path. -const MAX_ENCODE_BLOCK_ROWS: usize = 512; -const MIN_ENCODE_BLOCK_ROWS: usize = 4; -const ENCODE_SGEMM_MIN_ROWS: usize = 32; +#[cfg(target_arch = "x86_64")] +pub(crate) fn supports_transposed_pq_encoding() -> bool { + is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") +} -fn encode_sgemm_min_rows(workers: usize) -> usize { - ENCODE_SGEMM_MIN_ROWS.max(workers.max(1) * MIN_ENCODE_BLOCK_ROWS) +#[cfg(target_arch = "aarch64")] +pub(crate) fn supports_transposed_pq_encoding() -> bool { + true } +#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] +pub(crate) fn supports_transposed_pq_encoding() -> bool { + false +} + +/// Row block for the automatic batch-encode paths. 512 rows keeps scratch +/// buffers hot and yields enough blocks to occupy the worker pool. Blocks +/// only divide work; they never select another encoder. +const MAX_ENCODE_BLOCK_ROWS: usize = 512; + fn encode_block_rows(rows: usize, workers: usize) -> usize { rows.div_ceil(workers.max(1)) .clamp(1, MAX_ENCODE_BLOCK_ROWS) } +/// Squared-L2 argmin kernel over a transposed sub-codebook `t` laid out as +/// `[dsub][ksub]` (stride-1 over `j`). `scores` is a reusable `ksub`-sized +/// scratch buffer that some kernels do not touch. +type ScoreArgminKernel = fn(&[f32], &[f32], usize, &mut [f32]) -> u8; + +/// Pick the kernel for one sub-quantizer. The choice depends only on `dsub` +/// and CPU features, never on the batch size, which is what makes +/// [`ProductQuantizer::encode_batch_blocked`] batch invariant. Every kernel +/// returned here is byte-identical to [`score_argmin_scalar`]. +fn score_argmin_kernel(dsub: usize, ksub: usize) -> ScoreArgminKernel { + #[cfg(target_arch = "aarch64")] + { + if dsub == 4 && ksub.is_multiple_of(4) { + return score_argmin_neon_d4_entry; + } + if dsub >= 4 && ksub.is_multiple_of(16) { + return score_argmin_neon_generic_entry; + } + } + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + if dsub == 4 && ksub.is_multiple_of(8) { + return score_argmin_avx2_d4_entry; + } + return score_argmin_avx2_generic_entry; + } + } + let _ = (dsub, ksub); + score_argmin_scalar +} + +/// Run `kernel` after checking the slice contract shared by all kernels. +#[inline] +fn score_argmin( + kernel: ScoreArgminKernel, + q: &[f32], + t: &[f32], + ksub: usize, + scores: &mut [f32], +) -> u8 { + assert_eq!(q.len().checked_mul(ksub), Some(t.len())); + assert!(scores.len() >= ksub); + kernel(q, t, ksub, scores) +} + +/// Direct squared-difference accumulation, one `mul_add` per dimension so +/// every kernel rounds once per dimension in the same order. +#[inline(always)] +fn score_argmin_accumulate(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32]) -> u8 { + let scores = &mut scores[..ksub]; + scores.fill(0.0); + for (k, &qv) in q.iter().enumerate() { + let tk = &t[k * ksub..(k + 1) * ksub]; + for (s, &tv) in scores.iter_mut().zip(tk) { + let diff = qv - tv; + *s = diff.mul_add(diff, *s); + } + } + argmin_code(scores) +} + +/// Scalar oracle for every SIMD kernel. +fn score_argmin_scalar(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32]) -> u8 { + score_argmin_accumulate(q, t, ksub, scores) +} + +/// Generic AVX2+FMA kernel for any `dsub`: the same loop as the scalar +/// oracle compiled with the target features enabled so the stride-1 inner +/// loop vectorizes with `vfmadd`. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2", enable = "fma")] +unsafe fn score_argmin_avx2_generic(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32]) -> u8 { + score_argmin_accumulate(q, t, ksub, scores) +} + +#[cfg(target_arch = "x86_64")] +fn score_argmin_avx2_generic_entry(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32]) -> u8 { + // SAFETY: only handed out by `score_argmin_kernel` after AVX2 and FMA were + // detected on this CPU. + unsafe { score_argmin_avx2_generic(q, t, ksub, scores) } +} + +#[cfg(target_arch = "x86_64")] +fn score_argmin_avx2_d4_entry(q: &[f32], t: &[f32], ksub: usize, _scores: &mut [f32]) -> u8 { + debug_assert_eq!(q.len(), 4); + debug_assert!(ksub.is_multiple_of(8)); + // SAFETY: only handed out by `score_argmin_kernel` after AVX2 and FMA were + // detected on this CPU; slice bounds are checked by `score_argmin`. + unsafe { score_argmin_avx2_d4(q, t, ksub) } +} + +#[cfg(target_arch = "aarch64")] +fn score_argmin_neon_generic_entry(q: &[f32], t: &[f32], ksub: usize, _scores: &mut [f32]) -> u8 { + debug_assert!(q.len() >= 4); + debug_assert!(ksub.is_multiple_of(16)); + // SAFETY: NEON is baseline on aarch64; slice bounds are checked by + // `score_argmin`. + unsafe { score_argmin_neon_generic(q, t, ksub) } +} + +#[cfg(target_arch = "aarch64")] +fn score_argmin_neon_d4_entry(q: &[f32], t: &[f32], ksub: usize, _scores: &mut [f32]) -> u8 { + debug_assert_eq!(q.len(), 4); + debug_assert!(ksub.is_multiple_of(4)); + // SAFETY: NEON is baseline on aarch64; slice bounds are checked by + // `score_argmin`. + unsafe { score_argmin_neon_d4(q, t, ksub) } +} + +/// Generic NEON kernel: score 16 centroids together, then update the running +/// SIMD argmin without materializing the score table. +#[cfg(target_arch = "aarch64")] +unsafe fn score_argmin_neon_generic(q: &[f32], t: &[f32], ksub: usize) -> u8 { + use std::arch::aarch64::*; + + unsafe { + let mut min_val = vdupq_n_f32(f32::MAX); + let mut min_idx = vdupq_n_u32(0); + let lane0: [u32; 4] = [0, 1, 2, 3]; + let mut cur_idx = vld1q_u32(lane0.as_ptr()); + let four = vdupq_n_u32(4); + + for j in (0..ksub).step_by(16) { + let q0 = vdupq_n_f32(q[0]); + let t0 = t.as_ptr().add(j); + let d0 = vsubq_f32(q0, vld1q_f32(t0)); + let d1 = vsubq_f32(q0, vld1q_f32(t0.add(4))); + let d2 = vsubq_f32(q0, vld1q_f32(t0.add(8))); + let d3 = vsubq_f32(q0, vld1q_f32(t0.add(12))); + let mut s0 = vmulq_f32(d0, d0); + let mut s1 = vmulq_f32(d1, d1); + let mut s2 = vmulq_f32(d2, d2); + let mut s3 = vmulq_f32(d3, d3); + + for (k, &qv) in q.iter().enumerate().skip(1) { + let qk = vdupq_n_f32(qv); + let tk = t.as_ptr().add(k * ksub + j); + let d0 = vsubq_f32(qk, vld1q_f32(tk)); + let d1 = vsubq_f32(qk, vld1q_f32(tk.add(4))); + let d2 = vsubq_f32(qk, vld1q_f32(tk.add(8))); + let d3 = vsubq_f32(qk, vld1q_f32(tk.add(12))); + s0 = vfmaq_f32(s0, d0, d0); + s1 = vfmaq_f32(s1, d1, d1); + s2 = vfmaq_f32(s2, d2, d2); + s3 = vfmaq_f32(s3, d3, d3); + } + + for score in [s0, s1, s2, s3] { + let mask = vcltq_f32(score, min_val); + min_val = vbslq_f32(mask, score, min_val); + min_idx = vbslq_u32(mask, cur_idx, min_idx); + cur_idx = vaddq_u32(cur_idx, four); + } + } + + let mut vals = [0.0f32; 4]; + let mut idxs = [0u32; 4]; + vst1q_f32(vals.as_mut_ptr(), min_val); + vst1q_u32(idxs.as_mut_ptr(), min_idx); + let mut best = idxs[0]; + let mut best_val = vals[0]; + for lane in 1..4 { + if vals[lane] < best_val || (vals[lane] == best_val && idxs[lane] < best) { + best_val = vals[lane]; + best = idxs[lane]; + } + } + best as u8 + } +} + +/// dsub=4 NEON kernel: 4 broadcast-FMA rows, SIMD min+index tracking, +/// horizontal reduce with smallest-index tie-break. +#[cfg(target_arch = "aarch64")] +#[inline] +unsafe fn score_argmin_neon_d4(q: &[f32], t: &[f32], ksub: usize) -> u8 { + use std::arch::aarch64::*; + + let t0 = t.as_ptr(); + let t1 = unsafe { t0.add(ksub) }; + let t2 = unsafe { t0.add(2 * ksub) }; + let t3 = unsafe { t0.add(3 * ksub) }; + + unsafe { + let q0 = vdupq_n_f32(q[0]); + let q1 = vdupq_n_f32(q[1]); + let q2 = vdupq_n_f32(q[2]); + let q3 = vdupq_n_f32(q[3]); + + let mut min_val = vdupq_n_f32(f32::MAX); + let mut min_idx = vdupq_n_u32(0); + let lane0: [u32; 4] = [0, 1, 2, 3]; + let mut cur_idx = vld1q_u32(lane0.as_ptr()); + let step = vdupq_n_u32(4); + + for j in (0..ksub).step_by(4) { + let d0 = vsubq_f32(q0, vld1q_f32(t0.add(j))); + let d1 = vsubq_f32(q1, vld1q_f32(t1.add(j))); + let d2 = vsubq_f32(q2, vld1q_f32(t2.add(j))); + let d3 = vsubq_f32(q3, vld1q_f32(t3.add(j))); + let mut s = vmulq_f32(d0, d0); + s = vfmaq_f32(s, d1, d1); + s = vfmaq_f32(s, d2, d2); + s = vfmaq_f32(s, d3, d3); + + // Strictly-smaller keeps the earliest index on equal scores. + let mask = vcltq_f32(s, min_val); + min_val = vbslq_f32(mask, s, min_val); + min_idx = vbslq_u32(mask, cur_idx, min_idx); + cur_idx = vaddq_u32(cur_idx, step); + } + + let mut vals = [0.0f32; 4]; + let mut idxs = [0u32; 4]; + vst1q_f32(vals.as_mut_ptr(), min_val); + vst1q_u32(idxs.as_mut_ptr(), min_idx); + let mut best = idxs[0]; + let mut best_val = vals[0]; + for l in 1..4 { + if vals[l] < best_val || (vals[l] == best_val && idxs[l] < best) { + best_val = vals[l]; + best = idxs[l]; + } + } + best as u8 + } +} + +/// dsub=4 AVX2 kernel: mirrors the NEON version 8-wide. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2", enable = "fma")] +unsafe fn score_argmin_avx2_d4(q: &[f32], t: &[f32], ksub: usize) -> u8 { + use std::arch::x86_64::*; + + let t0 = t.as_ptr(); + let t1 = unsafe { t0.add(ksub) }; + let t2 = unsafe { t0.add(2 * ksub) }; + let t3 = unsafe { t0.add(3 * ksub) }; + + unsafe { + let q0 = _mm256_set1_ps(q[0]); + let q1 = _mm256_set1_ps(q[1]); + let q2 = _mm256_set1_ps(q[2]); + let q3 = _mm256_set1_ps(q[3]); + + let mut min_val = _mm256_set1_ps(f32::MAX); + let mut min_idx = _mm256_setzero_si256(); + let mut cur_idx = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); + let step = _mm256_set1_epi32(8); + + for j in (0..ksub).step_by(8) { + let d0 = _mm256_sub_ps(q0, _mm256_loadu_ps(t0.add(j))); + let d1 = _mm256_sub_ps(q1, _mm256_loadu_ps(t1.add(j))); + let d2 = _mm256_sub_ps(q2, _mm256_loadu_ps(t2.add(j))); + let d3 = _mm256_sub_ps(q3, _mm256_loadu_ps(t3.add(j))); + let mut s = _mm256_mul_ps(d0, d0); + s = _mm256_fmadd_ps(d1, d1, s); + s = _mm256_fmadd_ps(d2, d2, s); + s = _mm256_fmadd_ps(d3, d3, s); + + // Strictly-smaller keeps the earliest index on equal scores. + let mask = _mm256_cmp_ps::<_CMP_LT_OQ>(s, min_val); + min_val = _mm256_blendv_ps(min_val, s, mask); + min_idx = _mm256_blendv_epi8(min_idx, cur_idx, _mm256_castps_si256(mask)); + cur_idx = _mm256_add_epi32(cur_idx, step); + } + + let mut vals = [0.0f32; 8]; + let mut idxs = [0i32; 8]; + _mm256_storeu_ps(vals.as_mut_ptr(), min_val); + _mm256_storeu_si256(idxs.as_mut_ptr().cast(), min_idx); + let mut best = idxs[0] as u32; + let mut best_val = vals[0]; + for l in 1..8 { + let idx = idxs[l] as u32; + if vals[l] < best_val || (vals[l] == best_val && idx < best) { + best_val = vals[l]; + best = idx; + } + } + best as u8 + } +} + #[cfg(test)] mod tests { use super::*; @@ -729,11 +1196,13 @@ mod tests { pq.train(&data, n); let cs = pq.code_size(); // m/2 = 4 - let mut codes = vec![0u8; n * cs]; - pq.encode_batch(&data, n, &mut codes); + let mut canonical = vec![0u8; n * cs]; + pq.encode_batch(&data, n, &mut canonical); + let mut blocked = vec![0u8; n * cs]; + pq.encode_batch_blocked(&data, n, &mut blocked); - // Verify codes are non-trivial (not all zeros) - assert!(codes.iter().any(|&b| b != 0)); + assert_eq!(blocked, canonical); + assert!(blocked.iter().any(|&b| b != 0)); } #[test] @@ -764,116 +1233,444 @@ mod tests { assert!((table_distance - decoded_distance).abs() < 1e-4); } - /// Reference per-vector encode used to pin the blocked batch path. - fn encode_per_vector(pq: &ProductQuantizer, data: &[f32], n: usize) -> Vec { + /// Scalar direct-L2 oracle for the transposed path: same transposed + /// table, always the scalar kernel, no threads. + fn encode_scalar_oracle(pq: &ProductQuantizer, data: &[f32], n: usize) -> Vec { let cs = pq.code_size(); + assert_eq!(cs, pq.m); + let (transposed, sub_stride) = pq.build_transposed_codebook(); + let mut scores = vec![0.0f32; pq.ksub]; let mut codes = vec![0u8; n * cs]; - for i in 0..n { - pq.encode( - &data[i * pq.d..(i + 1) * pq.d], - &mut codes[i * cs..(i + 1) * cs], - ); + for r in 0..n { + let row = &data[r * pq.d..(r + 1) * pq.d]; + for sub in 0..pq.m { + let range = pq.chunk_range(sub); + let dsub = range.len(); + let t = &transposed[sub * sub_stride..sub * sub_stride + dsub * pq.ksub]; + codes[r * cs + sub] = + score_argmin(score_argmin_scalar, &row[range], t, pq.ksub, &mut scores); + } } codes } - #[test] - fn test_encode_batch_blocked_matches_per_vector() { - let d = 32; - let m = 8; // dsub = 4: hits the SIMD kernels - let mut rng = StdRng::seed_from_u64(20260820); + const BATCH_SIZES: [usize; 11] = [1, 2, 7, 31, 32, 33, 511, 512, 513, 2048, 4096]; + + fn trained_pq(d: usize, m: usize, seed: u64) -> (ProductQuantizer, StdRng) { + let mut rng = StdRng::seed_from_u64(seed); let train: Vec = (0..3000 * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); let mut pq = ProductQuantizer::new(d, m); pq.train(&train, 3000); + (pq, rng) + } - // Below, exactly at, above, and misaligned against the block size. - for n in [ - 1, - 31, - 32, - 33, - MAX_ENCODE_BLOCK_ROWS, - MAX_ENCODE_BLOCK_ROWS + 7, - 2048, - ] { - let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); - let reference = encode_per_vector(&pq, &data, n); - let mut batch = vec![0u8; n * pq.code_size()]; - pq.encode_batch_blocked(&data, n, &mut batch); - assert_eq!(batch, reference, "n={n}"); + #[test] + fn test_encode_batch_8bit_transposed_matches_scalar_oracle() { + // dsub = 4 (d4 kernel), 5 and 8 (generic kernel). + for (d, m) in [(32, 8), (25, 5), (40, 5)] { + let (pq, mut rng) = trained_pq(d, m, 20260820 + d as u64); + for n in BATCH_SIZES { + let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let oracle = encode_scalar_oracle(&pq, &data, n); + let mut batch = vec![0u8; n * pq.code_size()]; + pq.encode_batch_8bit_transposed(&data, n, &mut batch); + assert_eq!(batch, oracle, "d={d} m={m} n={n}"); + } } } #[test] - fn test_encode_batch_blocked_non_uniform_chunks() { + fn test_centroids_finite_cache_updates_and_clones() { + let mut pq = ProductQuantizer::new(4, 1); + for value in [0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY, 1.0] { + let mut centroids = vec![0.0; 4 * 256]; + centroids[4 * 256 - 1] = value; + pq.set_centroids(centroids); + assert_eq!(pq.centroids_are_finite, value.is_finite()); + let cloned = pq.clone_without_transposed_cache(); + assert_eq!(cloned.centroids_are_finite, value.is_finite()); + if !supports_transposed_pq_encoding() { + let data = [1.0; 4]; + let mut expected = [0]; + if value.is_finite() { + pq.encode_batch_8bit_sgemm(&data, 1, &mut expected); + } else { + pq.encode_batch(&data, 1, &mut expected); + } + for encoder in [&pq, &cloned] { + let mut actual = [0]; + encoder.encode_batch_blocked(&data, 1, &mut actual); + assert_eq!(actual, expected); + } + } + } + } + + #[test] + fn test_transposed_codebook_cache_reuses_and_invalidates() { + let (mut pq, mut rng) = trained_pq(32, 8, 20260908); + let data: Vec = (0..32).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let mut expected = vec![0u8; pq.code_size()]; + + assert!(pq.transposed_codebook_cache.get().is_none()); + pq.encode_batch_8bit_transposed(&data, 1, &mut expected); + let cached = pq.transposed_codebook_cache.get().unwrap().0.as_ptr(); + + let mut actual = vec![0u8; pq.code_size()]; + pq.encode_batch_8bit_transposed(&data, 1, &mut actual); + assert_eq!(actual, expected); + assert_eq!( + pq.transposed_codebook_cache.get().unwrap().0.as_ptr(), + cached + ); + + let centroids = pq.centroids().to_vec(); + pq.set_centroids(centroids); + assert!(pq.transposed_codebook_cache.get().is_none()); + pq.encode_batch_8bit_transposed(&data, 1, &mut actual); + assert_eq!(actual, expected); + assert!(pq.transposed_codebook_cache.get().is_some()); + } + + #[test] + fn test_encode_batch_8bit_transposed_non_uniform_chunks() { let d = 13; let m = 3; let mut rng = StdRng::seed_from_u64(20260821); let train: Vec = (0..2000 * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); let mut pq = ProductQuantizer::with_nbits_balanced(d, m, 8); pq.train(&train, 2000); + // Chunks of 5, 4 and 4: generic kernel next to the d4 kernel. assert_eq!(pq.chunk_offsets, vec![0, 5, 9, 13]); let n = MAX_ENCODE_BLOCK_ROWS + 13; let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); - let reference = encode_per_vector(&pq, &data, n); + let oracle = encode_scalar_oracle(&pq, &data, n); let mut batch = vec![0u8; n * pq.code_size()]; - pq.encode_batch_blocked(&data, n, &mut batch); - assert_eq!(batch, reference); + pq.encode_batch_8bit_transposed(&data, n, &mut batch); + assert_eq!(batch, oracle); } #[test] - fn test_encode_batch_blocked_matches_canonical_large_offset() { + fn test_encode_batch_8bit_transposed_large_offset_is_correct() { + // The expanded form cancels here; the direct form must still pick the + // exact centroid, for every batch size. let mut pq = ProductQuantizer::new(4, 1); - pq.centroids = vec![100_000_016.0; pq.d * pq.ksub]; - pq.centroids[0..4].fill(100_000_008.0); - pq.centroids[4..8].fill(100_000_000.0); - - let data = vec![100_000_000.0; 32 * pq.d]; - let mut canonical = vec![0; 32]; - pq.encode_batch(&data, 32, &mut canonical); - let mut blocked = vec![0; 32]; - pq.encode_batch_blocked(&data, 32, &mut blocked); + let mut centroids = vec![100_000_016.0; pq.d * pq.ksub]; + centroids[0..4].fill(100_000_008.0); + centroids[4..8].fill(100_000_000.0); + pq.set_centroids(centroids); + + for n in BATCH_SIZES { + let data = vec![100_000_000.0; n * pq.d]; + let mut codes = vec![0; n]; + pq.encode_batch_8bit_transposed(&data, n, &mut codes); + assert!(codes.iter().all(|&code| code == 1), "n={n}"); + } + } - assert_eq!(blocked, canonical); + #[test] + fn test_encode_batch_8bit_transposed_is_batch_invariant() { + let mut pq = ProductQuantizer::new(4, 1); + let mut centroids = vec![100.0; pq.d * pq.ksub]; + centroids[0..4].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994, -0.31766486]); + centroids[4..8].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994, -0.31766483]); + pq.set_centroids(centroids); + + let max_n = *BATCH_SIZES.last().unwrap(); + let mut largest = vec![0; max_n]; + pq.encode_batch_8bit_transposed(&vec![0.0; max_n * 4], max_n, &mut largest); + for n in BATCH_SIZES { + let mut codes = vec![0; n]; + pq.encode_batch_8bit_transposed(&vec![0.0; n * 4], n, &mut codes); + assert_eq!(codes.as_slice(), &largest[..n], "n={n}"); + } } #[test] - fn test_encode_batch_blocked_is_batch_invariant() { + fn test_encode_batch_8bit_transposed_non_finite_semantics_are_batch_invariant() { let mut pq = ProductQuantizer::new(4, 1); - pq.centroids = vec![100.0; pq.d * pq.ksub]; - pq.centroids[0..4].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994, -0.31766486]); - pq.centroids[4..8].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994, -0.31766483]); + let mut centroids = vec![1.0; pq.d * pq.ksub]; + centroids[0..4].fill(f32::NAN); + centroids[4..8].fill(0.0); + centroids[8..12].fill(f32::INFINITY); + centroids[12..16].fill(f32::NEG_INFINITY); + pq.set_centroids(centroids); + + for n in BATCH_SIZES { + let mut codes = vec![0; n]; + pq.encode_batch_8bit_transposed(&vec![0.0; n * pq.d], n, &mut codes); + assert!(codes.iter().all(|&code| code == 1), "n={n}"); + } - let mut small = vec![0; 31]; - pq.encode_batch_blocked(&[0.0; 31 * 4], 31, &mut small); - let mut large = vec![0; 32]; - pq.encode_batch_blocked(&[0.0; 32 * 4], 32, &mut large); + pq.set_centroids(vec![f32::NAN; pq.d * pq.ksub]); + for n in BATCH_SIZES { + let mut codes = vec![u8::MAX; n]; + pq.encode_batch_8bit_transposed(&vec![0.0; n * pq.d], n, &mut codes); + assert!(codes.iter().all(|&code| code == 0), "n={n}"); + } + } - assert_eq!(small.as_slice(), &large[..31]); + #[test] + fn test_encode_batch_blocked_small_dsub_uses_canonical_path() { + for dsub in [1, 2, 3] { + let m = 4; + let d = m * dsub; + let (pq, mut rng) = trained_pq(d, m, 20260903 + dsub as u64); + for n in [1, 33, 513] { + let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let mut canonical = vec![0u8; n * m]; + pq.encode_batch(&data, n, &mut canonical); + let mut blocked = vec![0u8; n * m]; + pq.encode_batch_blocked(&data, n, &mut blocked); + assert_eq!(blocked, canonical, "dsub={dsub} n={n}"); + } + } } #[test] - fn test_encode_batch_blocked_nan_centroid_is_batch_invariant() { - let mut pq = ProductQuantizer::new(4, 1); - pq.centroids = vec![1.0; pq.d * pq.ksub]; - pq.centroids[0] = f32::NAN; - pq.centroids[4..8].fill(0.0); + fn test_encode_batch_8bit_transposed_is_thread_and_split_invariant() { + let d = 32; + let m = 8; + let (pq, mut rng) = trained_pq(d, m, 20260904); + let n = 3 * MAX_ENCODE_BLOCK_ROWS + 5; + let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let mut whole = vec![0u8; n * m]; + pq.encode_batch_8bit_transposed(&data, n, &mut whole); - let mut small = vec![0; 1]; - pq.encode_batch_blocked(&[0.0; 4], 1, &mut small); - let mut large = vec![0; 32]; - pq.encode_batch_blocked(&[0.0; 32 * 4], 32, &mut large); + for threads in [1, 2, 11, 16] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let mut codes = vec![0u8; n * m]; + pool.install(|| pq.encode_batch_8bit_transposed(&data, n, &mut codes)); + assert_eq!(codes, whole, "threads={threads}"); + } - assert_eq!(small[0], large[0]); + for batch in [1, 7, MAX_ENCODE_BLOCK_ROWS] { + let mut codes = vec![0u8; n * m]; + for start in (0..n).step_by(batch) { + let rows = batch.min(n - start); + pq.encode_batch_8bit_transposed( + &data[start * d..(start + rows) * d], + rows, + &mut codes[start * m..(start + rows) * m], + ); + } + assert_eq!(codes, whole, "batch={batch}"); + } } #[test] - fn test_encode_block_rows_avoids_tiny_sgemm_blocks() { + fn test_encode_batch_8bit_sgemm_is_thread_and_split_invariant() { + let d = 32; + let m = 8; + let (pq, mut rng) = trained_pq(d, m, 20260905); + let n = 3 * MAX_ENCODE_BLOCK_ROWS + 5; + let data: Vec = (0..n * d).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let mut whole = vec![0u8; n * m]; + pq.encode_batch_8bit_sgemm(&data, n, &mut whole); + + for threads in [1, 2, 11, 16] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let mut codes = vec![0u8; n * m]; + pool.install(|| pq.encode_batch_8bit_sgemm(&data, n, &mut codes)); + assert_eq!(codes, whole, "threads={threads}"); + } + + for batch in [1, 7, MAX_ENCODE_BLOCK_ROWS] { + let mut codes = vec![0u8; n * m]; + for start in (0..n).step_by(batch) { + let rows = batch.min(n - start); + pq.encode_batch_8bit_sgemm( + &data[start * d..(start + rows) * d], + rows, + &mut codes[start * m..(start + rows) * m], + ); + } + assert_eq!(codes, whole, "batch={batch}"); + } + } + + #[test] + fn test_encode_block_rows_clamps() { assert_eq!(encode_block_rows(2730, 12), 228); assert_eq!(encode_block_rows(32768, 12), MAX_ENCODE_BLOCK_ROWS); - assert_eq!(encode_sgemm_min_rows(8), 32); - assert_eq!(encode_sgemm_min_rows(32), 128); + assert_eq!(encode_block_rows(5, 16), 1); + assert_eq!(encode_block_rows(0, 16), 1); + assert_eq!(encode_block_rows(100, 0), 100); + } + + /// Every kernel `score_argmin_kernel` can hand out on this machine, plus + /// the scalar oracle. + fn kernels_under_test(dsub: usize, ksub: usize) -> Vec<(&'static str, ScoreArgminKernel)> { + let mut kernels: Vec<(&'static str, ScoreArgminKernel)> = vec![ + ("scalar", score_argmin_scalar), + ("dispatched", score_argmin_kernel(dsub, ksub)), + ]; + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + kernels.push(("avx2_generic", score_argmin_avx2_generic_entry)); + if dsub == 4 && ksub.is_multiple_of(8) { + kernels.push(("avx2_d4", score_argmin_avx2_d4_entry)); + } + } + } + #[cfg(target_arch = "aarch64")] + { + if dsub >= 4 && ksub.is_multiple_of(16) { + kernels.push(("neon_generic", score_argmin_neon_generic_entry)); + } + if dsub == 4 && ksub.is_multiple_of(4) { + kernels.push(("neon_d4", score_argmin_neon_d4_entry)); + } + } + kernels + } + + #[test] + fn test_score_argmin_kernels_match_scalar_oracle() { + let ksub = 256; + let mut rng = StdRng::seed_from_u64(20260905); + for dsub in [4, 5, 8, 12, 16] { + let mut t: Vec = (0..dsub * ksub) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(); + // Inject exact duplicates so some queries hit exact ties. + for j in (0..ksub).step_by(37) { + for k in 0..dsub { + t[k * ksub + (j + 1) % ksub] = t[k * ksub + j]; + } + } + let mut scores = vec![0.0f32; ksub]; + for _ in 0..500 { + let q: Vec = (0..dsub).map(|_| rng.gen_range(-1.0f32..1.0)).collect(); + let expected = score_argmin(score_argmin_scalar, &q, &t, ksub, &mut scores); + for (name, kernel) in kernels_under_test(dsub, ksub) { + let got = score_argmin(kernel, &q, &t, ksub, &mut scores); + assert_eq!(got, expected, "kernel {name} dsub {dsub}"); + } + } + // Queries equal to a centroid resolve to its smallest duplicate. + for j in (0..ksub).step_by(37) { + let q: Vec = (0..dsub).map(|k| t[k * ksub + j]).collect(); + for (name, kernel) in kernels_under_test(dsub, ksub) { + assert_eq!( + score_argmin(kernel, &q, &t, ksub, &mut scores) as usize, + j, + "kernel {name} dsub {dsub} j {j}" + ); + } + } + } + } + + #[test] + fn test_score_argmin_non_finite_semantics() { + let ksub = 256; + for dsub in [4, 5] { + let q = vec![0.0f32; dsub]; + let mut t = vec![1.0f32; dsub * ksub]; + for k in 0..dsub { + t[k * ksub] = f32::NAN; + t[k * ksub + 1] = 0.0; + t[k * ksub + 2] = f32::INFINITY; + t[k * ksub + 3] = f32::NEG_INFINITY; + } + let mut scores = vec![0.0f32; ksub]; + for (name, kernel) in kernels_under_test(dsub, ksub) { + assert_eq!( + score_argmin(kernel, &q, &t, ksub, &mut scores), + 1, + "kernel {name} dsub {dsub}" + ); + } + + t.fill(f32::NAN); + for (name, kernel) in kernels_under_test(dsub, ksub) { + assert_eq!( + score_argmin(kernel, &q, &t, ksub, &mut scores), + 0, + "all NaN: kernel {name} dsub {dsub}" + ); + } + } + } + + #[test] + fn test_score_argmin_tie_prefers_smallest_index() { + let ksub = 256; + let dsub = 4; + let q = [0.25f32, -0.5, 0.75, -0.125]; + let mut scores = vec![0.0f32; ksub]; + + // All centroids identical -> every score ties -> index 0 wins. + let all_equal = vec![0.5f32; dsub * ksub]; + for (name, kernel) in kernels_under_test(dsub, ksub) { + assert_eq!( + score_argmin(kernel, &q, &all_equal, ksub, &mut scores), + 0, + "{name}" + ); + } + + // Exact ties at indices 1 and 8 (different AVX2 lanes) and at 8 and 9 + // (the same lane group): the smallest index must win each time. + for (tie_a, tie_b) in [(1usize, 8usize), (8, 9), (7, 8), (255, 1)] { + let mut t = vec![3.0f32; dsub * ksub]; + for k in 0..dsub { + t[k * ksub + tie_a] = q[k]; + t[k * ksub + tie_b] = q[k]; + } + let expected = tie_a.min(tie_b); + for (name, kernel) in kernels_under_test(dsub, ksub) { + assert_eq!( + score_argmin(kernel, &q, &t, ksub, &mut scores) as usize, + expected, + "{name} tie ({tie_a}, {tie_b})" + ); + } + } + } + + #[test] + fn test_score_argmin_scalar_matches_squared_distance() { + let q = [0.4f32, -0.7, 0.2, 0.9, -0.3]; + let ksub = 17; + let t: Vec = (0..q.len() * ksub) + .map(|i| ((i * 13 % 29) as f32 - 14.0) / 7.0) + .collect(); + let distances: Vec = (0..ksub) + .map(|j| { + q.iter() + .enumerate() + .map(|(k, value)| (value - t[k * ksub + j]).powi(2)) + .sum() + }) + .collect(); + let mut scores = vec![0.0; ksub]; + assert_eq!( + score_argmin(score_argmin_scalar, &q, &t, ksub, &mut scores), + argmin_code(&distances) + ); + } + + #[test] + #[should_panic] + fn test_score_argmin_rejects_short_transposed_codebook() { + let mut scores = vec![0.0; 256]; + score_argmin( + score_argmin_scalar, + &[0.0; 4], + &[0.0; 4 * 256 - 1], + 256, + &mut scores, + ); } #[test] @@ -907,8 +1704,7 @@ mod tests { let m = 1; let ksub = 256; let mut pq = ProductQuantizer::new(d, m); - pq.centroids = vec![0.0; ksub * d]; - + let mut centroids = vec![0.0; ksub * d]; let mut query = vec![0.0; d]; for i in 0..d { let value = if i.is_multiple_of(2) { @@ -917,9 +1713,9 @@ mod tests { -1.0e10_f32 + i as f32 }; query[i] = value; - pq.centroids[i] = value; + centroids[i] = value; } - pq.rebuild_norms_cache(); + pq.set_centroids(centroids); let mut table = vec![0.0; m * ksub]; pq.compute_distance_table(&query, MetricType::L2, &mut table); diff --git a/core/src/vamana.rs b/core/src/vamana.rs index fe037c7..f64ec30 100644 --- a/core/src/vamana.rs +++ b/core/src/vamana.rs @@ -346,8 +346,8 @@ impl VamanaGraph { params: DiskAnnBuildParams, ) -> io::Result<(Self, VamanaBuildStats)> { validate_build_inputs(vectors, count, dimension, params)?; - if pq.d != dimension - || !matches!(pq.nbits, 4 | 8) + if pq.d() != dimension + || !matches!(pq.nbits(), 4 | 8) || !pq.has_valid_layout() || pq_codes.len() != count.saturating_mul(pq.code_size()) { @@ -427,8 +427,8 @@ impl VamanaGraph { )); } if let Some((pq, pq_codes)) = pq_build { - if pq.d != dimension - || !matches!(pq.nbits, 4 | 8) + if pq.d() != dimension + || !matches!(pq.nbits(), 4 | 8) || !pq.has_valid_layout() || pq_codes.len() != count.saturating_mul(pq.code_size()) { @@ -1274,30 +1274,30 @@ impl<'a> PqBuildDistance<'a> { metric: MetricType, ) -> io::Result { let table_len = pq - .m - .checked_mul(pq.ksub) - .and_then(|value| value.checked_mul(pq.ksub)) + .m() + .checked_mul(pq.ksub()) + .and_then(|value| value.checked_mul(pq.ksub())) .ok_or_else(|| invalid_input("Vamana PQ build-distance table size overflows usize"))?; let mut centroid_distances = Vec::new(); centroid_distances .try_reserve_exact(table_len) .map_err(|_| invalid_input("Vamana PQ build-distance table allocation failed"))?; centroid_distances.resize(table_len, 0.0); - for sub in 0..pq.m { + for sub in 0..pq.m() { let chunk_dim = pq.chunk_dim(sub); let sub_base = pq.centroid_chunk_base(sub); - let table_base = sub * pq.ksub * pq.ksub; - for left in 0..pq.ksub { + let table_base = sub * pq.ksub() * pq.ksub(); + for left in 0..pq.ksub() { let left_start = sub_base + left * chunk_dim; - for right in left..pq.ksub { + for right in left..pq.ksub() { let right_start = sub_base + right * chunk_dim; let distance = fvec_distance( - &pq.centroids[left_start..left_start + chunk_dim], - &pq.centroids[right_start..right_start + chunk_dim], + &pq.centroids()[left_start..left_start + chunk_dim], + &pq.centroids()[right_start..right_start + chunk_dim], metric, ); - centroid_distances[table_base + left * pq.ksub + right] = distance; - centroid_distances[table_base + right * pq.ksub + left] = distance; + centroid_distances[table_base + left * pq.ksub() + right] = distance; + centroid_distances[table_base + right * pq.ksub() + left] = distance; } } } @@ -1310,9 +1310,9 @@ impl<'a> PqBuildDistance<'a> { Ok(Self { codes, code_size, - m: pq.m, - ksub: pq.ksub, - nbits: pq.nbits, + m: pq.m(), + ksub: pq.ksub(), + nbits: pq.nbits(), centroid_distances, }) } @@ -2498,13 +2498,15 @@ mod tests { #[test] fn vamana_pq_build_distance_matches_decoded_centroid_distance_for_4bit_codes() { let mut pq = ProductQuantizer::with_nbits(4, 2, 4); - pq.centroids = (0..pq.d * pq.ksub) - .map(|index| index as f32 * 0.125) - .collect(); + pq.set_centroids( + (0..pq.d() * pq.ksub()) + .map(|index| index as f32 * 0.125) + .collect(), + ); let codes = [0x21, 0x43]; let distance = PqBuildDistance::new(&pq, &codes, 2, MetricType::L2).unwrap(); - let mut left = vec![0.0; pq.d]; - let mut right = vec![0.0; pq.d]; + let mut left = vec![0.0; pq.d()]; + let mut right = vec![0.0; pq.d()]; pq.decode(&codes[..1], &mut left); pq.decode(&codes[1..], &mut right); @@ -2516,13 +2518,15 @@ mod tests { #[test] fn vamana_pq_build_distance_and_pruning_follow_inner_product_semantics() { let mut pq = ProductQuantizer::with_nbits(2, 1, 4); - pq.centroids = (0..pq.d * pq.ksub) - .map(|index| index as f32 * 0.25 - 1.0) - .collect(); + pq.set_centroids( + (0..pq.d() * pq.ksub()) + .map(|index| index as f32 * 0.25 - 1.0) + .collect(), + ); let codes = [0x01, 0x03]; let distance = PqBuildDistance::new(&pq, &codes, 2, MetricType::InnerProduct).unwrap(); - let mut left = vec![0.0; pq.d]; - let mut right = vec![0.0; pq.d]; + let mut left = vec![0.0; pq.d()]; + let mut right = vec![0.0; pq.d()]; pq.decode(&codes[..1], &mut left); pq.decode(&codes[1..], &mut right); diff --git a/core/tests/storage_format_fixtures.rs b/core/tests/storage_format_fixtures.rs index df3d0e4..f2af87f 100644 --- a/core/tests/storage_format_fixtures.rs +++ b/core/tests/storage_format_fixtures.rs @@ -350,8 +350,9 @@ fn build_diskann_fixture() -> Vec { ..DiskAnnBuildParams::default() }, ); - index.pq.centroids = (0..256).map(|code| code as f32 * 0.25).collect(); - index.pq.rebuild_norms_cache(); + index + .pq + .set_centroids((0..256).map(|code| code as f32 * 0.25).collect()); index.ids = vec![7]; index.vectors = vec![0.0]; write_diskann_fixture(index) @@ -375,10 +376,13 @@ fn build_diskann_compact_multipage_fixture() -> Vec { ..DiskAnnBuildParams::default() }, ); - index.pq.centroids = (0..dimension) - .flat_map(|coordinate| (0..256).map(move |code| code as f32 + coordinate as f32 * 0.001)) - .collect(); - index.pq.rebuild_norms_cache(); + index.pq.set_centroids( + (0..dimension) + .flat_map(|coordinate| { + (0..256).map(move |code| code as f32 + coordinate as f32 * 0.001) + }) + .collect(), + ); index.ids = (0..count).map(|node| 10_000 + node as i64 * 7).collect(); index.vectors = (0..count) .flat_map(|node| { @@ -423,12 +427,13 @@ fn build_diskann_interleaved_4bit_fixture() -> Vec { ..DiskAnnBuildParams::default() }, ); - index.pq.centroids = (0..dimension) - .flat_map(|coordinate| { - (0..16).map(move |code| code as f32 * 0.5 + coordinate as f32 * 0.001) - }) - .collect(); - index.pq.rebuild_norms_cache(); + index.pq.set_centroids( + (0..dimension) + .flat_map(|coordinate| { + (0..16).map(move |code| code as f32 * 0.5 + coordinate as f32 * 0.001) + }) + .collect(), + ); index.ids = (0..count).map(|node| -500 + node as i64 * 11).collect(); index.vectors = (0..count) .flat_map(|node| { @@ -454,8 +459,9 @@ fn build_diskann_raw_row_ids_fixture() -> Vec { ..DiskAnnBuildParams::default() }, ); - index.pq.centroids = (0..256).map(|code| code as f32).collect(); - index.pq.rebuild_norms_cache(); + index + .pq + .set_centroids((0..256).map(|code| code as f32).collect()); index.ids = vec![i64::MIN, 0, i64::MAX]; index.vectors = vec![0.0, 1.0, 2.0]; write_diskann_fixture(index) @@ -488,8 +494,11 @@ fn build_ivf_flat_fixture() -> Vec { fn build_ivf_pq_fixture() -> Vec { let mut index = IVFPQIndex::new(1, 2, 1, MetricType::L2, false); index.set_quantizer_centroids(vec![0.0, 10.0]); - index.pq.centroids = (0..index.pq.ksub).map(|code| code as f32 * 0.25).collect(); - index.pq.rebuild_norms_cache(); + index.pq.set_centroids( + (0..index.pq.ksub()) + .map(|code| code as f32 * 0.25) + .collect(), + ); index.ids = vec![vec![20, 10], vec![30]]; index.codes = vec![vec![1, 0], vec![0]]; @@ -501,10 +510,11 @@ fn build_ivf_pq_fixture() -> Vec { fn build_ivf_pq_4bit_fixture() -> Vec { let mut index = IVFPQIndex::with_nbits(2, 2, 2, 4, MetricType::L2, false); index.set_quantizer_centroids(vec![0.0, 0.0, 10.0, 10.0]); - index.pq.centroids = (0..index.pq.m) - .flat_map(|_| (0..index.pq.ksub).map(|code| code as f32 * 0.5)) - .collect(); - index.pq.rebuild_norms_cache(); + index.pq.set_centroids( + (0..index.pq.m()) + .flat_map(|_| (0..index.pq.ksub()).map(|code| code as f32 * 0.5)) + .collect(), + ); index.ids = vec![vec![8, 5], vec![30]]; index.codes = vec![vec![0x11, 0x00], vec![0x00]]; diff --git a/docs/api.html b/docs/api.html index b1e004d..79bedf8 100644 --- a/docs/api.html +++ b/docs/api.html @@ -50,6 +50,7 @@

Shared lifecycle

01Create a Trainer
Parse and validate options
02Submit one or more
training batches
03Finish training and
create a one-shot Writer
04Add row IDs / vectors
and write the file
05Detect file magic
and execute searches
  • Vectors are contiguous f32 values; length must equal vector_count × dimension.
  • Training data may arrive in batches. Every IVF trainer keeps a deterministic reservoir of at most max(65,536, 64 × resolved nlist) vectors. DiskANN starts from a 50,000-row cap and lowers it when necessary so the retained sample, optional cosine-normalized copy, codebook, and parallel PQ-training scratch fit diskann.memory-budget-bytes. Sampling is independent of batch boundaries.
  • The Python and Java one-shot train helpers infer dimension from the matrix and use its row count for automatic nlist. When the matrix is only a sample, pass the final corpus size as expected-vector-count. Streaming Trainer APIs require a concrete dimension before their first batch.
  • A Writer may receive production vectors in multiple batches. Row-ID count must equal vector count.
  • Readers expose metadata, single-query search, batch search, and Roaring64-filtered variants.
  • Files carry their type and resolved model sections. Callers do not pass index options again when opening a Reader.
IVF coarse assignment is approximate by default for large centroid matricesWhen dimension × nlist ≥ 1,000,000, ivf.coarse-assignment=auto uses a Vamana graph while training and adding vectors. Search still selects lists by exact centroid distance, so graph assignment can lower recall at small nprobe and does not guarantee that a vector is found by a self-query with nprobe=1. Set ivf.coarse-assignment=exact to disable the graph, preserve exact nearest-centroid assignment, and avoid graph startup cost for small non-empty batches. Empty batches never build the graph.
+
IVF-PQ encoding defaults to autoFor 8-bit PQ (ksub=256) where every subvector has at least four dimensions, ivf.pq-encoding=auto uses the transposed direct-L2 encoder on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked SGEMM for finite codebooks. Unsupported shapes, and non-finite codebooks on the SGEMM fallback, use the canonical encoder. Codes can differ across these backends; NaN and high-dynamic-range inputs also follow the selected backend's arithmetic. Set ivf.pq-encoding=canonical to reproduce ProductQuantizer::encode_batch on the same CPU and runtime backend; it is substantially slower. Neither mode promises byte-identical codes across CPU feature sets. The choice affects builds only, is not stored, and does not change the index format or search path.
diff --git a/docs/ivf-pq.html b/docs/ivf-pq.html index 8ceeb65..d39de04 100644 --- a/docs/ivf-pq.html +++ b/docs/ivf-pq.html @@ -48,6 +48,7 @@

Usage

options.put("pq.code-ratio", "0.0625"); // optional; default 0.0625 options.put("use-opq", "false"); // optional; default false // options.put("ivf.coarse-assignment", "exact"); // optional; default auto +// options.put("ivf.pq-encoding", "canonical"); // optional; default auto try (VectorIndexTraining training = VectorIndexTrainer.train(options, trainingVectors, trainingCount); @@ -72,7 +73,8 @@

Usage

Parameters

-
ParameterRequirement / defaultPurposeTuning meaning
dimensionInferred by Java/Python one-shot training; otherwise > 0Input dimension dThe inferred or explicit pq.m must divide it
nlistAuto from expected-vector-count, or explicit > 0IVF partition countControls list length, centroid cost, and automatic nprobe
pq.code-ratioDefault 0.0625; finite and in (0, 0.25]Target ratio between PQ-code bytes and raw f32-vector bytesThe closest valid divisor is selected; max-bytes-per-vector can supply the target instead
pq.mOptional expert override; > 0; d % m == 0Concrete subspace count and 8-bit code bytesTakes precedence over automatic sizing; use after representative recall measurements
use-opqAuto enables when target-recall ≥ 0.9; explicit true/false winsLearn and apply an orthogonal rotationTrades training and query work for possible recall improvement
ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
nprobeAutomatic by default; explicit expert overrideLists readAuto accounts for K, average list size, and filter selectivity
+
ParameterRequirement / defaultPurposeTuning meaning
dimensionInferred by Java/Python one-shot training; otherwise > 0Input dimension dThe inferred or explicit pq.m must divide it
nlistAuto from expected-vector-count, or explicit > 0IVF partition countControls list length, centroid cost, and automatic nprobe
pq.code-ratioDefault 0.0625; finite and in (0, 0.25]Target ratio between PQ-code bytes and raw f32-vector bytesThe closest valid divisor is selected; max-bytes-per-vector can supply the target instead
pq.mOptional expert override; > 0; d % m == 0Concrete subspace count and 8-bit code bytesTakes precedence over automatic sizing; use after representative recall measurements
use-opqAuto enables when target-recall ≥ 0.9; explicit true/false winsLearn and apply an orthogonal rotationTrades training and query work for possible recall improvement
ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
ivf.pq-encodingauto by default; optional canonicalControls build-time PQ encodingFor 8-bit PQ with every dsub >= 4, auto uses transposed direct-L2 on x86 with AVX2+FMA and on AArch64, and blocked SGEMM on other CPUs when the codebook is finite; unsupported shapes and non-finite codebooks on the SGEMM fallback use canonical. canonical reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend at substantially higher build cost
nprobeAutomatic by default; explicit expert overrideLists readAuto accounts for K, average list size, and filter selectivity
+

All PQ encoding backends produce the same index format. The accelerated automatic paths require 8-bit PQ (ksub=256) and at least four dimensions in every subvector. Eligible shapes use transposed direct squared L2 on x86 with AVX2+FMA and on AArch64; other CPUs use blocked SGEMM when the codebook is finite. Unsupported shapes and non-finite codebooks on the SGEMM fallback use canonical. SGEMM and canonical use the expanded form, so codes can differ near ties; NaN centroids and high-dynamic-range inputs can also behave differently because the expanded form can cancel. The automatic backend therefore is not byte-stable across CPU feature sets. Canonical mode reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend, but its SIMD norm reductions also do not promise byte-identical codes across CPU feature sets.

diff --git a/docs/releases.html b/docs/releases.html index 75598c0..65a2fb2 100644 --- a/docs/releases.html +++ b/docs/releases.html @@ -51,7 +51,8 @@

Upcoming: 0.5.0

The repository is currently developing the 0.5.0 line. Until an ASF vote passes and the signed source archive appears under Apache downloads, code and packages from this line are development artifacts rather than an Apache release.

IVF-SQ now pools residual training bounds, fuses residual encoding, transposes output in bounded parallel batches, and reuses query heaps with conservative L2 block pruning. Unified readers and language bindings also cache decoded partitions within the existing reader memory budget. The SIFT1M/GIST1M/GloVe benchmarks record build, native query, and recall results with reproducible commands.

IVSQ v1 files remain compatible. Existing indexes receive reader optimizations without a rebuild; retrain and rebuild to use the new quantization bounds. Set the reader memory budget to zero to disable the SQ cache; direct Rust IVFSQIndexReader::open remains uncached. See reader options and upgrade details.

-
Rust IVF API migrationVersion 0.5.0 makes quantizer_centroids private on IVFFlatIndex, IVFPQIndex, IVFSQIndex, and IVFRQIndex. Replace direct reads with quantizer_centroids() and direct assignments with set_quantizer_centroids(...). The setter validates the centroid shape, rejects replacement after vectors are added, and refreshes cached derived state. IVF variants of VectorIndexConfig also require use_approximate_coarse_assignment; set it to true for the automatic 0.5.0 behavior or false for exact nearest-centroid assignment. Option-map callers can select the same policy with ivf.coarse-assignment=auto|exact. The policy is fixed when the writer is created; direct IVF indexes do not expose a post-training policy switch. These are source-level changes; the stored index format is unchanged.
+
Rust IVF API migrationVersion 0.5.0 makes quantizer_centroids private on IVFFlatIndex, IVFPQIndex, IVFSQIndex, and IVFRQIndex. Replace direct reads with quantizer_centroids() and direct assignments with set_quantizer_centroids(...). The setter validates the centroid shape, rejects replacement after vectors are added, and refreshes cached derived state. IVF variants of VectorIndexConfig also require use_approximate_coarse_assignment; set it to true for the automatic 0.5.0 behavior or false for exact nearest-centroid assignment. VectorIndexConfig::IvfPq additionally requires canonical_pq_encoding; use false for the default automatic encoder. Option-map callers can select these policies with ivf.coarse-assignment=auto|exact and ivf.pq-encoding=auto|canonical. The policies are fixed when the writer is created; direct IVF indexes do not expose post-training policy switches. ProductQuantizer::centroids and its derived norm cache are also private in 0.5.0; replace direct codebook reads with centroids() and replacements with set_centroids(...). The setter validates the codebook shape and refreshes the norm and transposed-codebook caches. The PQ layout fields d, m, nbits, dsub, ksub, and chunk_offsets are private; use the same-named read-only methods. To change the layout, construct a new ProductQuantizer so derived caches cannot outlive their layout. These are source-level changes; the stored index format is unchanged.
+
IVF-PQ build encoding changes by defaultFor 8-bit PQ (ksub=256) where every subvector has at least four dimensions, the default automatic encoder uses transposed direct-L2 on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked SGEMM expanded-form encoding for finite codebooks. Unsupported shapes, and non-finite codebooks on the SGEMM fallback, use canonical encoding. Codes can differ across these backends at near ties; NaN centroids and high-dynamic-range finite inputs can also behave differently because the expanded form can cancel. Set ivf.pq-encoding=canonical to reproduce ProductQuantizer::encode_batch on the same CPU and runtime backend. Canonical mode does not promise byte-identical codes across CPU feature sets and is substantially slower; the index format and query path are unchanged.
diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java index 761762e..b9bc6a3 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java @@ -346,6 +346,7 @@ public void run() { private static void testReaderAndWriterApiCompile() { Map options = ivfPqOptions(2, 4); + options.put("ivf.pq-encoding", "canonical"); VectorIndexReader closedReader = VectorIndexReader.fromNativePointerForTesting(0L); closedReader.close(); closedReader.close(); diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java index a7fd445..e33692d 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java @@ -423,7 +423,14 @@ public void run() { private static void testSupportedIndexRoundtrips() { runRoundtrip("ivf_flat", ivfFlatOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); - runRoundtrip("ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 2, 8); + Map automaticPq = + ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST); + automaticPq.put("ivf.pq-encoding", "auto"); + runRoundtrip("ivf_pq", automaticPq, 2, 8); + Map canonicalPq = + ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST); + canonicalPq.put("ivf.pq-encoding", "canonical"); + runRoundtrip("ivf_pq", canonicalPq, 2, 8); runRoundtrip("ivf_rq", ivfRqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST), 0, 0); runRoundtrip( "ivf_sq", diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py index 950f714..d82490b 100644 --- a/python/tests/test_vindex.py +++ b/python/tests/test_vindex.py @@ -278,6 +278,18 @@ def test_python_ffi_roundtrips_supported_indexes(): "nlist": "4", "metric": "l2", "use-opq": "false", + "ivf.pq-encoding": "auto", + }, + 16, + ), + ( + { + "index.type": "ivf_pq", + "dimension": "16", + "nlist": "4", + "metric": "l2", + "use-opq": "false", + "ivf.pq-encoding": "canonical", }, 16, ),