feat: THP and improved result collections experiment

Author: sdd

Cargo.toml

@@ -21,6 +21,8 @@ all-features = true
 cargo_asm = ["dep:rand", "dep:rand_chacha"]
 csv = ["dep:csv"]
 default = ["tracing", "fixed"]
+buffered_result_collection = ["dep:smallvec"]
+huge_pages = ["dep:libc"]
 leaf_view_chunked = []
 leaf_nta_prefetch = []
 modified_van_emde_boas = []
@@ -46,6 +48,7 @@ serde = [
 simd = []
 simulator = []
 rkyv_08 = ["dep:rkyv_08"]
+small_n_result_collectors = ["dep:smallvec"]
 test_utils = ["dep:rand", "dep:rand_chacha", "dep:rayon", "dep:fixed"]
 tracing = ["dep:tracing", "dep:tracing-subscriber"]
 unreliable_select_nth_unstable = []
@@ -90,6 +93,14 @@ features = ["laz-parallel"]
 version = "0.4"
 optional = true
 
+[dependencies.libc]
+version = "0.2"
+optional = true
+
+[dependencies.smallvec]
+version = "1"
+optional = true
+
 [dependencies.kiddo_v5]
 package = "kiddo"
 path = "../kiddo-v5"
@@ -223,6 +234,11 @@ name = "v6_stem_strategies_focus"
 harness = false
 required-features = ["test_utils"]
 
+[[bench]]
+name = "v6_result_collection_focus"
+harness = false
+required-features = ["test_utils"]
+
 #[[example]]
 #name = "avx2-check"
 #path = "examples/avx2-check.rs"

benches/v6_result_collection_focus.rs

@@ -0,0 +1,141 @@
+use codspeed_criterion_compat::{
+    black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput,
+};
+use kiddo::dist::SquaredEuclidean;
+use kiddo::kd_tree::leaf_strategies::VecOfArenas;
+use kiddo::kd_tree::KdTree;
+use kiddo::stem_strategies::donnelly_2_pf::DonnellyPf;
+use rand::Rng;
+use rand::SeedableRng;
+use rand_chacha::ChaCha8Rng;
+use std::num::NonZeroUsize;
+
+const K: usize = 3;
+const B: usize = 32;
+const DEFAULT_QUERY_COUNT: usize = 1_00;
+const DEFAULT_POINT_COUNT: usize = 1usize << 24;
+const DEFAULT_MAX_QTY: usize = 16;
+const DEFAULT_MAX_DIST: f64 = 0.0025;
+const POINT_SEED: u64 = 0x5eed_0000_0000_0301;
+const QUERY_SEED: u64 = 0x5eed_0000_0000_0302;
+
+type ArenaLeaves = VecOfArenas<f64, u32, K, B>;
+type DonnellyPfTree = KdTree<f64, u32, DonnellyPf<3, 64, 8, K>, ArenaLeaves, K, B>;
+
+fn read_usize_env(var: &str, default: usize) -> usize {
+    std::env::var(var)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(default)
+}
+
+fn read_f64_env(var: &str, default: f64) -> f64 {
+    std::env::var(var)
+        .ok()
+        .and_then(|value| value.parse::<f64>().ok())
+        .unwrap_or(default)
+}
+
+fn build_points(point_count: usize) -> Vec<[f64; K]> {
+    let mut rng = ChaCha8Rng::seed_from_u64(POINT_SEED);
+    (0..point_count).map(|_| rng.random::<[f64; K]>()).collect()
+}
+
+fn build_queries(query_count: usize) -> Vec<[f64; K]> {
+    let mut rng = ChaCha8Rng::seed_from_u64(QUERY_SEED);
+    (0..query_count).map(|_| rng.random::<[f64; K]>()).collect()
+}
+
+fn run_sorted_nearest_n_within_queries(
+    tree: &DonnellyPfTree,
+    queries: &[[f64; K]],
+    max_dist: f64,
+    max_qty: NonZeroUsize,
+) -> (usize, u64, f64) {
+    let mut checksum_len = 0usize;
+    let mut checksum_item = 0u64;
+    let mut checksum_dist = 0.0f64;
+
+    for query in queries {
+        let results = tree.nearest_n_within::<SquaredEuclidean<f64>>(
+            black_box(query),
+            max_dist,
+            max_qty,
+            true,
+        );
+        checksum_len += results.len();
+
+        for result in results {
+            checksum_item = checksum_item.wrapping_add(result.item as u64);
+            checksum_dist += result.distance;
+        }
+    }
+
+    (checksum_len, checksum_item, checksum_dist)
+}
+
+fn run_best_n_within_queries(
+    tree: &DonnellyPfTree,
+    queries: &[[f64; K]],
+    max_dist: f64,
+    max_qty: NonZeroUsize,
+) -> (usize, u64, f64) {
+    let mut checksum_len = 0usize;
+    let mut checksum_item = 0u64;
+    let mut checksum_dist = 0.0f64;
+
+    for query in queries {
+        let results =
+            tree.best_n_within::<SquaredEuclidean<f64>>(black_box(query), max_dist, max_qty);
+        checksum_len += results.len();
+
+        for result in results.into_vec() {
+            checksum_item = checksum_item.wrapping_add(result.item as u64);
+            checksum_dist += result.distance;
+        }
+    }
+
+    (checksum_len, checksum_item, checksum_dist)
+}
+
+fn v6_result_collection_focus(c: &mut Criterion) {
+    let query_count = read_usize_env("KIDDO_BENCH_QUERIES", DEFAULT_QUERY_COUNT);
+    let point_count = read_usize_env("KIDDO_BENCH_POINTS", DEFAULT_POINT_COUNT);
+    let max_qty =
+        NonZeroUsize::new(read_usize_env("KIDDO_BENCH_MAX_QTY", DEFAULT_MAX_QTY)).unwrap();
+    let max_dist = read_f64_env("KIDDO_BENCH_MAX_DIST", DEFAULT_MAX_DIST);
+
+    let points = build_points(point_count);
+    let queries = build_queries(query_count);
+    let tree: DonnellyPfTree = KdTree::new_from_slice(&points);
+
+    let mut group = c.benchmark_group("v6 result collection focus");
+    group.throughput(Throughput::Elements(query_count as u64));
+
+    group.bench_function(
+        BenchmarkId::new("sorted nearest_n_within / Donnelly PF", point_count),
+        |b| {
+            b.iter(|| {
+                black_box(run_sorted_nearest_n_within_queries(
+                    &tree, &queries, max_dist, max_qty,
+                ))
+            });
+        },
+    );
+
+    group.bench_function(
+        BenchmarkId::new("best_n_within / Donnelly PF", point_count),
+        |b| {
+            b.iter(|| {
+                black_box(run_best_n_within_queries(
+                    &tree, &queries, max_dist, max_qty,
+                ))
+            });
+        },
+    );
+
+    group.finish();
+}
+
+criterion_group!(benches, v6_result_collection_focus);
+criterion_main!(benches);

benches/v6_stem_strategies_focus.rs

@@ -6,13 +6,13 @@ use kiddo::kd_tree::leaf_strategies::VecOfArenas;
 use kiddo::kd_tree::KdTree;
 use kiddo::stem_strategies::donnelly_2_pf::DonnellyPf;
 use kiddo::stem_strategies::eytzinger_pf_far::EytzingerPfFar;
+use kiddo::stem_strategies::Eytzinger;
 #[cfg(all(
     feature = "simd",
     target_arch = "x86_64",
     any(target_feature = "avx2", target_feature = "avx512f")
 ))]
 use kiddo::stem_strategies::{Block3, DonnellyMarkerSimd};
-use kiddo::stem_strategies::{Donnelly, Eytzinger};
 use rand::Rng;
 use rand::SeedableRng;
 use rand_chacha::ChaCha8Rng;

justfile

@@ -118,6 +118,22 @@ bench-v6-stem-strategies-focus FEATURES='simd,test_utils,logging_off' POINTS='41
     RUSTFLAGS='-C target-cpu=native' \
     cargo criterion --bench v6_stem_strategies_focus --features {{FEATURES}}
 
+bench-v6-stem-strategies-big FEATURES='simd,test_utils,logging_off' POINTS='16777216' QUERIES='10000':
+    RUSTC_WRAPPER= \
+    KIDDO_BENCH_POINTS={{POINTS}} \
+    KIDDO_BENCH_QUERIES={{QUERIES}} \
+    RUSTFLAGS='-C target-cpu=native' \
+    cargo criterion --bench v6_stem_strategies_focus --features {{FEATURES}}
+
+bench-v6-result-collection-focus FEATURES='simd,test_utils,logging_off' POINTS='16777216' QUERIES='100' MAX_QTY='16' MAX_DIST='0.0025':
+    RUSTC_WRAPPER= \
+    KIDDO_BENCH_POINTS={{POINTS}} \
+    KIDDO_BENCH_QUERIES={{QUERIES}} \
+    KIDDO_BENCH_MAX_QTY={{MAX_QTY}} \
+    KIDDO_BENCH_MAX_DIST={{MAX_DIST}} \
+    RUSTFLAGS='-C target-cpu=native' \
+    cargo criterion --bench v6_result_collection_focus --features {{FEATURES}}
+
 profile-v6-stem-exact-stats FEATURES='simd,test_utils,logging_off' POINTS='4194304' QUERIES='10000' REPEATS='1':
     RUSTC_WRAPPER= \
     KIDDO_PROFILE_POINTS={{POINTS}} \