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

mod fallback;

#[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
pub(crate) mod avx512;

use crate::dist::DistanceMetric;
use crate::leaf_view::{LeafArena, LeafView};
use crate::{Axis, Content};
// #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
// pub(crate) use avx512::nearest_one_avx512_unchecked;

pub(crate) use fallback::{
    nearest_one_with_query_wide_arena_fallback, nearest_one_with_query_wide_fallback,
};

#[inline(always)]
pub(crate) fn nearest_one_with_query_wide_arena<AX, T, D, const K: usize>(
    arena: &LeafArena<'_, AX, T, K>,
    query_wide: &[D::Output; K],
    best_dist: &mut D::Output,
    best_item: &mut T,
) where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + 'static,
{
    #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
    if unsafe {
        try_nearest_one_arena_avx512::<AX, T, D, K>(arena, query_wide, best_dist, best_item)
    } {
        return;
    }

    nearest_one_with_query_wide_arena_fallback::<AX, T, D, K>(
        arena, query_wide, best_dist, best_item,
    );
}

#[inline(always)]
pub(crate) fn nearest_one_with_query_wide<AX, T, D, const K: usize, const B: usize>(
    leaf: &LeafView<'_, AX, T, K, B>,
    query_wide: &[D::Output; K],
    best_dist: &mut D::Output,
    best_item: &mut T,
) where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + 'static,
{
    #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
    if unsafe { try_nearest_one_avx512::<AX, T, D, K, B>(leaf, query_wide, best_dist, best_item) } {
        return;
    }

    nearest_one_with_query_wide_fallback::<AX, T, D, K, B>(leaf, query_wide, best_dist, best_item);
}

#[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
#[inline(always)]
unsafe fn try_nearest_one_avx512<AX, T, D, const K: usize, const B: usize>(
    leaf: &LeafView<'_, AX, T, K, B>,
    query_wide: &[D::Output; K],
    best_dist: &mut D::Output,
    best_item: &mut T,
) -> bool
where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + 'static,
{
    D::try_nearest_one_leaf_avx512(leaf, query_wide, best_dist, best_item)
}

#[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx512f"))]
#[inline(always)]
unsafe fn try_nearest_one_arena_avx512<AX, T, D, const K: usize>(
    arena: &LeafArena<'_, AX, T, K>,
    query_wide: &[D::Output; K],
    best_dist: &mut D::Output,
    best_item: &mut T,
) -> bool
where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + 'static,
{
    D::try_nearest_one_arena_avx512(arena, query_wide, best_dist, best_item)
}