perf: refactor donnelly block simd code

Author: sdd

Cargo.toml

@@ -218,6 +218,11 @@ name = "v6_stem_strategies"
 harness = false
 required-features = ["test_utils"]
 
+[[bench]]
+name = "v6_stem_strategies_focus"
+harness = false
+required-features = ["test_utils"]
+
 #[[example]]
 #name = "avx2-check"
 #path = "examples/avx2-check.rs"
@@ -412,6 +417,16 @@ name = "profile_v6_leaf_strategies"
 path = "src/bin/profile_v6_leaf_strategies.rs"
 required-features = ["profile_v5"]
 
+[[bin]]
+name = "profile_v6_stem_exact_stats"
+path = "src/bin/profile_v6_stem_exact_stats.rs"
+required-features = ["simd", "test_utils"]
+
+[[bin]]
+name = "repro_donnelly_block3_exact_divergence"
+path = "src/bin/repro_donnelly_block3_exact_divergence.rs"
+required-features = ["simd", "test_utils"]
+
 [[test]]
 name = "donnelly_simd_regressions"
 path = "tests/donnelly_simd_regressions.rs"

benches/v6_stem_strategies_focus.rs

@@ -0,0 +1,169 @@
+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 kiddo::stem_strategies::eytzinger_pf_far::EytzingerPfFar;
+#[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;
+
+const K: usize = 3;
+const B: usize = 32;
+const DEFAULT_QUERY_COUNT: usize = 10_000;
+const DEFAULT_POINT_COUNT: usize = 1usize << 22;
+const POINT_SEED: u64 = 0x5eed_0000_0000_0201;
+const QUERY_SEED: u64 = 0x5eed_0000_0000_0202;
+
+type ArenaLeaves = VecOfArenas<f64, u32, K, B>;
+type EytzingerTree = KdTree<f64, u32, Eytzinger<K>, ArenaLeaves, K, B>;
+type EytzingerPfFarTree = KdTree<f64, u32, EytzingerPfFar<K, 8>, ArenaLeaves, K, B>;
+type DonnellyPfTree = KdTree<f64, u32, DonnellyPf<3, 64, 8, K>, ArenaLeaves, K, B>;
+#[cfg(all(
+    feature = "simd",
+    target_arch = "x86_64",
+    any(target_feature = "avx2", target_feature = "avx512f")
+))]
+type DonnellySimdTree = KdTree<f64, u32, DonnellyMarkerSimd<Block3, 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 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_nearest_queries_eytzinger(tree: &EytzingerTree, queries: &[[f64; K]]) -> (f64, u64) {
+    let mut checksum_dist = 0.0f64;
+    let mut checksum_item = 0u64;
+
+    for query in queries {
+        let (dist, item) = tree.nearest_one::<SquaredEuclidean<f64>>(black_box(query));
+        checksum_dist += dist;
+        checksum_item = checksum_item.wrapping_add(item as u64);
+    }
+
+    (checksum_dist, checksum_item)
+}
+
+fn run_nearest_queries_eytzinger_pf_far(
+    tree: &EytzingerPfFarTree,
+    queries: &[[f64; K]],
+) -> (f64, u64) {
+    let mut checksum_dist = 0.0f64;
+    let mut checksum_item = 0u64;
+
+    for query in queries {
+        let (dist, item) = tree.nearest_one::<SquaredEuclidean<f64>>(black_box(query));
+        checksum_dist += dist;
+        checksum_item = checksum_item.wrapping_add(item as u64);
+    }
+
+    (checksum_dist, checksum_item)
+}
+
+fn run_nearest_queries_donnelly(tree: &DonnellyPfTree, queries: &[[f64; K]]) -> (f64, u64) {
+    let mut checksum_dist = 0.0f64;
+    let mut checksum_item = 0u64;
+
+    for query in queries {
+        let (dist, item) = tree.nearest_one::<SquaredEuclidean<f64>>(black_box(query));
+        checksum_dist += dist;
+        checksum_item = checksum_item.wrapping_add(item as u64);
+    }
+
+    (checksum_dist, checksum_item)
+}
+
+#[cfg(all(
+    feature = "simd",
+    target_arch = "x86_64",
+    any(target_feature = "avx2", target_feature = "avx512f")
+))]
+fn run_nearest_queries_donnelly_simd(tree: &DonnellySimdTree, queries: &[[f64; K]]) -> (f64, u64) {
+    let mut checksum_dist = 0.0f64;
+    let mut checksum_item = 0u64;
+
+    for query in queries {
+        let (dist, item) = tree.nearest_one::<SquaredEuclidean<f64>>(black_box(query));
+        checksum_dist += dist;
+        checksum_item = checksum_item.wrapping_add(item as u64);
+    }
+
+    (checksum_dist, checksum_item)
+}
+
+fn v6_stem_strategies_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 points = build_points(point_count);
+    let queries = build_queries(query_count);
+
+    let eytzinger_tree: EytzingerTree = KdTree::new_from_slice(&points);
+    let eytzinger_pf_far_tree: EytzingerPfFarTree = KdTree::new_from_slice(&points);
+    let donnelly_tree: DonnellyPfTree = KdTree::new_from_slice(&points);
+    #[cfg(all(
+        feature = "simd",
+        target_arch = "x86_64",
+        any(target_feature = "avx2", target_feature = "avx512f")
+    ))]
+    let donnelly_simd_tree: DonnellySimdTree = KdTree::new_from_slice(&points);
+
+    let mut group = c.benchmark_group("v6 nearest_one stem strategies focus");
+    group.throughput(Throughput::Elements(query_count as u64));
+
+    group.bench_function(BenchmarkId::new("Eytzinger", point_count), |b| {
+        b.iter(|| black_box(run_nearest_queries_eytzinger(&eytzinger_tree, &queries)));
+    });
+
+    group.bench_function(BenchmarkId::new("Eytzinger PF Far", point_count), |b| {
+        b.iter(|| {
+            black_box(run_nearest_queries_eytzinger_pf_far(
+                &eytzinger_pf_far_tree,
+                &queries,
+            ))
+        });
+    });
+
+    group.bench_function(BenchmarkId::new("Donnelly PF", point_count), |b| {
+        b.iter(|| black_box(run_nearest_queries_donnelly(&donnelly_tree, &queries)));
+    });
+
+    #[cfg(all(
+        feature = "simd",
+        target_arch = "x86_64",
+        any(target_feature = "avx2", target_feature = "avx512f")
+    ))]
+    group.bench_function(BenchmarkId::new("Donnelly Block SIMD", point_count), |b| {
+        b.iter(|| {
+            black_box(run_nearest_queries_donnelly_simd(
+                &donnelly_simd_tree,
+                &queries,
+            ))
+        });
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, v6_stem_strategies_focus);
+criterion_main!(benches);

justfile

@@ -99,12 +99,49 @@ asm-k6-nearest-one-donnelly-block3-fill-avx512-clean:
 asm-k6-nearest-one-donnelly-block3-pending-select-avx512-clean:
     RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "donnelly_block3_pending_select_f64_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_donnelly_block3_pending_select_avx512_clean.asm
 
+asm-k6-nearest-one-donnelly-block3-pending-fast-path-avx512-clean:
+    RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "donnelly_block3_pending_fast_path_f64_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_donnelly_block3_pending_fast_path_avx512_clean.asm
+
+asm-k6-nearest-one-donnelly-block3-exact-step-avx512-clean:
+    RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "donnelly_block3_exact_step_f64_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_donnelly_block3_exact_step_avx512_clean.asm
+
+asm-k6-nearest-one-donnelly-voarena-v3-avx512-clean:
+    RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "v6_nearest_one_donnelly_vec_of_arenas_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_donnelly_vec_of_arenas_v3_avx512_clean.asm
+
+asm-k6-nearest-one-donnelly-blocksimd-voarena-v3-avx512-clean:
+    RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "v6_nearest_one_donnelly_blocksimd_vec_of_arenas_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_donnelly_blocksimd_vec_of_arenas_v3_avx512_clean.asm
+
+bench-v6-stem-strategies-focus FEATURES='simd,test_utils,logging_off' POINTS='4194304' 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}}
+
+profile-v6-stem-exact-stats FEATURES='simd,test_utils,logging_off' POINTS='4194304' QUERIES='10000' REPEATS='1':
+    RUSTC_WRAPPER= \
+    KIDDO_PROFILE_POINTS={{POINTS}} \
+    KIDDO_PROFILE_QUERIES={{QUERIES}} \
+    KIDDO_PROFILE_QUERY_BATCH_REPEATS={{REPEATS}} \
+    RUSTFLAGS='-C target-cpu=native' \
+    cargo run --release --bin profile_v6_stem_exact_stats --features {{FEATURES}}
+
+repro-donnelly-block3-exact-divergence FEATURES='simd,test_utils,logging_off' POINTS='4194304' QUERIES='10000':
+    RUSTC_WRAPPER= \
+    KIDDO_REPRO_POINTS={{POINTS}} \
+    KIDDO_REPRO_QUERIES={{QUERIES}} \
+    RUSTFLAGS='-C target-cpu=native' \
+    cargo run --release --bin repro_donnelly_block3_exact_divergence --features {{FEATURES}}
+
 asm-k6-nearest-one-eytz-v3-core-avx512:
     cargo asm --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "v6_nearest_one_eytzinger_arithmetic_core_cargo_asm_hook" > v6_nearest_one_eytzinger_v3_core_avx512.asm
 
 asm-k6-nearest-one-eytz-v3-leaf-avx512:
     cargo asm --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "kiddo::kd_tree::leaf_view_chunked::nearest_one::avx512::leaf_nearest_one_chunked_nozero_f64_k3::<f64, kiddo::dist::squared_euclidean::SquaredEuclidean<f64>, usize>" > v6_nearest_one_eytzinger_v3_leaf_avx512.asm
 
+asm-k6-nearest-one-arena-leaf-v3-avx512-clean:
+    RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "v6_nearest_one_arena_leaf_cargo_asm_hook" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_arena_leaf_v3_avx512_clean.asm
+
 asm-k6-nearest-one-eytz-v3-leaf-avx512-clean:
     RUSTC_WRAPPER= cargo asm --simplify --features simd,cargo_asm,logging_off --lib --target-cpu=native -C="opt-level=2" -C="target-cpu=native" "kiddo::kd_tree::leaf_view_chunked::nearest_one::avx512::leaf_nearest_one_chunked_nozero_f64_k3::<f64, kiddo::dist::squared_euclidean::SquaredEuclidean<f64>, usize>" | python3 scripts/clean_cargo_asm.py > v6_nearest_one_eytzinger_v3_leaf_avx512_clean.asm