From 824515f0670691477196a156af9452ba15864441 Mon Sep 17 00:00:00 2001 From: Scott Donnelly Date: Wed, 8 Jan 2025 23:41:11 +0000 Subject: [PATCH] feat(wip): rkyv v0.8 --- Cargo.toml | 16 +++ examples/immutable-rkyv_08-deserialize.rs | 72 +++++++++++++ examples/immutable-rkyv_08-serialize.rs | 63 +++++++++++ .../common/generate_immutable_nearest_one.rs | 2 +- src/immutable/float/kdtree.rs | 85 +++++++++++++-- src/immutable/float/mod.rs | 3 + src/immutable/float/query/nearest_one.rs | 20 ++++ src/immutable/float/rkyv_aligned_vec.rs | 101 ++++++++++++++++++ src/lib.rs | 1 + 9 files changed, 355 insertions(+), 8 deletions(-) create mode 100644 examples/immutable-rkyv_08-deserialize.rs create mode 100644 examples/immutable-rkyv_08-serialize.rs create mode 100644 src/immutable/float/rkyv_aligned_vec.rs diff --git a/Cargo.toml b/Cargo.toml index 458f86d..efe527a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,6 +95,11 @@ optional = true default-features = false features = ["alloc", "copy_unsafe", "size_64"] +[dependencies.rkyv_08] +package = "rkyv" +version = "0.8.9" +optional = true + [dependencies.serde] version = "1" optional = true @@ -127,6 +132,7 @@ las = ["dep:las"] serde = ["dep:serde", "serde/derive", "dep:serde_derive", "dep:serde_with", "fixed/serde", "aligned-vec/serde"] simd = [] rkyv = ["dep:rkyv"] +rkyv_08 = ["dep:rkyv_08"] test_utils = ["dep:rand", "dep:rand_chacha", "dep:rayon"] tracing = ["dep:tracing", "dep:tracing-subscriber"] @@ -215,6 +221,16 @@ name = "immutable-rkyv-deserialize" path = "examples/immutable-rkyv-deserialize.rs" required-features = ["rkyv"] +[[example]] +name = "immutable-rkyv_08-serialize" +path = "examples/immutable-rkyv_08-serialize.rs" +required-features = ["rkyv_08"] + +[[example]] +name = "immutable-rkyv_08-deserialize" +path = "examples/immutable-rkyv_08-deserialize.rs" +required-features = ["rkyv_08"] + [[example]] name = "pointcloud-las" path = "examples/pointcloud-las.rs" diff --git a/examples/immutable-rkyv_08-deserialize.rs b/examples/immutable-rkyv_08-deserialize.rs new file mode 100644 index 0000000..b3ab8e1 --- /dev/null +++ b/examples/immutable-rkyv_08-deserialize.rs @@ -0,0 +1,72 @@ +use elapsed::ElapsedDuration; +use memmap::MmapOptions; +use std::error::Error; +use std::fs::File; +use std::time::Instant; +#[cfg(feature = "tracing")] +use tracing::Level; +#[cfg(feature = "tracing")] +use tracing_subscriber::fmt; + +use rkyv_08::{from_bytes_unchecked, rancor::Error as RkyvError}; + +// use kiddo::immutable::float::kdtree::ArchivedImmutableKdTree; +use kiddo::immutable::float::kdtree::ImmutableKdTree; +use kiddo::SquaredEuclidean; + +type Tree = ImmutableKdTree; + +fn main() -> Result<(), Box> +where +{ + #[cfg(feature = "tracing")] + let subscriber = fmt().with_max_level(Level::TRACE).without_time().finish(); + #[cfg(feature = "tracing")] + tracing::subscriber::set_global_default(subscriber)?; + + let query = [0.123f64, 0.456f64, 0.789f64]; + + let start = Instant::now(); + + // memmap the file into a buffer + let buf = + unsafe { MmapOptions::new().map(&File::open("./examples/immutable-test-tree-r08.rkyv")?)? }; + + // TODO: unsatisfied trait bounds when trying to call nearest_one on archived tree + + // safe API + // let archived_tree = rkyv_08::access::, RkyvError>(&buf[..]).unwrap(); + + // faster unsafe API + // let archived_tree = + // unsafe { rkyv_08::access_unchecked::>(&buf) }; + + // perform a query + // let nearest_neighbour = archived_tree.nearest_one::(&query); + + // println!( + // "total elapsed: {}\n\n", + // ElapsedDuration::new(start.elapsed()) + // ); + // println!( + // "Nearest item to query (archived): {:?}", + // nearest_neighbour.item + // ); + + // full deserialization + let tree = unsafe { from_bytes_unchecked::(&buf) }?; + + // perform a query + let nearest_neighbour = tree.nearest_one::(&query); + + println!( + "Nearest item to query (deserialized): {:?}", + nearest_neighbour.item + ); + println!( + "total elapsed: {}\n\n", + ElapsedDuration::new(start.elapsed()) + ); + + Ok(()) +} diff --git a/examples/immutable-rkyv_08-serialize.rs b/examples/immutable-rkyv_08-serialize.rs new file mode 100644 index 0000000..ed3f329 --- /dev/null +++ b/examples/immutable-rkyv_08-serialize.rs @@ -0,0 +1,63 @@ +use elapsed::ElapsedDuration; +// use memmap::MmapOptions; +use rand::Rng; +use rand_chacha::rand_core::SeedableRng; +use rkyv_08::{rancor::Error as RkyvError, to_bytes}; +use std::error::Error; +use std::fs::File; +use std::io::Write; +use std::time::Instant; +#[cfg(feature = "tracing")] +use tracing::Level; +#[cfg(feature = "tracing")] +use tracing_subscriber::fmt; +use ubyte::ToByteUnit; + +use kiddo::float::distance::SquaredEuclidean; +use kiddo::immutable::float::kdtree::ImmutableKdTree; + +const NUM_ITEMS: usize = 50_000_000; + +type Tree = ImmutableKdTree; + +fn main() -> Result<(), Box> { + #[cfg(feature = "tracing")] + let subscriber = fmt().with_max_level(Level::TRACE).without_time().finish(); + #[cfg(feature = "tracing")] + tracing::subscriber::set_global_default(subscriber)?; + + let query = [0.123f64, 0.456f64, 0.789f64]; + + // build and serialize a large ImmutableKdTree + let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1); + let content_to_add: Vec<[f64; 3]> = (0..NUM_ITEMS).map(|_| rng.gen::<[f64; 3]>()).collect(); + + let start = Instant::now(); + let tree: Tree = ImmutableKdTree::new_from_slice(&content_to_add); + println!( + "Populated ImmutableKdTree with {} items ({})", + tree.size(), + ElapsedDuration::new(start.elapsed()) + ); + + let nearest_neighbour = tree.nearest_one::(&query); + + println!("Nearest item to query: {:?}", nearest_neighbour.item); + + let start = Instant::now(); + + let buf = to_bytes::(&tree)?; + + let mut file = File::create("./examples/immutable-test-tree-r08.rkyv")?; + file.write_all(&buf) + .expect("Could not write serialized rkyv to file"); + + let file_size = file.metadata().unwrap().len().bytes(); + println!( + "Serialized k-d tree to rkyv file 'immutable-test-tree-r08.rkyv' ({}). File size: {:.2}", + ElapsedDuration::new(start.elapsed()), + file_size + ); + + Ok(()) +} diff --git a/src/immutable/common/generate_immutable_nearest_one.rs b/src/immutable/common/generate_immutable_nearest_one.rs index 421aba6..16dca00 100644 --- a/src/immutable/common/generate_immutable_nearest_one.rs +++ b/src/immutable/common/generate_immutable_nearest_one.rs @@ -72,7 +72,7 @@ macro_rules! generate_immutable_nearest_one { use cmov::Cmov; use $crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless; - if level > self.max_stem_level { + if level > Into::::into(self.max_stem_level) || self.stems.is_empty() { self.search_leaf_for_nearest_one::(query, nearest, leaf_idx as usize); return; } diff --git a/src/immutable/float/kdtree.rs b/src/immutable/float/kdtree.rs index 47dd1ce..ba6593c 100644 --- a/src/immutable/float/kdtree.rs +++ b/src/immutable/float/kdtree.rs @@ -11,11 +11,8 @@ //! As with the vanilla tree, [`f64`] or [`f32`] are supported currently for co-ordinate //! values, or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled -pub use crate::float::kdtree::Axis; -use crate::float_leaf_slice::leaf_slice::{LeafSlice, LeafSliceFloat, LeafSliceFloatChunk}; -#[cfg(feature = "modified_van_emde_boas")] -use crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless; -use crate::traits::Content; +use std::{cmp::PartialEq, fmt::Debug}; + use aligned_vec::{avec, AVec, ConstAlign, CACHELINE_ALIGN}; use array_init::array_init; use az::{Az, Cast}; @@ -25,8 +22,14 @@ use ordered_float::OrderedFloat; use rkyv::vec::ArchivedVec; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use std::cmp::PartialEq; -use std::fmt::Debug; + +pub use crate::float::kdtree::Axis; +use crate::float_leaf_slice::leaf_slice::{LeafSlice, LeafSliceFloat, LeafSliceFloatChunk}; +#[cfg(feature = "rkyv_08")] +use crate::immutable::float::rkyv_aligned_vec::EncodeAVec; +#[cfg(feature = "modified_van_emde_boas")] +use crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless; +use crate::traits::Content; /// Immutable floating point k-d tree /// @@ -45,8 +48,14 @@ use std::fmt::Debug; /// /// A convenient type alias exists for ImmutableKdTree with some sensible defaults set: [`kiddo::ImmutableKdTree`](`crate::ImmutableKdTree`). #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr( + feature = "rkyv_08", + derive(rkyv_08::Archive, rkyv_08::Serialize, rkyv_08::Deserialize) +)] +#[cfg_attr(feature = "rkyv_08", rkyv(crate=rkyv_08))] #[derive(Clone, Debug, PartialEq)] pub struct ImmutableKdTree { + #[cfg_attr(feature = "rkyv_08", rkyv(with = EncodeAVec))] pub(crate) stems: AVec, #[cfg_attr(feature = "serde", serde(with = "crate::custom_serde::array_of_vecs"))] @@ -128,6 +137,37 @@ where } } +#[cfg(feature = "rkyv")] +impl From> + for ImmutableKdTree +where + A: Axis + LeafSliceFloat + LeafSliceFloatChunk, + T: Content, + usize: Cast, +{ + /// Creates an [`ImmutableKdTree`] from an [`ImmutableKdTreeRK`] + fn from(orig: ImmutableKdTreeRK) -> Self { + let ImmutableKdTreeRK { + stems, + leaf_points, + leaf_items, + leaf_extents, + max_stem_level, + } = orig; + + let (ptr, length, capacity) = stems.into_raw_parts(); + let stems = unsafe { AVec::from_raw_parts(ptr, CACHELINE_ALIGN, length, capacity) }; + + ImmutableKdTree { + stems, + leaf_points, + leaf_items, + leaf_extents, + max_stem_level, + } + } +} + /// rkyv zero-copy deserializable version of an `ImmutableKdTree`. /// /// Convert an `ImmutableKdTreeRK` into this in order to perform queries. @@ -212,6 +252,37 @@ where } } +#[cfg(feature = "rkyv_08")] +impl ArchivedImmutableKdTree +where + A: Axis + LeafSliceFloat + LeafSliceFloatChunk + rkyv_08::Archive, + T: Content + rkyv_08::Archive, + usize: Cast, +{ + /// Returns the current number of elements stored in the tree + #[inline] + pub fn size(&self) -> usize { + self.leaf_items.len() + } + + /// Returns a LeafSlice for a given leaf index + #[inline] + pub(crate) fn get_leaf_slice(&self, leaf_idx: usize) -> LeafSlice { + let extents = unsafe { self.leaf_extents.get_unchecked(leaf_idx) }; + let start = Into::::into(extents.0) as usize; + let end = Into::::into(extents.1) as usize; + + // 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); + + LeafSlice::new( + array_init::array_init(|i| &self.leaf_points[i][start..end]), + &self.leaf_items[start..end], + ) + } +} + impl From<&[[A; K]]> for ImmutableKdTree where diff --git a/src/immutable/float/mod.rs b/src/immutable/float/mod.rs index 5d7cc9f..23f3881 100644 --- a/src/immutable/float/mod.rs +++ b/src/immutable/float/mod.rs @@ -24,3 +24,6 @@ pub mod kdtree; #[doc(hidden)] pub mod query; + +#[cfg(feature = "rkyv_08")] +mod rkyv_aligned_vec; diff --git a/src/immutable/float/query/nearest_one.rs b/src/immutable/float/query/nearest_one.rs index bf0d7a0..28723d0 100644 --- a/src/immutable/float/query/nearest_one.rs +++ b/src/immutable/float/query/nearest_one.rs @@ -71,6 +71,26 @@ where ); } +#[cfg(feature = "rkyv_08")] +use crate::immutable::float::kdtree::ArchivedImmutableKdTree; +#[cfg(feature = "rkyv_08")] +impl ArchivedImmutableKdTree +where + A: Axis + LeafSliceFloat + LeafSliceFloatChunk + rkyv_08::Archive, + T: Content + rkyv_08::Archive, + usize: Cast, +{ + generate_immutable_float_nearest_one!( + "use std::fs::File; + use memmap::MmapOptions; + + use kiddo::immutable::float::kdtree::AlignedArchivedImmutableKdTree; + + let mmap = unsafe { MmapOptions::new().map(&File::open(\"./examples/immutable-doctest-tree.rkyv\").unwrap()).unwrap() }; + let tree = unsafe { access_unchecked::>>(&buf) };" + ); +} + #[cfg(test)] mod tests { use crate::float::distance::SquaredEuclidean; diff --git a/src/immutable/float/rkyv_aligned_vec.rs b/src/immutable/float/rkyv_aligned_vec.rs new file mode 100644 index 0000000..6130469 --- /dev/null +++ b/src/immutable/float/rkyv_aligned_vec.rs @@ -0,0 +1,101 @@ +use aligned_vec::{AVec, CACHELINE_ALIGN}; +use rkyv_08::rancor::Fallible; +use rkyv_08::with::{ArchiveWith, DeserializeWith, SerializeWith}; +use rkyv_08::{ + ser::{Allocator, Writer}, + vec::{ArchivedVec, VecResolver}, + Archive, Place, Serialize, +}; +use std::marker::PhantomData; + +pub(crate) struct EncodeAVec { + _p: PhantomData, +} + +pub(crate) struct AVecResolver { + len: usize, + inner: VecResolver, +} + +impl ArchiveWith> for EncodeAVec { + type Archived = ArchivedVec; + type Resolver = AVecResolver; + + fn resolve_with(_: &AVec, resolver: Self::Resolver, out: Place) { + ArchivedVec::resolve_from_len(resolver.len, resolver.inner, out); + } +} + +impl, S> SerializeWith, S> for EncodeAVec +where + S: Fallible + Allocator + Writer + ?Sized, +{ + fn serialize_with(avec: &AVec, serializer: &mut S) -> Result { + Ok(AVecResolver { + len: avec.len(), + inner: ArchivedVec::serialize_from_slice(avec.as_slice(), serializer)?, + }) + } +} + +impl DeserializeWith, AVec, D> for EncodeAVec +where + T: Archive + Clone, + D: Fallible + ?Sized, +{ + fn deserialize_with(vec: &ArchivedVec, _: &mut D) -> Result, D::Error> { + let mut result = AVec::with_capacity(CACHELINE_ALIGN, vec.len()); + + for item in vec.as_slice() { + result.push(item.clone()); + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use aligned_vec::{AVec, CACHELINE_ALIGN}; + use std::ops::Rem; + + use rkyv_08::{ + access_unchecked, deserialize, rancor::Error, Archive, Archived, Deserialize, Serialize, + }; + + use crate::immutable::float::rkyv_aligned_vec::EncodeAVec; + + #[test] + fn roundtrip_avec() { + #[derive(Archive, Debug, Serialize, Deserialize, PartialEq)] + #[rkyv(crate=rkyv_08)] + struct Obj { + #[rkyv(with = EncodeAVec)] + pub inner: AVec, + } + + let original = Obj { + inner: AVec::from_slice(CACHELINE_ALIGN, &[10, 20, 30, 40]), + }; + + let buf = rkyv_08::to_bytes::(&original).unwrap(); + + // check that the archived values are the same + let archived: &ArchivedObj = unsafe { access_unchecked::>(buf.as_ref()) }; + assert_eq!(original.inner.as_slice(), archived.inner.as_slice()); + + // check that the deserialized values are the same + let deserialized: Obj = deserialize::(archived).unwrap(); + assert_eq!(original, deserialized); + + // check that the deserialized AVec inside Obj has the alignment we expect + let inner_ptr_adr_deser = deserialized.inner.as_ptr() as *const () as usize; + assert_eq!(inner_ptr_adr_deser.rem(CACHELINE_ALIGN), 0); + + // check that the archived AVec has the alignment we expect + // NOTE: right now this works sometimes but not others - it + // seems like the Archived AVec items are 32-byte aligned rather than 128 + let inner_ptr_adr_archived = archived.inner.as_ptr() as *const () as usize; + assert_eq!(inner_ptr_adr_archived.rem(CACHELINE_ALIGN), 0); + } +} diff --git a/src/lib.rs b/src/lib.rs index 5647348..979cc21 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ +#![feature(vec_into_raw_parts)] #![cfg_attr(feature = "simd", feature(slice_as_chunks))] #![warn(rustdoc::missing_crate_level_docs)] #![deny(rustdoc::invalid_codeblock_attributes)]