kiddo 6.0.0

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! ARM aarch64 SIMD implementations (NEON) for Donnelly block marker strategy.

use std::ptr::NonNull;

/// Compare query value against 7 pivots in a Block3 (f64, NEON)
///
/// Returns the child index (0-7) based on how many pivots the query exceeds.
///
/// # Safety
/// - `stems_ptr` must point to at least 8 contiguous f64 values starting at
///   `cache_line_base * 8`.
/// - The caller must ensure the pointer is valid for reads and properly aligned
///   for `f64` loads on the target.
#[inline(always)]
pub unsafe fn compare_block3_f64_neon(
    stems_ptr: NonNull<u8>,
    cache_line_base: usize,
    query_val: f64,
) -> u8 {
    use core::arch::aarch64::*;

    let ptr = stems_ptr.as_ptr().add(cache_line_base * 8) as *const f64;
    let query_vec = vdupq_n_f64(query_val);

    let pivots_0 = vld1q_f64(ptr);
    let pivots_1 = vld1q_f64(ptr.add(2));
    let pivots_2 = vld1q_f64(ptr.add(4));
    let pivots_3 = vld1q_f64(ptr.add(6));

    let cmp_0 = vcgeq_f64(query_vec, pivots_0);
    let cmp_1 = vcgeq_f64(query_vec, pivots_1);
    let cmp_2 = vcgeq_f64(query_vec, pivots_2);
    let cmp_3 = vcgeq_f64(query_vec, pivots_3);

    let ones_0 = vshrq_n_u64(cmp_0, 63);
    let ones_1 = vshrq_n_u64(cmp_1, 63);
    let ones_2 = vshrq_n_u64(cmp_2, 63);
    let ones_3 = vshrq_n_u64(cmp_3, 63);

    let count_0 = vaddvq_u64(ones_0);
    let count_1 = vaddvq_u64(ones_1);
    let count_2 = vaddvq_u64(ones_2);
    let count_3 = vaddvq_u64(ones_3);

    (count_0 + count_1 + count_2 + count_3) as u8
}

/// Compare query value against 7 pivots in a Block3 (f32, NEON)
///
/// Returns the child index (0-7) based on how many pivots the query exceeds.
///
/// # Safety
/// - `stems_ptr` must point to at least 8 contiguous f32 values starting at
///   `cache_line_base * 4`.
/// - The caller must ensure the pointer is valid for reads and properly aligned
///   for `f32` loads on the target.
#[inline(always)]
pub unsafe fn compare_block3_f32_neon(
    stems_ptr: NonNull<u8>,
    cache_line_base: usize,
    query_val: f32,
) -> u8 {
    use core::arch::aarch64::*;

    let ptr = stems_ptr.as_ptr().add(cache_line_base * 4) as *const f32;
    let query_vec = vdupq_n_f32(query_val);

    let pivots_0 = vld1q_f32(ptr);
    let pivots_1 = vld1q_f32(ptr.add(4));

    let cmp_0 = vcgeq_f32(query_vec, pivots_0);
    let cmp_1 = vcgeq_f32(query_vec, pivots_1);

    let ones_0 = vshrq_n_u32(cmp_0, 31);
    let ones_1 = vshrq_n_u32(cmp_1, 31);

    let count_0 = vaddvq_u32(ones_0);
    let count_1 = vaddvq_u32(ones_1);

    (count_0 + count_1) as u8
}

/// Compare query value against 15 pivots in a Block4 (f32, NEON)
///
/// Returns the child index (0-15) based on how many pivots the query exceeds.
///
/// # Safety
/// - `stems_ptr` must point to at least 16 contiguous f32 values starting at
///   `cache_line_base * 4`.
/// - The caller must ensure the pointer is valid for reads and properly aligned
///   for `f32` loads on the target.
#[inline(always)]
pub unsafe fn compare_block4_f32_neon(
    stems_ptr: NonNull<u8>,
    cache_line_base: usize,
    query_val: f32,
) -> u8 {
    use core::arch::aarch64::*;

    let ptr = stems_ptr.as_ptr().add(cache_line_base * 4) as *const f32;
    let query_vec = vdupq_n_f32(query_val);

    let pivots_0 = vld1q_f32(ptr);
    let pivots_1 = vld1q_f32(ptr.add(4));
    let pivots_2 = vld1q_f32(ptr.add(8));
    let pivots_3 = vld1q_f32(ptr.add(12));

    let cmp_0 = vcgeq_f32(query_vec, pivots_0);
    let cmp_1 = vcgeq_f32(query_vec, pivots_1);
    let cmp_2 = vcgeq_f32(query_vec, pivots_2);
    let cmp_3 = vcgeq_f32(query_vec, pivots_3);

    let ones_0 = vshrq_n_u32(cmp_0, 31);
    let ones_1 = vshrq_n_u32(cmp_1, 31);
    let ones_2 = vshrq_n_u32(cmp_2, 31);
    let ones_3 = vshrq_n_u32(cmp_3, 31);

    let count_0 = vaddvq_u32(ones_0);
    let count_1 = vaddvq_u32(ones_1);
    let count_2 = vaddvq_u32(ones_2);
    let count_3 = vaddvq_u32(ones_3);

    (count_0 + count_1 + count_2 + count_3) as u8
}