lattice_embed/simd/
mod.rs

1//! SIMD vector operations for embedding similarity and compressed-vector search.
2//!
3//! Dispatch uses the best supported target kernel with scalar fallbacks; WASM
4//! SIMD128 is selected at compile time. The public SIMD surface is unstable.
5//!
6//! See docs/simd.md for the kernel family, dispatch, and quantization design.
7
8mod binary;
9mod cosine;
10mod distance;
11mod dot_product;
12mod int4;
13mod manhattan;
14mod normalize;
15mod quantized;
16mod tier;
17
18#[cfg(test)]
19mod tests;
20
21pub use binary::BinaryVector;
22pub use cosine::{
23    batch_cosine_one_vs_many, batch_cosine_similarity, cosine_similarity, cosine_similarity_fused,
24    cosine_similarity_pre_normalized, query_norm,
25};
26pub use distance::{euclidean_distance, squared_euclidean_distance};
27pub use dot_product::{
28    DotBatch4Kernel, DotKernel, batch_dot_product, dot_product, dot_product_batch4,
29    resolved_dot_product_batch4_kernel, resolved_dot_product_kernel,
30};
31pub use int4::{Int4Params, Int4Vector};
32pub use manhattan::manhattan_distance;
33pub use normalize::normalize;
34pub use quantized::{
35    I8DotKernel, QuantizationParams, QuantizedVector, cosine_similarity_i8, dot_product_i8,
36    dot_product_i8_raw, resolved_i8_dot_kernel,
37};
38pub use tier::{
39    NormalizationHint, PreparedQuery, PreparedQueryWithMeta, QuantizationTier, QuantizedData,
40    approximate_cosine_distance, approximate_cosine_distance_prepared,
41    approximate_cosine_distance_prepared_with_meta, approximate_dot_product,
42    approximate_dot_product_prepared, approximate_int4_batch_prepared,
43    approximate_int4_batch_prepared_into, approximate_int8_batch_prepared,
44    approximate_int8_batch_prepared_into, batch_approximate_cosine_distance_prepared,
45    batch_approximate_cosine_distance_prepared_into, is_unit_norm, prepare_query,
46    prepare_query_with_norm, try_approximate_cosine_distance_prepared,
47    try_approximate_dot_product_prepared,
48};
49
50use std::sync::OnceLock;
51
52/// **Unstable**: SIMD dispatch internals; fields may be added as new ISAs are supported.
53///
54/// SIMD configuration with runtime feature detection.
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub struct SimdConfig {
57    /// **Unstable**: AVX-512F support available (x86_64).
58    pub avx512f_enabled: bool,
59    /// **Unstable**: AVX2 support available (x86_64).
60    pub avx2_enabled: bool,
61    /// **Unstable**: FMA (Fused Multiply-Add) support available (x86_64).
62    pub fma_enabled: bool,
63    /// **Unstable**: AVX-512F + AVX-512VNNI support available (x86_64).
64    pub avx512vnni_enabled: bool,
65    /// **Unstable**: NEON support available (aarch64/ARM64).
66    pub neon_enabled: bool,
67    /// **Unstable**: whether aarch64 FEAT_DotProd is available for SDOT/UDOT dispatch.
68    ///
69    /// Always false off aarch64. See [`docs/simd.md`](../../docs/simd.md#dispatch-model) for availability rules.
70    pub dotprod_enabled: bool,
71}
72
73impl Default for SimdConfig {
74    fn default() -> Self {
75        Self::detect()
76    }
77}
78
79impl SimdConfig {
80    /// **Unstable**: feature detection details may change as ISA support expands.
81    pub fn detect() -> Self {
82        #[cfg(target_arch = "x86_64")]
83        {
84            let avx512f_enabled = is_x86_feature_detected!("avx512f");
85
86            Self {
87                avx512f_enabled,
88                avx2_enabled: is_x86_feature_detected!("avx2"),
89                fma_enabled: is_x86_feature_detected!("fma"),
90                avx512vnni_enabled: avx512f_enabled
91                    && is_x86_feature_detected!("avx512bw")
92                    && is_x86_feature_detected!("avx512vnni"),
93                neon_enabled: false,
94                dotprod_enabled: false,
95            }
96        }
97        #[cfg(target_arch = "aarch64")]
98        {
99            // NEON is mandatory on aarch64, always available.
100            // FEAT_DotProd (dotprod) is optional: required on Armv8.4+,
101            // optional on Armv8.2/v8.3. Detect at runtime.
102            Self {
103                avx512f_enabled: false,
104                avx2_enabled: false,
105                fma_enabled: false,
106                avx512vnni_enabled: false,
107                neon_enabled: true,
108                dotprod_enabled: std::arch::is_aarch64_feature_detected!("dotprod"),
109            }
110        }
111        #[cfg(target_arch = "wasm32")]
112        {
113            // No runtime detection on wasm32: `simd128` is either compiled in
114            // for the whole module (via `-C target-feature=+simd128`) or not.
115            // `cfg!` reads the same compile-time flag the `#[cfg(...)]` gates
116            // on the SIMD kernel functions themselves key off, so this stays
117            // consistent with which kernels actually exist in the binary.
118            Self {
119                avx512f_enabled: false,
120                avx2_enabled: false,
121                fma_enabled: false,
122                avx512vnni_enabled: false,
123                neon_enabled: false,
124                dotprod_enabled: false,
125            }
126        }
127        #[cfg(not(any(
128            target_arch = "x86_64",
129            target_arch = "aarch64",
130            target_arch = "wasm32"
131        )))]
132        {
133            Self {
134                avx512f_enabled: false,
135                avx2_enabled: false,
136                fma_enabled: false,
137                avx512vnni_enabled: false,
138                neon_enabled: false,
139                dotprod_enabled: false,
140            }
141        }
142    }
143
144    /// **Unstable**: reports whether this wasm32 artifact was built with SIMD128.
145    ///
146    /// See [`docs/simd.md`](../../docs/simd.md#dispatch-model) for the compile-time dispatch model.
147    #[inline]
148    pub fn simd128_enabled(&self) -> bool {
149        cfg!(all(target_arch = "wasm32", target_feature = "simd128"))
150    }
151
152    /// **Unstable**: check if any SIMD is available; logic may expand with new ISAs.
153    #[inline]
154    pub fn simd_available(&self) -> bool {
155        self.avx512f_enabled
156            || self.avx512vnni_enabled
157            || self.avx2_enabled
158            || self.neon_enabled
159            || self.simd128_enabled()
160    }
161
162    /// Force scalar-only mode (useful for testing).
163    #[cfg(test)]
164    pub fn scalar_only() -> Self {
165        Self {
166            avx512f_enabled: false,
167            avx2_enabled: false,
168            fma_enabled: false,
169            avx512vnni_enabled: false,
170            neon_enabled: false,
171            dotprod_enabled: false,
172        }
173    }
174}
175
176// Process-wide SIMD configuration (detected once).
177static SIMD_CONFIG: OnceLock<SimdConfig> = OnceLock::new();
178
179/// **Unstable**: SIMD dispatch internal; shape may change as new backends are added.
180///
181/// The config is detected once per process and cached.
182#[inline]
183pub fn simd_config() -> SimdConfig {
184    *SIMD_CONFIG.get_or_init(SimdConfig::detect)
185}
lattice_embed/simd/mod.rs

lattice_embed/simd/
mod.rs
