usearch 2.21.3 → 2.25.0

package/numkong/rust/capabilities.rs

@@ -0,0 +1,96 @@
+//! Runtime CPU capability detection.
+//!
+//! Capability bits (`cap::*`) are probed at library load time by examining
+//! CPUID / `getauxval` / HWCAP records on the host CPU; [`available`] returns the
+//! resulting bitmask so downstream code can select the best kernel without
+//! recompiling.
+//!
+//! This module provides:
+//!
+//! - [`available`]: Query the bitmask of supported SIMD instruction sets
+//! - [`configure_thread`]: Enable optimal SIMD settings for the current thread
+//! - [`uses_dynamic_dispatch`]: Check if the library selects kernels at runtime
+//! - [`cap`]: Constants for individual capability bits (NEON, SKYLAKE, etc.)
+
+#[link(name = "numkong")]
+extern "C" {
+    fn nk_configure_thread(capabilities: u64) -> i32;
+    fn nk_uses_dynamic_dispatch() -> i32;
+    fn nk_capabilities() -> u64;
+}
+
+/// Returns the bitmask of available CPU capabilities.
+/// Use with `cap::*` constants to check for specific features.
+///
+/// # Example
+/// ```
+/// use numkong::{capabilities, cap};
+///
+/// let caps = capabilities::available();
+/// if caps & cap::NEON != 0 {
+///     println!("NEON is available");
+/// }
+/// if caps & cap::SKYLAKE != 0 {
+///     println!("AVX-512 (Skylake) is available");
+/// }
+/// ```
+pub fn available() -> u64 {
+    unsafe { nk_capabilities() }
+}
+
+/// Configures the current thread for optimal SIMD performance.
+/// On x86, this enables AMX tile state via `arch_prctl`. On other platforms this is a no-op.
+/// Must be called once per thread before using AMX (Advanced Matrix Extensions) operations.
+pub fn configure_thread() -> bool {
+    // Pass !0 to enable all capabilities including AMX
+    unsafe { nk_configure_thread(!0) != 0 }
+}
+
+/// Returns `true` if the library uses dynamic dispatch for function selection.
+pub fn uses_dynamic_dispatch() -> bool {
+    unsafe { nk_uses_dynamic_dispatch() != 0 }
+}
+
+/// Capability bit masks in chronological order (by first commercial silicon).
+pub mod cap {
+    pub const SERIAL: u64 = 1 << 0; // Always: Fallback
+    pub const NEON: u64 = 1 << 1; // ARM NEON
+    pub const HASWELL: u64 = 1 << 2; // Intel AVX2
+    pub const SKYLAKE: u64 = 1 << 3; // Intel AVX-512
+    pub const NEONHALF: u64 = 1 << 4; // ARM NEON FP16
+    pub const NEONSDOT: u64 = 1 << 5; // ARM NEON i8 dot
+    pub const NEONFHM: u64 = 1 << 6; // ARM NEON FP16 FML
+    pub const ICELAKE: u64 = 1 << 7; // Intel AVX-512 VNNI
+    pub const GENOA: u64 = 1 << 8; // AMD AVX-512 BF16
+    pub const NEONBFDOT: u64 = 1 << 9; // ARM NEON BF16
+    pub const SVE: u64 = 1 << 10; // ARM SVE
+    pub const SVEHALF: u64 = 1 << 11; // ARM SVE FP16
+    pub const SVESDOT: u64 = 1 << 12; // ARM SVE i8 dot
+    pub const ALDER: u64 = 1 << 13; // Intel AVX2+VNNI
+    pub const SVEBFDOT: u64 = 1 << 14; // ARM SVE BF16
+    pub const SVE2: u64 = 1 << 15; // ARM SVE2
+    pub const V128RELAXED: u64 = 1 << 16; // WASM Relaxed SIMD
+    pub const SAPPHIRE: u64 = 1 << 17; // Intel AVX-512 FP16
+    pub const SAPPHIREAMX: u64 = 1 << 18; // Intel Sapphire AMX
+    pub const RVV: u64 = 1 << 19; // RISC-V Vector
+    pub const RVVHALF: u64 = 1 << 20; // RISC-V Zvfh
+    pub const RVVBF16: u64 = 1 << 21; // RISC-V Zvfbfwma
+    pub const GRANITEAMX: u64 = 1 << 22; // Intel Granite AMX FP16
+    pub const TURIN: u64 = 1 << 23; // AMD Turin AVX-512 CD
+    pub const SME: u64 = 1 << 24; // ARM SME
+    pub const SME2: u64 = 1 << 25; // ARM SME2
+    pub const SMEF64: u64 = 1 << 26; // ARM SME F64
+    pub const SMEFA64: u64 = 1 << 27; // ARM SME FA64
+    pub const SVE2P1: u64 = 1 << 28; // ARM SVE2.1
+    pub const SME2P1: u64 = 1 << 29; // ARM SME2.1
+    pub const SMEHALF: u64 = 1 << 30; // ARM SME F16F16
+    pub const SMEBF16: u64 = 1 << 31; // ARM SME B16B16
+    pub const SMELUT2: u64 = 1 << 32; // ARM SME LUTv2
+    pub const RVVBB: u64 = 1 << 33; // RISC-V Zvbb
+    pub const SIERRA: u64 = 1 << 34; // Intel AVXVNNIINT8
+    pub const SMEBI32: u64 = 1 << 35; // ARM SME BI32I32
+    pub const LOONGSONASX: u64 = 1 << 36; // LoongArch LASX 256-bit SIMD
+    pub const POWERVSX: u64 = 1 << 37; // Power VSX 128-bit SIMD
+    pub const DIAMOND: u64 = 1 << 38; // Intel AVX10.2
+    pub const NEONFP8: u64 = 1 << 39; // ARM NEON FP8
+}

package/numkong/rust/cast.rs

@@ -0,0 +1,314 @@
+//! Type casting between scalar formats.
+//!
+//! This module provides:
+//!
+//! - [`CastDtype`]: Trait marking types eligible for bulk casting
+//! - [`cast`]: Bulk-converts a slice from one scalar format to another
+
+use crate::types::{bf16, bf16c, e2m3, e3m2, e4m3, e5m2, f16, f16c, f32c, f64c, StorageElement};
+
+#[link(name = "numkong")]
+extern "C" {
+    fn nk_cast(
+        from: *const core::ffi::c_void,
+        from_type: u32,
+        n: usize,
+        to: *mut core::ffi::c_void,
+        to_type: u32,
+    );
+}
+
+/// Internal dtype codes matching `nk_dtype_t` from C.
+/// Not exposed to users.
+mod dtype {
+    pub(crate) const F64: u32 = 1 << 10;
+    pub(crate) const F32: u32 = 1 << 11;
+    pub(crate) const F16: u32 = 1 << 12;
+    pub(crate) const BF16: u32 = 1 << 13;
+    pub(crate) const E4M3: u32 = 1 << 14;
+    pub(crate) const E5M2: u32 = 1 << 15;
+    pub(crate) const E2M3: u32 = 1 << 18;
+    pub(crate) const E3M2: u32 = 1 << 19;
+    pub(crate) const F64C: u32 = 1 << 20;
+    pub(crate) const F32C: u32 = 1 << 21;
+    pub(crate) const F16C: u32 = 1 << 22;
+    pub(crate) const BF16C: u32 = 1 << 23;
+    pub(crate) const I8: u32 = 1 << 2;
+    pub(crate) const I16: u32 = 1 << 3;
+    pub(crate) const I32: u32 = 1 << 4;
+    pub(crate) const I64: u32 = 1 << 5;
+    pub(crate) const U8: u32 = 1 << 6;
+    pub(crate) const U16: u32 = 1 << 7;
+    pub(crate) const U32: u32 = 1 << 8;
+    pub(crate) const U64: u32 = 1 << 9;
+}
+
+// Sealed trait pattern to prevent external implementations
+mod private {
+    pub trait Sealed {}
+    impl Sealed for f64 {}
+    impl Sealed for f32 {}
+    impl Sealed for super::f16 {}
+    impl Sealed for super::bf16 {}
+    impl Sealed for super::e4m3 {}
+    impl Sealed for super::e5m2 {}
+    impl Sealed for super::e2m3 {}
+    impl Sealed for super::e3m2 {}
+    impl Sealed for super::f64c {}
+    impl Sealed for super::f32c {}
+    impl Sealed for super::f16c {}
+    impl Sealed for super::bf16c {}
+    impl Sealed for i8 {}
+    impl Sealed for i16 {}
+    impl Sealed for i32 {}
+    impl Sealed for i64 {}
+    impl Sealed for u8 {}
+    impl Sealed for u16 {}
+    impl Sealed for u32 {}
+    impl Sealed for u64 {}
+}
+
+/// Trait for types that can participate in cast operations.
+///
+/// This trait is sealed - users cannot implement it for their own types.
+pub trait CastDtype: private::Sealed + StorageElement {
+    #[doc(hidden)]
+    fn dtype_code() -> u32;
+}
+
+impl CastDtype for f64 {
+    fn dtype_code() -> u32 {
+        dtype::F64
+    }
+}
+impl CastDtype for f32 {
+    fn dtype_code() -> u32 {
+        dtype::F32
+    }
+}
+impl CastDtype for f16 {
+    fn dtype_code() -> u32 {
+        dtype::F16
+    }
+}
+impl CastDtype for bf16 {
+    fn dtype_code() -> u32 {
+        dtype::BF16
+    }
+}
+impl CastDtype for e4m3 {
+    fn dtype_code() -> u32 {
+        dtype::E4M3
+    }
+}
+impl CastDtype for e5m2 {
+    fn dtype_code() -> u32 {
+        dtype::E5M2
+    }
+}
+impl CastDtype for e2m3 {
+    fn dtype_code() -> u32 {
+        dtype::E2M3
+    }
+}
+impl CastDtype for e3m2 {
+    fn dtype_code() -> u32 {
+        dtype::E3M2
+    }
+}
+impl CastDtype for f64c {
+    fn dtype_code() -> u32 {
+        dtype::F64C
+    }
+}
+impl CastDtype for f32c {
+    fn dtype_code() -> u32 {
+        dtype::F32C
+    }
+}
+impl CastDtype for f16c {
+    fn dtype_code() -> u32 {
+        dtype::F16C
+    }
+}
+impl CastDtype for bf16c {
+    fn dtype_code() -> u32 {
+        dtype::BF16C
+    }
+}
+impl CastDtype for i8 {
+    fn dtype_code() -> u32 {
+        dtype::I8
+    }
+}
+impl CastDtype for i16 {
+    fn dtype_code() -> u32 {
+        dtype::I16
+    }
+}
+impl CastDtype for i32 {
+    fn dtype_code() -> u32 {
+        dtype::I32
+    }
+}
+impl CastDtype for i64 {
+    fn dtype_code() -> u32 {
+        dtype::I64
+    }
+}
+impl CastDtype for u8 {
+    fn dtype_code() -> u32 {
+        dtype::U8
+    }
+}
+impl CastDtype for u16 {
+    fn dtype_code() -> u32 {
+        dtype::U16
+    }
+}
+impl CastDtype for u32 {
+    fn dtype_code() -> u32 {
+        dtype::U32
+    }
+}
+impl CastDtype for u64 {
+    fn dtype_code() -> u32 {
+        dtype::U64
+    }
+}
+
+/// Cast source slice elements to destination slice.
+///
+/// Converts elements from source type `S` to destination type `D` using
+/// hardware-accelerated SIMD operations when available.
+///
+/// # Arguments
+/// * `source` - Source slice of elements to cast
+/// * `dest` - Destination slice to receive cast elements (must be same length as source)
+///
+/// # Returns
+/// * `Some(())` if successful
+/// * `None` if slices have different lengths
+///
+/// # Example
+/// ```ignore
+/// use numkong::{f16, cast};
+///
+/// let f16_data: Vec<f16> = vec![f16::from_f32(1.0), f16::from_f32(2.0)];
+/// let mut f32_data: Vec<f32> = vec![0.0; f16_data.len()];
+/// cast(&f16_data, &mut f32_data);
+/// ```
+pub fn cast<S: CastDtype, D: CastDtype>(source: &[S], dest: &mut [D]) -> Option<()> {
+    if source.len() != dest.len() {
+        return None;
+    }
+    unsafe {
+        nk_cast(
+            source.as_ptr() as *const core::ffi::c_void,
+            S::dtype_code(),
+            source.len(),
+            dest.as_mut_ptr() as *mut core::ffi::c_void,
+            D::dtype_code(),
+        );
+    }
+    Some(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::{
+        assert_close, bf16, bf16c, e2m3, e3m2, e4m3, e5m2, f16, f16c, f32c, f64c, FloatLike,
+        StorageElement, TestableType,
+    };
+
+    fn check_cast_roundtrip<T: FloatLike + TestableType + CastDtype>(values: &[f32]) {
+        let src: Vec<T> = values.iter().map(|&v| T::from_f32(v)).collect();
+        let mut dst = vec![0.0f32; src.len()];
+        cast(&src, &mut dst).unwrap();
+        for (i, (&expected, &actual)) in values.iter().zip(dst.iter()).enumerate() {
+            assert_close(
+                actual as f64,
+                expected as f64,
+                T::atol(),
+                T::rtol(),
+                &format!("cast_roundtrip<{}>[{i}]", core::any::type_name::<T>()),
+            );
+        }
+    }
+
+    #[test]
+    fn cast_roundtrip() {
+        check_cast_roundtrip::<f16>(&[1.0, 0.5, -1.0]);
+        check_cast_roundtrip::<bf16>(&[1.0, 0.5, -1.0]);
+        check_cast_roundtrip::<e4m3>(&[1.0, 0.5, -1.0]);
+        check_cast_roundtrip::<e5m2>(&[1.0, 0.5, -1.0]);
+        check_cast_roundtrip::<e2m3>(&[1.0, 0.5, -1.0]);
+        check_cast_roundtrip::<e3m2>(&[1.0, 0.5, -1.0]);
+    }
+
+    #[test]
+    fn cast_f32_to_f16() {
+        let src = [1.0f32, -1.0];
+        let mut dst = [f16(0); 2];
+        cast(&src, &mut dst).unwrap();
+        assert_eq!([dst[0].0, dst[1].0], [0x3C00, 0xBC00]);
+    }
+
+    #[test]
+    fn cast_length_mismatch() {
+        let src = [f16(0x3C00)];
+        let mut dst = [0.0f32; 2];
+        assert!(cast(&src, &mut dst).is_none());
+    }
+
+    #[test]
+    fn cast_real_to_complex() {
+        let src = [1.25f32, -2.5];
+        let mut dst = [f32c::zero(); 2];
+        cast(&src, &mut dst).unwrap();
+        assert_eq!(dst[0], f32c::from_real_imag(1.25, 0.0));
+        assert_eq!(dst[1], f32c::from_real_imag(-2.5, 0.0));
+
+        let src = [f16::from_f32(3.0), f16::from_f32(-4.0)];
+        let mut dst = [f16c::zero(); 2];
+        cast(&src, &mut dst).unwrap();
+        assert_eq!(
+            dst,
+            [
+                f16c::from_real_imag(f16::from_f32(3.0), f16::ZERO),
+                f16c::from_real_imag(f16::from_f32(-4.0), f16::ZERO),
+            ]
+        );
+    }
+
+    #[test]
+    fn cast_complex_to_real() {
+        let src = [
+            f64c::from_real_imag(1.25, 9.0),
+            f64c::from_real_imag(-2.5, -7.0),
+        ];
+        let mut dst = [0.0f64; 2];
+        cast(&src, &mut dst).unwrap();
+        assert_eq!(dst, [1.25, -2.5]);
+    }
+
+    #[test]
+    fn cast_complex_to_complex() {
+        let src = [
+            f32c::from_real_imag(1.25, -2.5),
+            f32c::from_real_imag(-3.5, 4.25),
+        ];
+        let mut widened = [f64c::zero(); 2];
+        cast(&src, &mut widened).unwrap();
+        assert_eq!(widened[0], f64c::from_real_imag(1.25, -2.5));
+        assert_eq!(widened[1], f64c::from_real_imag(-3.5, 4.25));
+
+        let mut narrowed = [bf16c::zero(); 2];
+        cast(&widened, &mut narrowed).unwrap();
+        assert_eq!(narrowed[0].re.to_f32(), bf16::from_f32(1.25).to_f32());
+        assert_eq!(narrowed[0].im.to_f32(), bf16::from_f32(-2.5).to_f32());
+        assert_eq!(narrowed[1].re.to_f32(), bf16::from_f32(-3.5).to_f32());
+        assert_eq!(narrowed[1].im.to_f32(), bf16::from_f32(4.25).to_f32());
+    }
+}