Skip to content

Commit

Permalink
fix: benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
@nishaq503
nishaq503 committed Aug 8, 2024
1 parent 720b234 commit 1a5e7e7
Show file tree
Hide file tree
Showing 8 changed files with 279 additions and 189 deletions.
11 changes: 10 additions & 1 deletion crates/abd-clam/benches/ann_benchmarks.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,11 @@
mod utils;

use abd_clam::{cakes::OffBall, partition::ParPartition, Ball, Cluster, Metric};
use abd_clam::{
cakes::{OffBall, SquishyBall},
partition::ParPartition,
Ball, Cluster, Metric,
};
use criterion::*;

fn ann_benchmarks(c: &mut Criterion) {
Expand Down Expand Up @@ -39,6 +43,8 @@ fn ann_benchmarks(c: &mut Criterion) {
let mut perm_data = data.clone();
let perm_root = OffBall::par_from_ball_tree(root.clone(), &mut perm_data);

let (dec_root, dec_data) = SquishyBall::par_new_tree(&mut data.clone(), &criteria, seed);

utils::compare_permuted(
c,
data_name,
Expand All @@ -47,10 +53,13 @@ fn ann_benchmarks(c: &mut Criterion) {
&root,
&perm_data,
&perm_root,
&dec_data,
&dec_root,
queries,
&radii,
&ks,
true,
false,
);
}
}
Expand Down
11 changes: 10 additions & 1 deletion crates/abd-clam/benches/genomic_search.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,11 @@
mod utils;

use abd_clam::{cakes::OffBall, partition::ParPartition, Ball, Cluster, FlatVec, Metric};
use abd_clam::{
cakes::{OffBall, SquishyBall},
partition::ParPartition,
Ball, Cluster, FlatVec, Metric,
};
use criterion::*;
use rand::prelude::*;

Expand Down Expand Up @@ -62,6 +66,8 @@ fn genomic_search(c: &mut Criterion) {
let mut perm_data = data.clone();
let perm_root = OffBall::par_from_ball_tree(root.clone(), &mut perm_data);

let (dec_root, dec_data) = SquishyBall::par_new_tree(&mut data.clone(), &criteria, seed);

utils::compare_permuted(
c,
"genomic-search",
Expand All @@ -70,10 +76,13 @@ fn genomic_search(c: &mut Criterion) {
&root,
&perm_data,
&perm_root,
&dec_data,
&dec_root,
&queries,
&radii,
&ks,
true,
true,
);
}
}
Expand Down
68 changes: 51 additions & 17 deletions crates/abd-clam/benches/utils/compare_permuted.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@
use abd_clam::{
cakes::{
cluster::{ParSearchable, Searchable},
Algorithm, OffBall,
Algorithm, Decodable, Encodable, OffBall, ParCompressible, ParDecompressible, SquishyBall,
},
dataset::ParDataset,
linear_search::ParLinearSearch,
Ball,
};
use criterion::*;
use distances::Number;
Expand All @@ -29,26 +30,27 @@ use distances::Number;
///
/// - `I`: The type of the items in the dataset.
/// - `U`: The type of the scalars used to measure distances.
/// - `C`: The type of the cluster for the original dataset.
/// - `D`: The type of the original dataset.
/// - `Dp`: The type of the permuted dataset.
pub fn compare_permuted<I, U, C, D>(
/// - `Co`: The type of the original dataset.
pub fn compare_permuted<I, U, Co, Dec>(
c: &mut Criterion,
data_name: &str,
metric_name: &str,
data: &D,
root: &C,
perm_data: &D,
perm_root: &OffBall<I, U, D, C>,
data: &Co,
root: &Ball<I, U, Co>,
perm_data: &Co,
perm_root: &OffBall<I, U, Co, Ball<I, U, Co>>,
dec_data: &Dec,
dec_root: &SquishyBall<I, U, Co, Dec, Ball<I, U, Co>>,
queries: &[I],
radii: &[U],
ks: &[usize],
par_only: bool,
squishy: bool,
) where
I: Send + Sync,
I: Encodable + Decodable + Send + Sync,
U: Number,
D: ParDataset<I, U>,
C: ParSearchable<I, U, D>,
Co: ParCompressible<I, U> + ParLinearSearch<I, U>,
Dec: ParDecompressible<I, U>,
{
let algs = vec![
Algorithm::KnnRepeatedRnn(ks[0], U::ONE.double()),
Expand All @@ -66,20 +68,36 @@ pub fn compare_permuted<I, U, C, D>(
let alg = Algorithm::RnnClustered(radius);

if !par_only {
group.bench_with_input(BenchmarkId::new("Linear", radius), &radius, |b, _| {
b.iter_with_large_drop(|| alg.batch_linear_search(data, queries));
});
group.bench_with_input(BenchmarkId::new("Ball", radius), &radius, |b, _| {
b.iter_with_large_drop(|| root.batch_search(data, queries, alg));
});
group.bench_with_input(BenchmarkId::new("OffsetBall", radius), &radius, |b, _| {
group.bench_with_input(BenchmarkId::new("OffBall", radius), &radius, |b, _| {
b.iter_with_large_drop(|| perm_root.batch_search(perm_data, queries, alg));
});
if squishy {
group.bench_with_input(BenchmarkId::new("SquishyBall", radius), &radius, |b, _| {
b.iter_with_large_drop(|| dec_root.batch_search(dec_data, queries, alg));
});
}
}

group.bench_with_input(BenchmarkId::new("ParLinear", radius), &radius, |b, _| {
b.iter_with_large_drop(|| alg.par_batch_linear_search(data, queries));
});
group.bench_with_input(BenchmarkId::new("ParBall", radius), &radius, |b, _| {
b.iter_with_large_drop(|| root.par_batch_search(data, queries, alg));
});
group.bench_with_input(BenchmarkId::new("ParOffsetBall", radius), &radius, |b, _| {
group.bench_with_input(BenchmarkId::new("ParOffBall", radius), &radius, |b, _| {
b.iter_with_large_drop(|| perm_root.par_batch_search(perm_data, queries, alg));
});
if squishy {
group.bench_with_input(BenchmarkId::new("ParSquishyBall", radius), &radius, |b, _| {
b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, queries, alg));
});
}
}
group.finish();

Expand All @@ -94,20 +112,36 @@ pub fn compare_permuted<I, U, C, D>(
let alg = alg.with_params(U::ZERO, k);

if !par_only {
group.bench_with_input(BenchmarkId::new("Linear", k), &k, |b, _| {
b.iter_with_large_drop(|| alg.batch_linear_search(data, queries));
});
group.bench_with_input(BenchmarkId::new("Ball", k), &k, |b, _| {
b.iter_with_large_drop(|| root.batch_search(data, queries, alg));
});
group.bench_with_input(BenchmarkId::new("OffsetBall", k), &k, |b, _| {
group.bench_with_input(BenchmarkId::new("OffBall", k), &k, |b, _| {
b.iter_with_large_drop(|| perm_root.batch_search(perm_data, queries, alg));
});
if squishy {
group.bench_with_input(BenchmarkId::new("SquishyBall", k), &k, |b, _| {
b.iter_with_large_drop(|| dec_root.batch_search(dec_data, queries, alg));
});
}
}

group.bench_with_input(BenchmarkId::new("ParLinear", k), &k, |b, _| {
b.iter_with_large_drop(|| alg.par_batch_linear_search(data, queries));
});
group.bench_with_input(BenchmarkId::new("ParBall", k), &k, |b, _| {
b.iter_with_large_drop(|| root.par_batch_search(data, queries, alg));
});
group.bench_with_input(BenchmarkId::new("ParOffsetBall", k), &k, |b, _| {
group.bench_with_input(BenchmarkId::new("ParOffBall", k), &k, |b, _| {
b.iter_with_large_drop(|| perm_root.par_batch_search(perm_data, queries, alg));
});
if squishy {
group.bench_with_input(BenchmarkId::new("ParSquishyBall", k), &k, |b, _| {
b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, queries, alg));
});
}
}
group.finish();
}
Expand Down
19 changes: 14 additions & 5 deletions crates/abd-clam/benches/vector_search.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,11 @@
mod utils;

use abd_clam::{cakes::OffBall, partition::ParPartition, Ball, Cluster, FlatVec, Metric};
use abd_clam::{
cakes::{OffBall, SquishyBall},
partition::ParPartition,
Ball, Cluster, FlatVec, Metric,
};
use criterion::*;
use rand::prelude::*;

Expand All @@ -16,9 +20,9 @@ const METRICS: &[(&str, fn(&Vec<f32>, &Vec<f32>) -> f32)] = &[
];

fn vector_search(c: &mut Criterion) {
let cardinality = 100_000;
let dimensionality = 100;
let max_val = 10.0;
let cardinality = 1_000_000;
let dimensionality = 10;
let max_val = 2.0;
let min_val = -max_val;
let seed = 42;
let rows = symagen::random_data::random_tabular_seedable(cardinality, dimensionality, min_val, max_val, seed);
Expand All @@ -36,7 +40,7 @@ fn vector_search(c: &mut Criterion) {
};

let seed = Some(seed);
let radii = vec![0.01, 0.05, 0.1, 0.5];
let radii = vec![0.001, 0.005, 0.01, 0.1];
let ks = vec![1, 10, 100];
for &(metric_name, distance_fn) in METRICS {
let metric = Metric::new(distance_fn, true);
Expand All @@ -48,6 +52,8 @@ fn vector_search(c: &mut Criterion) {
let mut perm_data = data.clone();
let perm_root = OffBall::par_from_ball_tree(root.clone(), &mut perm_data);

let (dec_root, dec_data) = SquishyBall::par_new_tree(&mut data.clone(), &criteria, seed);

utils::compare_permuted(
c,
"vector-search",
Expand All @@ -56,10 +62,13 @@ fn vector_search(c: &mut Criterion) {
&root,
&perm_data,
&perm_root,
&dec_data,
&dec_root,
&queries,
&radii,
&ks,
false,
false,
);
}
}
Expand Down
Loading

0 comments on commit 1a5e7e7

Please sign in to comment.