/rust/registry/src/index.crates.io-1949cf8c6b5b557f/diskann-vector-0.54.0/src/conversion.rs

Source
/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use std::convert::{AsMut, AsRef};

use diskann_wide::{arch::Target2, Architecture, Const, Constant, SIMDCast, SIMDVector};
use half::f16;

/// Perform a numeric cast on a slice of values.
///
/// This trait is intended to have the following numerical behavior:
///
/// 1. If a lossless conversion between types is available, use that.
/// 2. Otherwise, if the two type are floating point types, use a round-to-nearest strategy.
/// 3. Otherwise, try to behave like the Rust `as` numeric cast.
///
/// The main reason we can't just say "behave like "as"" is because Rust does not have
/// a native `f16` type, which this crate supports.
pub trait CastFromSlice<From> {
    fn cast_from_slice(self, from: From);
}

macro_rules! use_simd_cast_from_slice {
    ($from:ty => $to:ty) => {
        impl CastFromSlice<&[$from]> for &mut [$to] {
            #[inline(always)]
            fn cast_from_slice(self, from: &[$from]) {
                SliceCast::<$to, $from>::new().run(diskann_wide::ARCH, self, from)
            }
        }

        impl<const N: usize> CastFromSlice<&[$from; N]> for &mut [$to; N] {
            #[inline(always)]
            fn cast_from_slice(self, from: &[$from; N]) {
                SliceCast::<$to, $from>::new().run(diskann_wide::ARCH, self, from)
            }
        }
    };
}

use_simd_cast_from_slice!(f32 => f16);
use_simd_cast_from_slice!(f16 => f32);

/// A zero-sized type providing implementations of [`diskann_wide::arch::Target2`] to provide
/// platform-dependent conversions between slices of the two generic types.
#[derive(Debug, Default, Clone, Copy)]
pub struct SliceCast<To, From> {
    _marker: std::marker::PhantomData<(To, From)>,
}

impl<To, From> SliceCast<To, From> {
    pub fn new() -> Self {
        Self {
            _marker: std::marker::PhantomData,
        }
    }
}

// Non-SIMD Instantiations

impl<T, U> Target2<diskann_wide::arch::Scalar, (), T, U> for SliceCast<f16, f32>
where
    T: AsMut<[f16]>,
    U: AsRef<[f32]>,
{
    #[inline(always)]
    fn run(self, _: diskann_wide::arch::Scalar, mut to: T, from: U) {
        let to = to.as_mut();
        let from = from.as_ref();
        std::iter::zip(to.iter_mut(), from.iter()).for_each(|(to, from)| {
            *to = diskann_wide::cast_f32_to_f16(*from);
        })
    }
}

impl<T, U> Target2<diskann_wide::arch::Scalar, (), T, U> for SliceCast<f32, f16>
where
    T: AsMut<[f32]>,
    U: AsRef<[f16]>,
{
    #[inline(always)]
    fn run(self, _: diskann_wide::arch::Scalar, mut to: T, from: U) {
        let to = to.as_mut();
        let from = from.as_ref();
        std::iter::zip(to.iter_mut(), from.iter()).for_each(|(to, from)| {
            *to = diskann_wide::cast_f16_to_f32(*from);
        })
    }
}

// SIMD Instantiations
#[cfg(target_arch = "x86_64")]
impl<T, U, To, From> Target2<diskann_wide::arch::x86_64::V4, (), T, U> for SliceCast<To, From>
where
    T: AsMut<[To]>,
    U: AsRef<[From]>,
    diskann_wide::arch::x86_64::V4: SIMDConvert<To, From>,
{
    #[inline(always)]
    fn run(self, arch: diskann_wide::arch::x86_64::V4, mut to: T, from: U) {
        simd_convert(arch, to.as_mut(), from.as_ref())
    }
}

#[cfg(target_arch = "x86_64")]
impl<T, U, To, From> Target2<diskann_wide::arch::x86_64::V3, (), T, U> for SliceCast<To, From>
where
    T: AsMut<[To]>,
    U: AsRef<[From]>,
    diskann_wide::arch::x86_64::V3: SIMDConvert<To, From>,
{
    #[inline(always)]
    fn run(self, arch: diskann_wide::arch::x86_64::V3, mut to: T, from: U) {
        simd_convert(arch, to.as_mut(), from.as_ref())
    }
}

#[cfg(target_arch = "aarch64")]
impl<T, U, To, From> Target2<diskann_wide::arch::aarch64::Neon, (), T, U> for SliceCast<To, From>
where
    T: AsMut<[To]>,
    U: AsRef<[From]>,
    diskann_wide::arch::aarch64::Neon: SIMDConvert<To, From>,
{
    #[inline(always)]
    fn run(self, arch: diskann_wide::arch::aarch64::Neon, mut to: T, from: U) {
        simd_convert(arch, to.as_mut(), from.as_ref())
    }
}

/////////////////////////////
// General SIMD Conversion //
/////////////////////////////

/// A helper trait to fill in the gaps for the unrolled `simd_convert` method.
trait SIMDConvert<To, From>: Architecture {
    /// A constant encoding the the SIMD width of the underlying schema.
    type Width: Constant<Type = usize>;

    /// The SIMD Vector for the converted-to type.
    type WideTo: SIMDVector<Arch = Self, Scalar = To, ConstLanes = Self::Width>;

    /// The SIMD Vector for the converted-from type.
    type WideFrom: SIMDVector<Arch = Self, Scalar = From, ConstLanes = Self::Width>;

    /// The method that actually does the vector-wide conversion.
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo;

    /// Delegate routing for handling conversion lengths less than the vector width.
    ///
    /// The canonical implementation uses predicated loads, but implementations may wish
    /// to use a scalar loop instead.
    ///
    /// # Safety
    ///
    /// This trait will only be called when the following guarantees are made:
    ///
    /// * `pto` will point to properly aligned memory that is valid for writes on the
    ///   range `[pto, pto + len)`.
    /// * `pfrom` will point to properly aligned memory that is valid for reads on the
    ///   range `[pfrom, pfrom + len)`.
    /// * The memory ranges covered by `pto` and `pfrom` must not alias.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut To, pfrom: *const From, len: usize) {
        let from = Self::WideFrom::load_simd_first(self, pfrom, len);
        let to = Self::simd_convert(from);
        to.store_simd_first(pto, len);
    }

    /// !! Do not extend this function !!
    ///
    /// Due to limitations on how associated constants can be used, we need a function
    /// to access the SIMD width and rely on the compiler to constant propagate the result.
    #[inline(always)]
    fn get_simd_width() -> usize {
        Self::Width::value()
    }
}

#[inline(never)]
#[allow(clippy::panic)]
fn emit_length_error(xlen: usize, ylen: usize) -> ! {
    panic!(
        "lengths must be equal, instead got: xlen = {}, ylen = {}",
        xlen, ylen
    )
}

/// Convert each element of `from` into its corresponding position in `to` using the
/// conversion rule applied by `S`.
///
/// # Panics
///
/// Panics if `to.len() != from.len()`.
///
/// # Implementation Notes
///
/// This function will only call `A::handle_small` if the total length of the processed
/// slices is less that the underlying SIMD width.
///
/// Otherwise, we take advantage of unaligned operations to avoid dealing with
/// non-full-width chunks.
///
/// For example, if the SIMD width was 4 and the total length was 7, then it would be
/// processed in two chunks of 4 like so:
/// ```text
///      Chunk 0
/// |---------------|
///   0   1   2   3   4   5   6
///             |---------------|
///                   Chunk 1
/// ```
/// This overlapping can only happen at the very end of the slice and only if the length
/// of the slice is not a multiple of the SIMD width used.
#[inline(always)]
fn simd_convert<A, To, From>(arch: A, to: &mut [To], from: &[From])
where
    A: SIMDConvert<To, From>,
{
    let len = to.len();

    // Keep stack writes to a minimum by explicitly outlining error handling.
    if len != from.len() {
        emit_length_error(len, from.len())
    }

    // Get the SIMD width.
    //
    // We're relying on the compiler to constant propagate this.
    let width = A::get_simd_width();

    let pto = to.as_mut_ptr();
    let pfrom = from.as_ptr();

    // Too short, deal with the small case and return.
    if len < width {
        // SAFETY: We know `pto` and `pfrom` do not alias because of Rust's aliasing
        // rules on `to` and `from.
        //
        // Additionally, we've checked that both spans are valid for `len`.
        unsafe { arch.handle_small(pto, pfrom, len) };
        return;
    }

    const UNROLL: usize = 8;

    let mut i = 0;
    // SAFETY: We emit a bunch of unrolled load and store operations in this loop.
    //
    // All of these operations are safe because the bound `i + UNROLL * width <= len`
    // is checked.
    unsafe {
        while i + UNROLL * width <= len {
            let s0 = A::WideFrom::load_simd(arch, pfrom.add(i));
            A::simd_convert(s0).store_simd(pto.add(i));

            let s1 = A::WideFrom::load_simd(arch, pfrom.add(i + width));
            A::simd_convert(s1).store_simd(pto.add(i + width));

            let s2 = A::WideFrom::load_simd(arch, pfrom.add(i + 2 * width));
            A::simd_convert(s2).store_simd(pto.add(i + 2 * width));

            let s3 = A::WideFrom::load_simd(arch, pfrom.add(i + 3 * width));
            A::simd_convert(s3).store_simd(pto.add(i + 3 * width));

            let s0 = A::WideFrom::load_simd(arch, pfrom.add(i + 4 * width));
            A::simd_convert(s0).store_simd(pto.add(i + 4 * width));

            let s1 = A::WideFrom::load_simd(arch, pfrom.add(i + 5 * width));
            A::simd_convert(s1).store_simd(pto.add(i + 5 * width));

            let s2 = A::WideFrom::load_simd(arch, pfrom.add(i + 6 * width));
            A::simd_convert(s2).store_simd(pto.add(i + 6 * width));

            let s3 = A::WideFrom::load_simd(arch, pfrom.add(i + 7 * width));
            A::simd_convert(s3).store_simd(pto.add(i + 7 * width));

            i += UNROLL * width;
        }
    }

    while i + width <= len {
        // SAFETY: `i + width <= len` ensure that this read is in-bounds.
        let s0 = unsafe { A::WideFrom::load_simd(arch, pfrom.add(i)) };
        let t0 = A::simd_convert(s0);
        // SAFETY: `i + width <= len` ensure that this write is in-bounds.
        unsafe { t0.store_simd(pto.add(i)) };
        i += width;
    }

    // Check if we need to deal with any remaining elements.
    // If so, bump back `i` so we can process a whole chunk.
    if i != len {
        let offset = i - (width - len % width);

        // SAFETY: At this point, we know that `len >= width`, `i < len`, and
        // `len - i == len % width != 0`.
        //
        // Therefore, `offset` is inbounds and `offset + width == len`.
        let s0 = unsafe { A::WideFrom::load_simd(arch, pfrom.add(offset)) };
        let t0 = A::simd_convert(s0);

        // SAFETY: This write is safe for the same reason that the preceeding read is safe.
        unsafe { t0.store_simd(pto.add(offset)) };
    }
}

#[cfg(target_arch = "x86_64")]
impl SIMDConvert<f32, f16> for diskann_wide::arch::x86_64::V4 {
    type Width = Const<8>;
    type WideTo = <diskann_wide::arch::x86_64::V4 as Architecture>::f32x8;
    type WideFrom = <diskann_wide::arch::x86_64::V4 as Architecture>::f16x8;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.into()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f32, pfrom: *const f16, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f16_to_f32(*pfrom.add(i))
        }
    }
}

#[cfg(target_arch = "x86_64")]
impl SIMDConvert<f32, f16> for diskann_wide::arch::x86_64::V3 {
    type Width = Const<8>;
    type WideTo = <diskann_wide::arch::x86_64::V3 as Architecture>::f32x8;
    type WideFrom = <diskann_wide::arch::x86_64::V3 as Architecture>::f16x8;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.into()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f32, pfrom: *const f16, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f16_to_f32(*pfrom.add(i))
        }
    }
}

#[cfg(target_arch = "x86_64")]
impl SIMDConvert<f16, f32> for diskann_wide::arch::x86_64::V4 {
    type Width = Const<8>;
    type WideTo = <diskann_wide::arch::x86_64::V4 as Architecture>::f16x8;
    type WideFrom = <diskann_wide::arch::x86_64::V4 as Architecture>::f32x8;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.simd_cast()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f16, pfrom: *const f32, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f32_to_f16(*pfrom.add(i))
        }
    }
}

#[cfg(target_arch = "x86_64")]
impl SIMDConvert<f16, f32> for diskann_wide::arch::x86_64::V3 {
    type Width = Const<8>;
    type WideTo = <diskann_wide::arch::x86_64::V3 as Architecture>::f16x8;
    type WideFrom = <diskann_wide::arch::x86_64::V3 as Architecture>::f32x8;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.simd_cast()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f16, pfrom: *const f32, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f32_to_f16(*pfrom.add(i))
        }
    }
}

//---------//
// Aarch64 //
//---------//

#[cfg(target_arch = "aarch64")]
impl SIMDConvert<f32, f16> for diskann_wide::arch::aarch64::Neon {
    type Width = Const<4>;
    type WideTo = <diskann_wide::arch::aarch64::Neon as Architecture>::f32x4;
    type WideFrom = diskann_wide::arch::aarch64::f16x4;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.into()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f32, pfrom: *const f16, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f16_to_f32(*pfrom.add(i))
        }
    }
}

#[cfg(target_arch = "aarch64")]
impl SIMDConvert<f16, f32> for diskann_wide::arch::aarch64::Neon {
    type Width = Const<4>;
    type WideTo = diskann_wide::arch::aarch64::f16x4;
    type WideFrom = <diskann_wide::arch::aarch64::Neon as Architecture>::f32x4;

    #[inline(always)]
    fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
        from.simd_cast()
    }

    // SAFETY: We only access data in the valid range for `pto` and `pfrom`.
    #[inline(always)]
    unsafe fn handle_small(self, pto: *mut f16, pfrom: *const f32, len: usize) {
        for i in 0..len {
            *pto.add(i) = diskann_wide::cast_f32_to_f16(*pfrom.add(i))
        }
    }
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use rand::{
        distr::{Distribution, StandardUniform},
        rngs::StdRng,
        SeedableRng,
    };

    use super::*;

    ////////////////
    // Fuzz tests //
    ////////////////

    trait ReferenceConvert<From> {
        fn reference_convert(self, from: &[From]);
    }

    impl ReferenceConvert<f32> for &mut [f16] {
        fn reference_convert(self, from: &[f32]) {
            assert_eq!(self.len(), from.len());
            std::iter::zip(self.iter_mut(), from.iter()).for_each(|(d, s)| *d = f16::from_f32(*s));
        }
    }

    impl ReferenceConvert<f16> for &mut [f32] {
        fn reference_convert(self, from: &[f16]) {
            assert_eq!(self.len(), from.len());
            std::iter::zip(self.iter_mut(), from.iter()).for_each(|(d, s)| *d = (*s).into());
        }
    }

    fn test_cast_from_slice<To, From>(max_dim: usize, num_trials: usize, rng: &mut StdRng)
    where
        StandardUniform: Distribution<From>,
        To: Default + PartialEq + std::fmt::Debug + Copy,
        From: Default + Copy,
        for<'a, 'b> &'a mut [To]: CastFromSlice<&'b [From]> + ReferenceConvert<From>,
    {
        let distribution = StandardUniform {};
        for dim in 0..=max_dim {
            let mut src = vec![From::default(); dim];
            let mut dst = vec![To::default(); dim];
            let mut dst_reference = vec![To::default(); dim];

            for _ in 0..num_trials {
                src.iter_mut().for_each(|s| *s = distribution.sample(rng));
                dst.cast_from_slice(src.as_slice());
                dst_reference.reference_convert(&src);

                assert_eq!(dst, dst_reference);
            }
        }
    }

    #[test]
    fn test_f32_to_f16_fuzz() {
        let mut rng = StdRng::seed_from_u64(0x0a3bfe052a8ebf98);
        test_cast_from_slice::<f16, f32>(256, 10, &mut rng);
    }

    #[test]
    fn test_f16_to_f32_fuzz() {
        let mut rng = StdRng::seed_from_u64(0x83765b2816321eca);
        test_cast_from_slice::<f32, f16>(256, 10, &mut rng);
    }

    ////////////////
    // Miri Tests //
    ////////////////

    // With Miri, we need to be really careful that we do not hit methods that call
    // `cvtph_ps`...
    #[test]
    fn miri_test_f32_to_f16() {
        for dim in 0..256 {
            println!("processing dim {}", dim);

            let src = vec![f32::default(); dim];
            let mut dst = vec![f16::default(); dim];

            // Just check that all accesses are inbounds.
            dst.cast_from_slice(&src);

            // Scalar conversion.
            SliceCast::<f16, f32>::new().run(
                diskann_wide::arch::Scalar,
                dst.as_mut_slice(),
                src.as_slice(),
            );

            // SIMD conversion
            #[cfg(target_arch = "x86_64")]
            {
                if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() {
                    SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
                }

                if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() {
                    SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
                }
            }

            #[cfg(target_arch = "aarch64")]
            if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() {
                SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
            }
        }
    }

    #[test]
    fn miri_test_f16_to_f32() {
        for dim in 0..256 {
            println!("processing dim {}", dim);

            let src = vec![f16::default(); dim];
            let mut dst = vec![f32::default(); dim];

            // Just check that all accesses are inbounds.
            dst.cast_from_slice(&src);

            // Scalar conversion.
            SliceCast::<f32, f16>::new().run(
                diskann_wide::arch::Scalar,
                dst.as_mut_slice(),
                src.as_slice(),
            );

            // SIMD conversion
            #[cfg(target_arch = "x86_64")]
            {
                if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() {
                    SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
                }

                if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() {
                    SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
                }
            }

            #[cfg(target_arch = "aarch64")]
            if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() {
                SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
            }
        }
    }
}

Coverage Report

Created: 2026-11-03 08:22

Line	Count	Source
1		/*
2		* Copyright (c) Microsoft Corporation.
3		* Licensed under the MIT license.
4		*/
5
6		use std::convert::{AsMut, AsRef};
7
8		use diskann_wide::{arch::Target2, Architecture, Const, Constant, SIMDCast, SIMDVector};
9		use half::f16;
10
11		/// Perform a numeric cast on a slice of values.
12		///
13		/// This trait is intended to have the following numerical behavior:
14		///
15		/// 1. If a lossless conversion between types is available, use that.
16		/// 2. Otherwise, if the two type are floating point types, use a round-to-nearest strategy.
17		/// 3. Otherwise, try to behave like the Rust `as` numeric cast.
18		///
19		/// The main reason we can't just say "behave like "as"" is because Rust does not have
20		/// a native `f16` type, which this crate supports.
21		pub trait CastFromSlice<From> {
22		fn cast_from_slice(self, from: From);
23		}
24
25		macro_rules! use_simd_cast_from_slice {
26		($from:ty => $to:ty) => {
27		impl CastFromSlice<&[$from]> for &mut [$to] {
28		#[inline(always)]
29	0	fn cast_from_slice(self, from: &[$from]) {
30	0	SliceCast::<$to, $from>::new().run(diskann_wide::ARCH, self, from)
31	0	} Unexecuted instantiation: <&mut [f32] as diskann_vector::conversion::CastFromSlice<&[half::binary16::f16]>>::cast_from_slice Unexecuted instantiation: <&mut [half::binary16::f16] as diskann_vector::conversion::CastFromSlice<&[f32]>>::cast_from_slice
32		}
33
34		impl<const N: usize> CastFromSlice<&[$from; N]> for &mut [$to; N] {
35		#[inline(always)]
36	0	fn cast_from_slice(self, from: &[$from; N]) {
37	0	SliceCast::<$to, $from>::new().run(diskann_wide::ARCH, self, from)
38	0	} Unexecuted instantiation: <&mut [half::binary16::f16; _] as diskann_vector::conversion::CastFromSlice<&[f32; _]>>::cast_from_slice Unexecuted instantiation: <&mut [f32; _] as diskann_vector::conversion::CastFromSlice<&[half::binary16::f16; _]>>::cast_from_slice
39		}
40		};
41		}
42
43		use_simd_cast_from_slice!(f32 => f16);
44		use_simd_cast_from_slice!(f16 => f32);
45
46		/// A zero-sized type providing implementations of [`diskann_wide::arch::Target2`] to provide
47		/// platform-dependent conversions between slices of the two generic types.
48		#[derive(Debug, Default, Clone, Copy)]
49		pub struct SliceCast<To, From> {
50		_marker: std::marker::PhantomData<(To, From)>,
51		}
52
53		impl<To, From> SliceCast<To, From> {
54	0	pub fn new() -> Self {
55	0	Self {
56	0	_marker: std::marker::PhantomData,
57	0	}
58	0	} Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16>>::new Unexecuted instantiation: <diskann_vector::conversion::SliceCast<_, _>>::new
59		}
60
61		// Non-SIMD Instantiations
62
63		impl<T, U> Target2<diskann_wide::arch::Scalar, (), T, U> for SliceCast<f16, f32>
64		where
65		T: AsMut<[f16]>,
66		U: AsRef<[f32]>,
67		{
68		#[inline(always)]
69	0	fn run(self, _: diskann_wide::arch::Scalar, mut to: T, from: U) {
70	0	let to = to.as_mut();
71	0	let from = from.as_ref();
72	0	std::iter::zip(to.iter_mut(), from.iter()).for_each(\|(to, from)\| {
73	0	to = diskann_wide::cast_f32_to_f16(from);
74	0	})
75	0	}
76		}
77
78		impl<T, U> Target2<diskann_wide::arch::Scalar, (), T, U> for SliceCast<f32, f16>
79		where
80		T: AsMut<[f32]>,
81		U: AsRef<[f16]>,
82		{
83		#[inline(always)]
84	0	fn run(self, _: diskann_wide::arch::Scalar, mut to: T, from: U) {
85	0	let to = to.as_mut();
86	0	let from = from.as_ref();
87	0	std::iter::zip(to.iter_mut(), from.iter()).for_each(\|(to, from)\| {
88	0	to = diskann_wide::cast_f16_to_f32(from);
89	0	}) Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16> as diskann_wide::arch::Target2<diskann_wide::arch::emulated::Scalar, (), &mut [f32], &[half::binary16::f16]>>::run::{closure#0} Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16> as diskann_wide::arch::Target2<diskann_wide::arch::emulated::Scalar, (), _, _>>::run::{closure#0}
90	0	} Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16> as diskann_wide::arch::Target2<diskann_wide::arch::emulated::Scalar, (), &mut [f32], &[half::binary16::f16]>>::run Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16> as diskann_wide::arch::Target2<diskann_wide::arch::emulated::Scalar, (), &mut [f32], &[half::binary16::f16]>>::run Unexecuted instantiation: <diskann_vector::conversion::SliceCast<f32, half::binary16::f16> as diskann_wide::arch::Target2<diskann_wide::arch::emulated::Scalar, (), _, _>>::run
91		}
92
93		// SIMD Instantiations
94		#[cfg(target_arch = "x86_64")]
95		impl<T, U, To, From> Target2<diskann_wide::arch::x86_64::V4, (), T, U> for SliceCast<To, From>
96		where
97		T: AsMut<[To]>,
98		U: AsRef<[From]>,
99		diskann_wide::arch::x86_64::V4: SIMDConvert<To, From>,
100		{
101		#[inline(always)]
102	0	fn run(self, arch: diskann_wide::arch::x86_64::V4, mut to: T, from: U) {
103	0	simd_convert(arch, to.as_mut(), from.as_ref())
104	0	}
105		}
106
107		#[cfg(target_arch = "x86_64")]
108		impl<T, U, To, From> Target2<diskann_wide::arch::x86_64::V3, (), T, U> for SliceCast<To, From>
109		where
110		T: AsMut<[To]>,
111		U: AsRef<[From]>,
112		diskann_wide::arch::x86_64::V3: SIMDConvert<To, From>,
113		{
114		#[inline(always)]
115	0	fn run(self, arch: diskann_wide::arch::x86_64::V3, mut to: T, from: U) {
116	0	simd_convert(arch, to.as_mut(), from.as_ref())
117	0	}
118		}
119
120		#[cfg(target_arch = "aarch64")]
121		impl<T, U, To, From> Target2<diskann_wide::arch::aarch64::Neon, (), T, U> for SliceCast<To, From>
122		where
123		T: AsMut<[To]>,
124		U: AsRef<[From]>,
125		diskann_wide::arch::aarch64::Neon: SIMDConvert<To, From>,
126		{
127		#[inline(always)]
128		fn run(self, arch: diskann_wide::arch::aarch64::Neon, mut to: T, from: U) {
129		simd_convert(arch, to.as_mut(), from.as_ref())
130		}
131		}
132
133		/////////////////////////////
134		// General SIMD Conversion //
135		/////////////////////////////
136
137		/// A helper trait to fill in the gaps for the unrolled `simd_convert` method.
138		trait SIMDConvert<To, From>: Architecture {
139		/// A constant encoding the the SIMD width of the underlying schema.
140		type Width: Constant<Type = usize>;
141
142		/// The SIMD Vector for the converted-to type.
143		type WideTo: SIMDVector<Arch = Self, Scalar = To, ConstLanes = Self::Width>;
144
145		/// The SIMD Vector for the converted-from type.
146		type WideFrom: SIMDVector<Arch = Self, Scalar = From, ConstLanes = Self::Width>;
147
148		/// The method that actually does the vector-wide conversion.
149		fn simd_convert(from: Self::WideFrom) -> Self::WideTo;
150
151		/// Delegate routing for handling conversion lengths less than the vector width.
152		///
153		/// The canonical implementation uses predicated loads, but implementations may wish
154		/// to use a scalar loop instead.
155		///
156		/// # Safety
157		///
158		/// This trait will only be called when the following guarantees are made:
159		///
160		/// * `pto` will point to properly aligned memory that is valid for writes on the
161		/// range `[pto, pto + len)`.
162		/// * `pfrom` will point to properly aligned memory that is valid for reads on the
163		/// range `[pfrom, pfrom + len)`.
164		/// * The memory ranges covered by `pto` and `pfrom` must not alias.
165		#[inline(always)]
166	0	unsafe fn handle_small(self, pto: mut To, pfrom: const From, len: usize) {
167	0	let from = Self::WideFrom::load_simd_first(self, pfrom, len);
168	0	let to = Self::simd_convert(from);
169	0	to.store_simd_first(pto, len);
170	0	}
171
172		/// !! Do not extend this function !!
173		///
174		/// Due to limitations on how associated constants can be used, we need a function
175		/// to access the SIMD width and rely on the compiler to constant propagate the result.
176		#[inline(always)]
177	0	fn get_simd_width() -> usize {
178	0	Self::Width::value()
179	0	}
180		}
181
182		#[inline(never)]
183		#[allow(clippy::panic)]
184	0	fn emit_length_error(xlen: usize, ylen: usize) -> ! {
185	0	panic!(
186		"lengths must be equal, instead got: xlen = {}, ylen = {}",
187		xlen, ylen
188		)
189		}
190
191		/// Convert each element of `from` into its corresponding position in `to` using the
192		/// conversion rule applied by `S`.
193		///
194		/// # Panics
195		///
196		/// Panics if `to.len() != from.len()`.
197		///
198		/// # Implementation Notes
199		///
200		/// This function will only call `A::handle_small` if the total length of the processed
201		/// slices is less that the underlying SIMD width.
202		///
203		/// Otherwise, we take advantage of unaligned operations to avoid dealing with
204		/// non-full-width chunks.
205		///
206		/// For example, if the SIMD width was 4 and the total length was 7, then it would be
207		/// processed in two chunks of 4 like so:
208		/// ```text
209		/// Chunk 0
210		/// \|---------------\|
211		/// 0 1 2 3 4 5 6
212		/// \|---------------\|
213		/// Chunk 1
214		/// ```
215		/// This overlapping can only happen at the very end of the slice and only if the length
216		/// of the slice is not a multiple of the SIMD width used.
217		#[inline(always)]
218	0	fn simd_convert<A, To, From>(arch: A, to: &mut [To], from: &[From])
219	0	where
220	0	A: SIMDConvert<To, From>,
221		{
222	0	let len = to.len();
223
224		// Keep stack writes to a minimum by explicitly outlining error handling.
225	0	if len != from.len() {
226	0	emit_length_error(len, from.len())
227	0	}
228
229		// Get the SIMD width.
230		//
231		// We're relying on the compiler to constant propagate this.
232	0	let width = A::get_simd_width();
233
234	0	let pto = to.as_mut_ptr();
235	0	let pfrom = from.as_ptr();
236
237		// Too short, deal with the small case and return.
238	0	if len < width {
239		// SAFETY: We know `pto` and `pfrom` do not alias because of Rust's aliasing
240		// rules on `to` and `from.
241		//
242		// Additionally, we've checked that both spans are valid for `len`.
243	0	unsafe { arch.handle_small(pto, pfrom, len) };
244	0	return;
245	0	}
246
247		const UNROLL: usize = 8;
248
249	0	let mut i = 0;
250		// SAFETY: We emit a bunch of unrolled load and store operations in this loop.
251		//
252		// All of these operations are safe because the bound `i + UNROLL * width <= len`
253		// is checked.
254		unsafe {
255	0	while i + UNROLL * width <= len {
256	0	let s0 = A::WideFrom::load_simd(arch, pfrom.add(i));
257	0	A::simd_convert(s0).store_simd(pto.add(i));
258	0
259	0	let s1 = A::WideFrom::load_simd(arch, pfrom.add(i + width));
260	0	A::simd_convert(s1).store_simd(pto.add(i + width));
261	0
262	0	let s2 = A::WideFrom::load_simd(arch, pfrom.add(i + 2 * width));
263	0	A::simd_convert(s2).store_simd(pto.add(i + 2 * width));
264	0
265	0	let s3 = A::WideFrom::load_simd(arch, pfrom.add(i + 3 * width));
266	0	A::simd_convert(s3).store_simd(pto.add(i + 3 * width));
267	0
268	0	let s0 = A::WideFrom::load_simd(arch, pfrom.add(i + 4 * width));
269	0	A::simd_convert(s0).store_simd(pto.add(i + 4 * width));
270	0
271	0	let s1 = A::WideFrom::load_simd(arch, pfrom.add(i + 5 * width));
272	0	A::simd_convert(s1).store_simd(pto.add(i + 5 * width));
273	0
274	0	let s2 = A::WideFrom::load_simd(arch, pfrom.add(i + 6 * width));
275	0	A::simd_convert(s2).store_simd(pto.add(i + 6 * width));
276	0
277	0	let s3 = A::WideFrom::load_simd(arch, pfrom.add(i + 7 * width));
278	0	A::simd_convert(s3).store_simd(pto.add(i + 7 * width));
279	0
280	0	i += UNROLL * width;
281	0	}
282		}
283
284	0	while i + width <= len {
285	0	// SAFETY: `i + width <= len` ensure that this read is in-bounds.
286	0	let s0 = unsafe { A::WideFrom::load_simd(arch, pfrom.add(i)) };
287	0	let t0 = A::simd_convert(s0);
288	0	// SAFETY: `i + width <= len` ensure that this write is in-bounds.
289	0	unsafe { t0.store_simd(pto.add(i)) };
290	0	i += width;
291	0	}
292
293		// Check if we need to deal with any remaining elements.
294		// If so, bump back `i` so we can process a whole chunk.
295	0	if i != len {
296	0	let offset = i - (width - len % width);
297	0
298	0	// SAFETY: At this point, we know that `len >= width`, `i < len`, and
299	0	// `len - i == len % width != 0`.
300	0	//
301	0	// Therefore, `offset` is inbounds and `offset + width == len`.
302	0	let s0 = unsafe { A::WideFrom::load_simd(arch, pfrom.add(offset)) };
303	0	let t0 = A::simd_convert(s0);
304	0
305	0	// SAFETY: This write is safe for the same reason that the preceeding read is safe.
306	0	unsafe { t0.store_simd(pto.add(offset)) };
307	0	}
308	0	}
309
310		#[cfg(target_arch = "x86_64")]
311		impl SIMDConvert<f32, f16> for diskann_wide::arch::x86_64::V4 {
312		type Width = Const<8>;
313		type WideTo = <diskann_wide::arch::x86_64::V4 as Architecture>::f32x8;
314		type WideFrom = <diskann_wide::arch::x86_64::V4 as Architecture>::f16x8;
315
316		#[inline(always)]
317	0	fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
318	0	from.into()
319	0	}
320
321		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
322		#[inline(always)]
323	0	unsafe fn handle_small(self, pto: mut f32, pfrom: const f16, len: usize) {
324	0	for i in 0..len {
325	0	pto.add(i) = diskann_wide::cast_f16_to_f32(pfrom.add(i))
326		}
327	0	}
328		}
329
330		#[cfg(target_arch = "x86_64")]
331		impl SIMDConvert<f32, f16> for diskann_wide::arch::x86_64::V3 {
332		type Width = Const<8>;
333		type WideTo = <diskann_wide::arch::x86_64::V3 as Architecture>::f32x8;
334		type WideFrom = <diskann_wide::arch::x86_64::V3 as Architecture>::f16x8;
335
336		#[inline(always)]
337	0	fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
338	0	from.into()
339	0	}
340
341		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
342		#[inline(always)]
343	0	unsafe fn handle_small(self, pto: mut f32, pfrom: const f16, len: usize) {
344	0	for i in 0..len {
345	0	pto.add(i) = diskann_wide::cast_f16_to_f32(pfrom.add(i))
346		}
347	0	}
348		}
349
350		#[cfg(target_arch = "x86_64")]
351		impl SIMDConvert<f16, f32> for diskann_wide::arch::x86_64::V4 {
352		type Width = Const<8>;
353		type WideTo = <diskann_wide::arch::x86_64::V4 as Architecture>::f16x8;
354		type WideFrom = <diskann_wide::arch::x86_64::V4 as Architecture>::f32x8;
355
356		#[inline(always)]
357	0	fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
358	0	from.simd_cast()
359	0	}
360
361		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
362		#[inline(always)]
363	0	unsafe fn handle_small(self, pto: mut f16, pfrom: const f32, len: usize) {
364	0	for i in 0..len {
365	0	pto.add(i) = diskann_wide::cast_f32_to_f16(pfrom.add(i))
366		}
367	0	}
368		}
369
370		#[cfg(target_arch = "x86_64")]
371		impl SIMDConvert<f16, f32> for diskann_wide::arch::x86_64::V3 {
372		type Width = Const<8>;
373		type WideTo = <diskann_wide::arch::x86_64::V3 as Architecture>::f16x8;
374		type WideFrom = <diskann_wide::arch::x86_64::V3 as Architecture>::f32x8;
375
376		#[inline(always)]
377	0	fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
378	0	from.simd_cast()
379	0	}
380
381		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
382		#[inline(always)]
383	0	unsafe fn handle_small(self, pto: mut f16, pfrom: const f32, len: usize) {
384	0	for i in 0..len {
385	0	pto.add(i) = diskann_wide::cast_f32_to_f16(pfrom.add(i))
386		}
387	0	}
388		}
389
390		//---------//
391		// Aarch64 //
392		//---------//
393
394		#[cfg(target_arch = "aarch64")]
395		impl SIMDConvert<f32, f16> for diskann_wide::arch::aarch64::Neon {
396		type Width = Const<4>;
397		type WideTo = <diskann_wide::arch::aarch64::Neon as Architecture>::f32x4;
398		type WideFrom = diskann_wide::arch::aarch64::f16x4;
399
400		#[inline(always)]
401		fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
402		from.into()
403		}
404
405		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
406		#[inline(always)]
407		unsafe fn handle_small(self, pto: mut f32, pfrom: const f16, len: usize) {
408		for i in 0..len {
409		pto.add(i) = diskann_wide::cast_f16_to_f32(pfrom.add(i))
410		}
411		}
412		}
413
414		#[cfg(target_arch = "aarch64")]
415		impl SIMDConvert<f16, f32> for diskann_wide::arch::aarch64::Neon {
416		type Width = Const<4>;
417		type WideTo = diskann_wide::arch::aarch64::f16x4;
418		type WideFrom = <diskann_wide::arch::aarch64::Neon as Architecture>::f32x4;
419
420		#[inline(always)]
421		fn simd_convert(from: Self::WideFrom) -> Self::WideTo {
422		from.simd_cast()
423		}
424
425		// SAFETY: We only access data in the valid range for `pto` and `pfrom`.
426		#[inline(always)]
427		unsafe fn handle_small(self, pto: mut f16, pfrom: const f32, len: usize) {
428		for i in 0..len {
429		pto.add(i) = diskann_wide::cast_f32_to_f16(pfrom.add(i))
430		}
431		}
432		}
433
434		///////////
435		// Tests //
436		///////////
437
438		#[cfg(test)]
439		mod tests {
440		use rand::{
441		distr::{Distribution, StandardUniform},
442		rngs::StdRng,
443		SeedableRng,
444		};
445
446		use super::*;
447
448		////////////////
449		// Fuzz tests //
450		////////////////
451
452		trait ReferenceConvert<From> {
453		fn reference_convert(self, from: &[From]);
454		}
455
456		impl ReferenceConvert<f32> for &mut [f16] {
457		fn reference_convert(self, from: &[f32]) {
458		assert_eq!(self.len(), from.len());
459		std::iter::zip(self.iter_mut(), from.iter()).for_each(\|(d, s)\| d = f16::from_f32(s));
460		}
461		}
462
463		impl ReferenceConvert<f16> for &mut [f32] {
464		fn reference_convert(self, from: &[f16]) {
465		assert_eq!(self.len(), from.len());
466		std::iter::zip(self.iter_mut(), from.iter()).for_each(\|(d, s)\| d = (s).into());
467		}
468		}
469
470		fn test_cast_from_slice<To, From>(max_dim: usize, num_trials: usize, rng: &mut StdRng)
471		where
472		StandardUniform: Distribution<From>,
473		To: Default + PartialEq + std::fmt::Debug + Copy,
474		From: Default + Copy,
475		for<'a, 'b> &'a mut [To]: CastFromSlice<&'b [From]> + ReferenceConvert<From>,
476		{
477		let distribution = StandardUniform {};
478		for dim in 0..=max_dim {
479		let mut src = vec![From::default(); dim];
480		let mut dst = vec![To::default(); dim];
481		let mut dst_reference = vec![To::default(); dim];
482
483		for _ in 0..num_trials {
484		src.iter_mut().for_each(\|s\| *s = distribution.sample(rng));
485		dst.cast_from_slice(src.as_slice());
486		dst_reference.reference_convert(&src);
487
488		assert_eq!(dst, dst_reference);
489		}
490		}
491		}
492
493		#[test]
494		fn test_f32_to_f16_fuzz() {
495		let mut rng = StdRng::seed_from_u64(0x0a3bfe052a8ebf98);
496		test_cast_from_slice::<f16, f32>(256, 10, &mut rng);
497		}
498
499		#[test]
500		fn test_f16_to_f32_fuzz() {
501		let mut rng = StdRng::seed_from_u64(0x83765b2816321eca);
502		test_cast_from_slice::<f32, f16>(256, 10, &mut rng);
503		}
504
505		////////////////
506		// Miri Tests //
507		////////////////
508
509		// With Miri, we need to be really careful that we do not hit methods that call
510		// `cvtph_ps`...
511		#[test]
512		fn miri_test_f32_to_f16() {
513		for dim in 0..256 {
514		println!("processing dim {}", dim);
515
516		let src = vec![f32::default(); dim];
517		let mut dst = vec![f16::default(); dim];
518
519		// Just check that all accesses are inbounds.
520		dst.cast_from_slice(&src);
521
522		// Scalar conversion.
523		SliceCast::<f16, f32>::new().run(
524		diskann_wide::arch::Scalar,
525		dst.as_mut_slice(),
526		src.as_slice(),
527		);
528
529		// SIMD conversion
530		#[cfg(target_arch = "x86_64")]
531		{
532		if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() {
533		SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
534		}
535
536		if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() {
537		SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
538		}
539		}
540
541		#[cfg(target_arch = "aarch64")]
542		if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() {
543		SliceCast::<f16, f32>::new().run(arch, dst.as_mut_slice(), src.as_slice())
544		}
545		}
546		}
547
548		#[test]
549		fn miri_test_f16_to_f32() {
550		for dim in 0..256 {
551		println!("processing dim {}", dim);
552
553		let src = vec![f16::default(); dim];
554		let mut dst = vec![f32::default(); dim];
555
556		// Just check that all accesses are inbounds.
557		dst.cast_from_slice(&src);
558
559		// Scalar conversion.
560		SliceCast::<f32, f16>::new().run(
561		diskann_wide::arch::Scalar,
562		dst.as_mut_slice(),
563		src.as_slice(),
564		);
565
566		// SIMD conversion
567		#[cfg(target_arch = "x86_64")]
568		{
569		if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() {
570		SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
571		}
572
573		if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() {
574		SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
575		}
576		}
577
578		#[cfg(target_arch = "aarch64")]
579		if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() {
580		SliceCast::<f32, f16>::new().run(arch, dst.as_mut_slice(), src.as_slice())
581		}
582		}
583		}
584		}