kiddo 6.0.0-alpha.1

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

//! ARM aarch64 NEON SIMD implementations for query orchestration pruning.

/// SIMD prune block for f64 values (NEON)
///
/// Compares 8 rd_values against max_dist in parallel and returns a bitmask.
#[inline(always)]
#[allow(dead_code)]
pub(crate) fn simd_prune_block_f64(rd_values: &[f64; 8], max_dist: f64, sibling_mask: u8) -> u8 {
    unsafe {
        use core::arch::aarch64::*;

        let max_vec = vdupq_n_f64(max_dist);
        let rd_0 = vld1q_f64(rd_values.as_ptr());
        let rd_1 = vld1q_f64(rd_values.as_ptr().add(2));
        let rd_2 = vld1q_f64(rd_values.as_ptr().add(4));
        let rd_3 = vld1q_f64(rd_values.as_ptr().add(6));

        let cmp_0 = vcleq_f64(rd_0, max_vec);
        let cmp_1 = vcleq_f64(rd_1, max_vec);
        let cmp_2 = vcleq_f64(rd_2, max_vec);
        let cmp_3 = vcleq_f64(rd_3, max_vec);

        let weights_0 = [1u64, 2u64];
        let weights_1 = [4u64, 8u64];
        let weights_2 = [16u64, 32u64];
        let weights_3 = [64u64, 128u64];

        let mask_0 = vaddvq_u64(vandq_u64(cmp_0, vld1q_u64(weights_0.as_ptr())));
        let mask_1 = vaddvq_u64(vandq_u64(cmp_1, vld1q_u64(weights_1.as_ptr())));
        let mask_2 = vaddvq_u64(vandq_u64(cmp_2, vld1q_u64(weights_2.as_ptr())));
        let mask_3 = vaddvq_u64(vandq_u64(cmp_3, vld1q_u64(weights_3.as_ptr())));

        let mask = (mask_0 | mask_1 | mask_2 | mask_3) as u8;

        mask & sibling_mask
    }
}

/// SIMD prune block for f32 values (NEON)
///
/// Compares 8 rd_values against max_dist in parallel and returns a bitmask.
#[inline(always)]
#[allow(dead_code)]
pub(crate) fn simd_prune_block_f32(rd_values: &[f32; 8], max_dist: f32, sibling_mask: u8) -> u8 {
    unsafe {
        use core::arch::aarch64::*;

        let max_vec = vdupq_n_f32(max_dist);
        let rd_0 = vld1q_f32(rd_values.as_ptr());
        let rd_1 = vld1q_f32(rd_values.as_ptr().add(4));

        let cmp_0 = vcleq_f32(rd_0, max_vec);
        let cmp_1 = vcleq_f32(rd_1, max_vec);

        let weights_0 = [1u32, 2u32, 4u32, 8u32];
        let weights_1 = [16u32, 32u32, 64u32, 128u32];

        let mask_0 = vaddvq_u32(vandq_u32(cmp_0, vld1q_u32(weights_0.as_ptr())));
        let mask_1 = vaddvq_u32(vandq_u32(cmp_1, vld1q_u32(weights_1.as_ptr())));

        let mask = (mask_0 | mask_1) as u8;

        mask & sibling_mask
    }
}