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

lattice_embed/simd/
mod.rs
