From 08325aa79fdf9fd71ec19f4e9ed78fbb589cd0da Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Sun, 23 Jun 2024 09:12:07 -0700 Subject: [PATCH 1/4] baseline --- rust/lance-index/Cargo.toml | 4 +++ rust/lance-index/benches/pq_assignment.rs | 44 +++++++++++++++++++++++ rust/lance-index/src/vector/pq.rs | 5 +-- 3 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 rust/lance-index/benches/pq_assignment.rs diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index 2f88fa684f..0991ae6e00 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -76,6 +76,10 @@ harness = false name = "pq_dist_table" harness = false +[[bench]] +name = "pq_assignment" +harness = false + [[bench]] name = "hnsw" harness = false diff --git a/rust/lance-index/benches/pq_assignment.rs b/rust/lance-index/benches/pq_assignment.rs new file mode 100644 index 0000000000..bf448633f7 --- /dev/null +++ b/rust/lance-index/benches/pq_assignment.rs @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Benchmark of Building PQ code from Dense Vectors. + +use std::sync::Arc; + +use arrow_array::{types::Float32Type, FixedSizeListArray}; +use criterion::{criterion_group, criterion_main, Criterion}; +use lance_arrow::FixedSizeListArrayExt; +use lance_index::vector::pq::{ProductQuantizer, ProductQuantizerImpl}; +use lance_linalg::distance::DistanceType; +use lance_testing::datagen::generate_random_array_with_seed; + +const PQ: usize = 96; +const DIM: usize = 1536; +const TOTAL: usize = 32 * 1024; + +fn pq_transform(c: &mut Criterion) { + let codebook = Arc::new(generate_random_array_with_seed::( + 256 * DIM, + [88; 32], + )); + + let vectors = generate_random_array_with_seed::(DIM * TOTAL, [3; 32]); + let fsl = FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap(); + + for dt in [DistanceType::L2, DistanceType::Dot].iter() { + let pq = ProductQuantizerImpl::::new(PQ, 8, DIM, codebook.clone(), *dt); + + c.bench_function(format!("{},{}", dt, TOTAL).as_str(), |b| { + b.iter(|| { + let _ = pq.transform(&fsl).unwrap(); + }) + }); + } +} + +criterion_group!( + name=benches; + config = Criterion::default().significance_level(0.1).sample_size(10); + targets = pq_transform); + +criterion_main!(benches); diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index 37473aa0b7..da5200fb2e 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -14,7 +14,7 @@ use lance_arrow::*; use lance_core::{Error, Result}; use lance_linalg::distance::{dot_distance_batch, l2_distance_batch, DistanceType, Dot, L2}; use lance_linalg::kernels::argmin_value_float; -use lance_linalg::kmeans::kmeans_find_partitions; +use lance_linalg::kmeans::{compute_partitions, kmeans_find_partitions}; use lance_linalg::{distance::MetricType, MatrixView}; use rayon::prelude::*; use snafu::{location, Location}; @@ -360,12 +360,13 @@ where location: location!(), })?; + let sub_dim = dim / num_sub_vectors; let values = flatten_data .as_slice() .par_chunks(dim) .map(|vector| { vector - .chunks_exact(dim / num_sub_vectors) + .chunks_exact(sub_dim) .enumerate() .flat_map(|(sub_idx, sub_vector)| { let centroids = get_sub_vector_centroids( From f1138daa849f0e468bdcbf5dc9149d285fade638 Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Sun, 23 Jun 2024 09:17:23 -0700 Subject: [PATCH 2/4] use faster routine for pq assignment --- rust/lance-index/src/vector/pq.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index da5200fb2e..0fa659c561 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -14,7 +14,7 @@ use lance_arrow::*; use lance_core::{Error, Result}; use lance_linalg::distance::{dot_distance_batch, l2_distance_batch, DistanceType, Dot, L2}; use lance_linalg::kernels::argmin_value_float; -use lance_linalg::kmeans::{compute_partitions, kmeans_find_partitions}; +use lance_linalg::kmeans::compute_partitions; use lance_linalg::{distance::MetricType, MatrixView}; use rayon::prelude::*; use snafu::{location, Location}; @@ -376,9 +376,9 @@ where num_sub_vectors, sub_idx, ); - let parts = kmeans_find_partitions(centroids, sub_vector, 1, distance_type) - .expect("kmeans_find_partitions failed"); - parts.values().iter().map(|v| *v as u8).collect::>() + let parts = + compute_partitions(centroids, sub_vector, sub_dim, distance_type); + parts.iter().map(|v| v.unwrap() as u8).collect::>() }) .collect::>() }) From b328d085583012bc0ddf6788db8a588305bc032c Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Sun, 23 Jun 2024 09:26:33 -0700 Subject: [PATCH 3/4] 10% more --- rust/lance-index/src/vector/pq/utils.rs | 1 + rust/lance-linalg/src/kmeans.rs | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/lance-index/src/vector/pq/utils.rs b/rust/lance-index/src/vector/pq/utils.rs index 08161947ce..80d5ff34d8 100644 --- a/rust/lance-index/src/vector/pq/utils.rs +++ b/rust/lance-index/src/vector/pq/utils.rs @@ -54,6 +54,7 @@ pub fn num_centroids(num_bits: impl Into) -> usize { 2_usize.pow(num_bits.into()) } +#[inline] pub fn get_sub_vector_centroids( codebook: &[T], dimension: usize, diff --git a/rust/lance-linalg/src/kmeans.rs b/rust/lance-linalg/src/kmeans.rs index 85990dd3c6..78661fac2a 100644 --- a/rust/lance-linalg/src/kmeans.rs +++ b/rust/lance-linalg/src/kmeans.rs @@ -686,7 +686,7 @@ pub fn compute_partitions( vectors .par_chunks(dimension) .map(|vec| { - argmin_value(match distance_type { + argmin_value_float(match distance_type { DistanceType::L2 => l2_distance_batch(vec, centroids, dimension), DistanceType::Dot => dot_distance_batch(vec, centroids, dimension), _ => { From df7b647e44b6f157ffb613e6f3864f30fb39860e Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Sun, 23 Jun 2024 09:44:04 -0700 Subject: [PATCH 4/4] more --- rust/lance-index/src/vector/pq.rs | 8 +++---- rust/lance-linalg/src/kmeans.rs | 36 ++++++++++++++++++++----------- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index 0fa659c561..9eee17c681 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -14,7 +14,7 @@ use lance_arrow::*; use lance_core::{Error, Result}; use lance_linalg::distance::{dot_distance_batch, l2_distance_batch, DistanceType, Dot, L2}; use lance_linalg::kernels::argmin_value_float; -use lance_linalg::kmeans::compute_partitions; +use lance_linalg::kmeans::compute_partition; use lance_linalg::{distance::MetricType, MatrixView}; use rayon::prelude::*; use snafu::{location, Location}; @@ -368,7 +368,7 @@ where vector .chunks_exact(sub_dim) .enumerate() - .flat_map(|(sub_idx, sub_vector)| { + .map(|(sub_idx, sub_vector)| { let centroids = get_sub_vector_centroids( codebook.as_slice(), dim, @@ -376,9 +376,7 @@ where num_sub_vectors, sub_idx, ); - let parts = - compute_partitions(centroids, sub_vector, sub_dim, distance_type); - parts.iter().map(|v| v.unwrap() as u8).collect::>() + compute_partition(centroids, sub_vector, distance_type).map(|v| v as u8) }) .collect::>() }) diff --git a/rust/lance-linalg/src/kmeans.rs b/rust/lance-linalg/src/kmeans.rs index 78661fac2a..80f3735b04 100644 --- a/rust/lance-linalg/src/kmeans.rs +++ b/rust/lance-linalg/src/kmeans.rs @@ -685,22 +685,32 @@ pub fn compute_partitions( let dimension = dimension.as_(); vectors .par_chunks(dimension) - .map(|vec| { - argmin_value_float(match distance_type { - DistanceType::L2 => l2_distance_batch(vec, centroids, dimension), - DistanceType::Dot => dot_distance_batch(vec, centroids, dimension), - _ => { - panic!( - "KMeans::find_partitions: {} is not supported", - distance_type - ); - } - }) - .map(|(idx, _)| idx) - }) + .map(|vec| compute_partition(centroids, vec, distance_type)) .collect::>() } +#[inline] +pub fn compute_partition( + centroids: &[T], + vector: &[T], + distance_type: DistanceType, +) -> Option { + match distance_type { + DistanceType::L2 => { + argmin_value_float(l2_distance_batch(vector, centroids, vector.len())).map(|c| c.0) + } + DistanceType::Dot => { + argmin_value_float(dot_distance_batch(vector, centroids, vector.len())).map(|c| c.0) + } + _ => { + panic!( + "KMeans::compute_partition: distance type {} is not supported", + distance_type + ); + } + } +} + #[cfg(test)] mod tests { use std::iter::repeat;