feat(wip): more immutable dynamic enhancements

Cargo.toml

@@ -41,6 +41,7 @@ csv = "1"
 cmov = "0.3.1"
 rstest = "0.23.0"
 array-init = "2.1.0"
+aligned-vec = "0.6.1"
 
 [dev-dependencies]
 bincode = "1.3"

examples/pointcloud-csv.rs

@@ -10,7 +10,7 @@ use std::io::Write;
 use std::time::Instant;
 
 use csv::Reader;
-use kiddo::ImmutableKdTree;
+use kiddo::ImmutableKdDynamicTree;
 use kiddo::SquaredEuclidean;
 use rkyv::ser::serializers::{AlignedSerializer, BufferScratch, CompositeSerializer};
 use rkyv::ser::Serializer;
@@ -30,7 +30,7 @@ struct Point {
     z: f64,
 }
 
-type Tree = ImmutableKdTree<f64, 3>;
+type Tree = ImmutableKdDynamicTree<f64, 3>;
 
 fn main() -> Result<(), Box<dyn Error>> {
     #[cfg(feature = "tracing")]
@@ -61,8 +61,6 @@ fn main() -> Result<(), Box<dyn Error>> {
         ElapsedDuration::new(start.elapsed())
     );
 
-    println!("Tree Stats: {:?}", kdtree.generate_stats());
-
     // Test query on the newly created tree
     let query = [0.123f64, 0.456f64, 0.789f64];
     let nearest_neighbour = kdtree.nearest_one::<SquaredEuclidean>(&query);

examples/pointcloud-las.rs

@@ -9,7 +9,7 @@ use std::fs::File;
 use std::io::Write;
 use std::time::Instant;
 
-use kiddo::ImmutableKdTree;
+use kiddo::ImmutableKdDynamicTree;
 use kiddo::SquaredEuclidean;
 use las::Reader;
 
@@ -22,7 +22,7 @@ use tracing_subscriber::fmt;
 const BUFFER_LEN: usize = 10_000_000_000;
 const SCRATCH_LEN: usize = 1_000_000_000;
 
-type Tree = ImmutableKdTree<f32, 3>;
+type Tree = ImmutableKdDynamicTree<f32, 3>;
 
 fn main() -> Result<(), Box<dyn Error>> {
     #[cfg(feature = "tracing")]
@@ -53,8 +53,6 @@ fn main() -> Result<(), Box<dyn Error>> {
         ElapsedDuration::new(start.elapsed())
     );
 
-    println!("Tree Stats: {:?}", kdtree.generate_stats());
-
     // Test query on the newly created tree
     let query = [0.123f32, 0.456f32, 0.789f32];
     let nearest_neighbour = kdtree.nearest_one::<SquaredEuclidean>(&query);

src/float_leaf_slice/fallback.rs

@@ -2,6 +2,7 @@ use az::Cast;
 
 use crate::{float::kdtree::Axis, types::Content};
 
+#[inline]
 pub(crate) fn get_best_from_dists_autovec<A: Axis, T: Content>(
     acc: &[A],
     items: &[T],

src/float_leaf_slice/leaf_slice.rs

@@ -1,6 +1,8 @@
 use az::Cast;
 use std::slice::ChunksExact;
 
+const CHUNK_SIZE: usize = 32;
+
 #[cfg(all(
     feature = "simd",
     target_feature = "avx2",
@@ -34,6 +36,7 @@ where
     T: Content,
     usize: Cast<T>,
 {
+    #[inline]
     pub fn nearest_one<D>(&self, query: &[A; K], best_dist: &mut A, best_item: &mut T)
     where
         D: DistanceMetric<A, K>,
@@ -77,6 +80,7 @@ impl<'a, A: Axis, T: Content, const K: usize, const C: usize> Iterator
 {
     type Item = ([&'a [A; C]; K], &'a [T; C]);
 
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
         if self.items_iterator.is_empty() {
             None
@@ -93,6 +97,7 @@ impl<'a, A: Axis, T: Content, const K: usize, const C: usize> Iterator
 impl<'a, A: Axis, T: Content, const K: usize, const C: usize>
     LeafFixedSliceIterator<'a, A, T, K, C>
 {
+    #[inline]
     fn remainder(&self) -> ([&'a [A]; K], &'a [T]) {
         (
             array_init::array_init(|i| self.points_iterators[i].remainder()),
@@ -122,6 +127,7 @@ where
     T: Content,
     usize: Cast<T>,
 {
+    #[inline]
     pub fn new<'a>(content_points: [&'a [A]; K], content_items: &'a [T]) -> LeafSlice<'a, A, T, K> {
         let size = content_items.len();
         for arr in content_points {
@@ -134,6 +140,7 @@ where
         }
     }
 
+    #[inline]
     fn as_full_chunks<const C: usize>(&self) -> LeafFixedSliceIterator<A, T, K, C> {
         let points_iterators = self.content_points.map(|i| i.chunks_exact(C));
         let items_iterator = self.content_items.chunks_exact(C);
@@ -144,14 +151,15 @@ where
         }
     }
 
+    #[inline]
     pub fn nearest_one<D>(&self, query: &[A; K], best_dist: &mut A, best_item: &mut T)
     where
         D: DistanceMetric<A, K>,
     {
-        let chunk_iter = self.as_full_chunks::<64>();
+        let chunk_iter = self.as_full_chunks::<CHUNK_SIZE>();
         let (remainder_points, remainder_items) = chunk_iter.remainder();
         for chunk in chunk_iter {
-            let dists = A::dists_for_chunk::<D, 64>(chunk.0, query);
+            let dists = A::dists_for_chunk::<D, CHUNK_SIZE>(chunk.0, query);
             A::update_best_dist(dists, chunk.1, best_dist, best_item);
         }
 
@@ -180,6 +188,7 @@ where
     T: Content,
     usize: Cast<T>,
 {
+    #[inline]
     fn update_best_dist<const C: usize>(
         acc: [f64; C],
         items: &[T; C],
@@ -213,6 +222,7 @@ where
         }
     }
 
+    #[inline]
     fn dists_for_chunk<D, const C: usize>(chunk: [&[Self; C]; K], query: &[Self; K]) -> [Self; C]
     where
         D: DistanceMetric<Self, K>,
@@ -237,6 +247,7 @@ where
     T: Content,
     usize: Cast<T>,
 {
+    #[inline]
     fn update_best_dist<const C: usize>(
         acc: [f32; C],
         items: &[T; C],
@@ -268,6 +279,7 @@ where
         }
     }
 
+    #[inline]
     fn dists_for_chunk<D, const C: usize>(chunk: [&[Self; C]; K], query: &[Self; K]) -> [Self; C]
     where
         D: DistanceMetric<Self, K>,

src/immutable_dynamic/common/generate_immutable_dynamic_approx_nearest_one.rs

@@ -20,7 +20,7 @@ macro_rules! generate_immutable_dynamic_approx_nearest_one {
                 let mut level: usize = 0;
                 let mut leaf_idx: usize = 0;
 
-                while level <= self.max_stem_level {
+                while level <= self.max_stem_level as usize {
                     let val = *unsafe { self.stems.get_unchecked(curr_idx) };
                     let is_right_child = *unsafe { query.get_unchecked(dim) } >= val;
 
@@ -32,13 +32,13 @@ macro_rules! generate_immutable_dynamic_approx_nearest_one {
                     dim = (dim + 1) % K;
                 }
 
-                let leaf_extent: core::range::Range<usize> = unsafe { *self.leaf_extents.get_unchecked(leaf_idx) };
+                let (start, end) = unsafe { *self.leaf_extents.get_unchecked(leaf_idx) };
 
                 let leaf_slice = $crate::float_leaf_slice::leaf_slice::LeafSlice::new(
                     array_init::array_init(|i|
-                        &self.leaf_points[i][leaf_extent]
+                        &self.leaf_points[i][start as usize..end as usize]
                     ),
-                    &self.leaf_items[leaf_extent],
+                    &self.leaf_items[start as usize..end as usize],
                 );
 
                 leaf_slice.nearest_one::<D>(

src/immutable_dynamic/common/generate_immutable_dynamic_nearest_one.rs

@@ -10,19 +10,24 @@ macro_rules! generate_immutable_dynamic_nearest_one {
                     D: DistanceMetric<A, K>,
             {
                 let mut off = [A::zero(); K];
+                let mut result = NearestNeighbour {
+                    distance: A::max_value(),
+                    item: T::zero(),
+                };
+
                 self.nearest_one_recurse::<D>(
                     query,
                     0,
                     0,
-                    NearestNeighbour {
-                        distance: A::max_value(),
-                        item: T::zero(),
-                    },
+                    &mut result,
                     &mut off,
                     A::zero(),
                     0,
+                    // 0,
                     0,
-                )
+                );
+
+                result
             }
 
             #[allow(clippy::too_many_arguments)]
@@ -31,21 +36,22 @@ macro_rules! generate_immutable_dynamic_nearest_one {
                 query: &[A; K],
                 stem_idx: usize,
                 split_dim: usize,
-                mut nearest: NearestNeighbour<A, T>,
+                nearest: &mut NearestNeighbour<A, T>,
                 off: &mut [A; K],
                 rd: A,
                 mut level: usize,
+                // mut minor_level: u64,
                 mut leaf_idx: usize,
-            ) -> NearestNeighbour<A, T>
+            )
                 where
                     D: DistanceMetric<A, K>,
             {
                 // use cmov::Cmov;
                 use $crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless;
 
-                if level > self.max_stem_level {
-                    self.search_leaf_for_nearest::<D>(query, &mut nearest, leaf_idx as usize);
-                    return nearest;
+                if level > self.max_stem_level as usize {
+                    self.search_leaf_for_nearest::<D>(query, nearest, leaf_idx as usize);
+                    return;
                 }
 
                 let val = *unsafe { self.stems.get_unchecked(stem_idx as usize) };
@@ -55,85 +61,76 @@ macro_rules! generate_immutable_dynamic_nearest_one {
                 let closer_leaf_idx = leaf_idx + is_right_child;
                 let farther_leaf_idx = leaf_idx + (1 - is_right_child);
 
-                let closer_node_idx = modified_van_emde_boas_get_child_idx_v2_branchless(stem_idx, is_right_child == 1, level);
-                let further_node_idx =  modified_van_emde_boas_get_child_idx_v2_branchless(stem_idx, is_right_child == 0, level);
+                let closer_node_idx = modified_van_emde_boas_get_child_idx_v2_branchless(stem_idx, is_right_child == 1, /*minor_*/level);
+                let further_node_idx =  modified_van_emde_boas_get_child_idx_v2_branchless(stem_idx, is_right_child == 0, /*minor_*/level);
 
                 let mut rd = rd;
                 let old_off = off[split_dim];
                 let new_off = query[split_dim].saturating_dist(val);
 
                 level += 1;
                 let next_split_dim = (split_dim + 1).rem(K);
+                // minor_level += 1;
+                // minor_level.cmovnz(&0, u8::from(minor_level == 3));
 
-                let nearest_neighbour = self.nearest_one_recurse::<D>(
+                self.nearest_one_recurse::<D>(
                     query,
                     closer_node_idx,
                     next_split_dim,
                     nearest,
                     off,
                     rd,
                     level,
+                    // minor_level,
                     closer_leaf_idx,
                 );
 
-                if nearest_neighbour < nearest {
-                    nearest = nearest_neighbour;
-                }
-
                 rd = Axis::rd_update(rd, D::dist1(new_off, old_off));
 
                 if rd <= nearest.distance {
                     off[split_dim] = new_off;
-
-                    let result = self.nearest_one_recurse::<D>(
+                    self.nearest_one_recurse::<D>(
                         query,
                         further_node_idx,
                         next_split_dim,
                         nearest,
                         off,
                         rd,
                         level,
+                        // minor_level,
                         farther_leaf_idx,
                     );
                     off[split_dim] = old_off;
-
-                    if result < nearest {
-                        nearest = result;
-                    }
                 }
-
-                nearest
             }
 
             #[inline]
-             fn search_leaf_for_nearest<D>(
+            fn search_leaf_for_nearest<D>(
                 &self,
                 query: &[A; K],
                 nearest: &mut NearestNeighbour<A, T>,
                 leaf_idx: usize,
             ) where
                 D: DistanceMetric<A, K>,
             {
-                let leaf_extent = unsafe { self.leaf_extents.get_unchecked(leaf_idx) };
-                // let leaf_extent = self.leaf_extents[leaf_idx];
+                let (start, end) = unsafe { *self.leaf_extents.get_unchecked(leaf_idx) };
+
+                // Artificially extend size to be at least chunk length for faster processing
+                // TODO: why does this slow things down?
+                // let end = end.max(start + 32).min(self.leaf_items.len() as u32);
+
                 let leaf_slice = $crate::float_leaf_slice::leaf_slice::LeafSlice::new(
                     array_init::array_init(|i|
-                        &self.leaf_points[i][leaf_extent.start as usize..leaf_extent.end as usize]
+                        &self.leaf_points[i][start as usize..end as usize]
                     ),
-                    &self.leaf_items[leaf_extent.start as usize..leaf_extent.end as usize],
+                    &self.leaf_items[start as usize..end as usize],
                 );
 
-                let mut best_item = nearest.item;
-                let mut best_dist = nearest.distance;
-
                 leaf_slice.nearest_one::<D>(
                     query,
-                    &mut best_dist,
-                    &mut best_item
+                    &mut nearest.distance,
+                    &mut nearest.item
                 );
-
-                nearest.distance = best_dist;
-                nearest.item = best_item;
             }
         }
     };