Skip to main content

hermes_simd/
lib.rs

1//! High-performance, zero-overhead SIMD abstraction library.
2//!
3//! # Architecture
4//!
5//! `hermes-simd` is the public facade for the hermes-simd workspace:
6//!
7//! - [`hermes_simd_core`] — core abstractions, traits, views
8//! - [`hermes_simd_intrinsics`] — architecture-specific kernels
9//! - [`hermes_simd_macros`] — proc-macro code generation
10//!
11//! # Feature Flags
12//!
13//! | Feature | Description |
14//! |---------|-------------|
15//! | `std` (default) | Runtime CPU feature detection (`is_x86_feature_detected!`); without it dispatch uses compile-time `cfg!(target_feature)` only |
16//! | `mnemosyne-memory` (default) | Route `AlignedVec` allocation through the mnemosyne allocator |
17//! | `libnuma` | Linux NUMA affinity and residency probes via libnuma (links `-lnuma`); allocation still routes through Mnemosyne/the configured allocator |
18//!
19//! # Usage Examples
20//!
21//! **Dense sum (runtime dispatch):**
22//! ```rust
23//! use hermes_simd::sum;
24//! let data = vec![1.0f32; 1024];
25//! assert_eq!(sum(&data), 1024.0);
26//! ```
27//!
28//! **Masked dot product:**
29//! ```rust
30//! use hermes_simd::masked_dot;
31//! let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
32//! let b = vec![1.0f32; 5];
33//! let mask = vec![true, false, true, false, true];
34//! assert_eq!(masked_dot(&a, &b, &mask).unwrap(), 9.0); // 1+3+5
35//! ```
36
37#![cfg_attr(not(feature = "std"), no_std)]
38#![deny(missing_docs)]
39#![allow(
40    unused_unsafe,
41    clippy::too_many_arguments,
42    clippy::needless_range_loop,
43    clippy::assign_op_pattern,
44    clippy::manual_memcpy
45)]
46
47extern crate alloc;
48
49pub use hermes_simd_core::{
50    align::{Aligned, Alignment, Unaligned},
51    arch::SimdArch,
52    bitboard::{BitBoardKernel, BitBoardView},
53    // ComputeView extension
54    compute::{ComputeReduce, ComputeView},
55    cow::{ArchivedPacked4Cow, ArchivedSimdCow, Packed4CowResolver, SimdCow, SimdCowResolver},
56    current_numa_node,
57    execution::{ExecutionMode, Masked, Unmasked},
58    // Chunk iterators
59    iter::{SimdChunks, SimdChunksMut, ZipChunks, ZipChunksMut},
60    kernel::SimdKernel,
61    mask::BitMask,
62    refresh_numa_node,
63    scalar::{CastFrom, CastTo, FloatElement, Scalar as SimdScalar},
64    vec::AlignedVec,
65    verify_numa_locality,
66    view::{Mask, SimdError, SimdView, TileMatrixMultiply, TileView, Vector},
67    // Unary strategy ZSTs
68    Abs,
69    Add,
70    BitAnd,
71    BitOr,
72    BitXor,
73    Clamp,
74    Div,
75    Dot,
76    ElementOp,
77    Exclusive,
78    // Extended strategy ZSTs (v2)
79    FmaAdd,
80    Inclusive,
81    Mul,
82    Neg,
83    NumaBinding,
84    Popcount,
85    Product,
86    RecipSqrt,
87    // Operation strategy ZSTs and sealed traits — zero-cost, erased at monomorphization.
88    ReductionOp,
89    ScanAdd,
90    ScanMax,
91    ScanMin,
92    ScanMode,
93    ScanMul,
94    // Scan strategy ZSTs
95    ScanOp,
96    Sqrt,
97    Sub,
98    Sum,
99    UnaryOp,
100};
101
102// Re-export sparse types
103pub use hermes_simd_core::sparse::{
104    BlockedCoo,
105    BlockedCooData,
106    // Format-to-owned-storage mapping for Cow containers
107    CowFormat,
108    Csr,
109    CsrData,
110    DenseWithMask,
111    DenseWithMaskData,
112    OwnedBlockedCoo,
113    // Owned heap-backed sparse storage types
114    OwnedCsr,
115    OwnedDenseWithMask,
116    OwnedSellP,
117    SellP,
118    SellPData,
119    // Generic Clone-on-Write sparse container
120    SparseCow,
121    SparseFormat,
122    SparseOps,
123    SparseSpMv,
124    SparseView,
125    Validated,
126    ValidatedData,
127};
128
129// Re-export tiling
130pub use hermes_simd_core::tiling::{tiled_dot, tiled_gemv, TilingPolicy, TilingStrategy};
131
132// Re-export tensor views
133pub use hermes_simd_core::tensor::{ColMajor, RowMajor, TensorCow, TensorError, TensorView};
134
135// Re-export concrete ZST architecture markers
136pub use hermes_simd_intrinsics::{
137    Avx2, Avx512, AvxVnni, HybridSwarMagic, Hyperbola, KoggeStone, Magic, Neon, Scalar, SveArch,
138    Swar, SwarUtils,
139};
140
141#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
142pub use hermes_simd_intrinsics::{
143    AmxBatchSession, AmxBf16, AmxConfig, AmxInt8, AmxSession, AmxSessionError,
144};
145
146pub use eunomia::{
147    unpack_bf4_to_bf16, unpack_bf4_to_bf16_packed, unpack_bf8_to_bf16, unpack_f4_to_f32,
148    unpack_f4_to_f32_packed, unpack_f8_to_f32, Bf16, Bf4, Bf8, Packable4, Packed4Cow, Packed4Iter,
149    Packed4Slice, Packed4SliceMut, Packed4Vec, PackedBf4Cow, PackedBf4Slice, PackedBf4SliceMut,
150    PackedBf4Vec, PackedF4Cow, PackedF4Slice, PackedF4SliceMut, PackedF4Vec, F16, F32, F4, F64, F8,
151    I16, I32, I8,
152};
153
154// Re-export monomorphized vector register types and PreferredArch
155pub use hermes_simd_types::{
156    MaskBf16, MaskBf4, MaskBf8, MaskF16, MaskF32, MaskF4, MaskF64, MaskF8, MaskI16, MaskI32,
157    MaskI8, PreferredArch, ScalarBf16, ScalarBf4, ScalarBf8, ScalarF16, ScalarF32, ScalarF4,
158    ScalarF64, ScalarF8, ScalarI16, ScalarI32, ScalarI8, ScalarMaskF32, ScalarMaskF64, SimdBf16,
159    SimdBf4, SimdBf8, SimdF16, SimdF32, SimdF4, SimdF64, SimdF8, SimdI16, SimdI32, SimdI8,
160    SimdMaskBf16, SimdMaskBf4, SimdMaskBf8, SimdMaskF16, SimdMaskF32, SimdMaskF4, SimdMaskF64,
161    SimdMaskF8, SimdMaskI16, SimdMaskI32, SimdMaskI8, VectorBf16, VectorBf4, VectorBf8, VectorF16,
162    VectorF32, VectorF4, VectorF64, VectorF8, VectorI16, VectorI32, VectorI8,
163};
164
165#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
166pub use hermes_simd_types::{
167    Avx2Bf16, Avx2Bf4, Avx2Bf8, Avx2F16, Avx2F32, Avx2F4, Avx2F64, Avx2F8, Avx2I16, Avx2I32,
168    Avx2I8, Avx2MaskBf16, Avx2MaskF16, Avx2MaskF32, Avx2MaskF64, Avx512Bf16, Avx512Bf4, Avx512Bf8,
169    Avx512F16, Avx512F32, Avx512F4, Avx512F64, Avx512F8, Avx512I16, Avx512I32, Avx512I8,
170    Avx512MaskBf16, Avx512MaskF16, Avx512MaskF32, Avx512MaskF64,
171};
172
173#[cfg(target_arch = "aarch64")]
174pub use hermes_simd_types::{
175    NeonBf16, NeonBf4, NeonBf8, NeonF16, NeonF32, NeonF4, NeonF64, NeonF8, NeonI16, NeonI32,
176    NeonI8, NeonMaskBf16, NeonMaskF16, NeonMaskF32, NeonMaskF64,
177};
178
179/// Chess board attack generation kernels using bitboards and SWAR.
180pub mod attacks;
181/// Runtime CPU feature detection utilities.
182pub mod cpu;
183/// Dynamic dispatcher choosing optimal backends based on hardware/layout.
184pub mod dispatcher;
185
186/// Tiled matrix multiplication dispatch and kernel interfaces.
187pub mod tile_matmul;
188
189/// Runtime-dispatched SIMD abstractions and dynamic facade.
190pub mod dispatch;
191/// Explicit runtime target tokens and forced dispatch helpers.
192pub mod target;
193
194pub use attacks::{bishop_attacks, queen_attacks, rook_attacks};
195pub use cpu::{has_fma3, AmxSupport, Avx512Support, FmaSupport};
196pub use dispatcher::{AdaptiveDispatcher, DispatchDecision};
197pub use target::{dispatch_view_mut_to, dispatch_view_to, TargetId};
198pub use tile_matmul::{
199    dispatch_tile_matmul, gemm, unpack_int4, widen_I8_to_I16, widen_I8_to_I32, widen_i8_to_i16,
200    widen_i8_to_i32, TiledGemm,
201};
202
203// Re-export the generic dispatch operations. These monomorphize at call sites:
204// calling `sum::<f32>(data)` produces the f32 specialization.
205pub use dispatch::{
206    abs_max,
207    abs_sum,
208    argmax,
209    argmin,
210    axpy,
211    axpy_mul,
212    axpy_rows,
213    axpy_rows_batch,
214    dot,
215    elementwise_add,
216    elementwise_div,
217    elementwise_mul,
218    elementwise_sub,
219    gemv,
220    gemv_strided,
221    gemv_transpose,
222    gemv_transpose_strided,
223    interleaved_complex_dot,
224    interleaved_complex_dot_runtime,
225    interleaved_complex_mul_assign,
226    interleaved_complex_mul_assign_runtime,
227    masked_add,
228    masked_dot,
229    masked_sum,
230    max,
231    min,
232    ntt_butterfly_stage_u64,
233    // Generic free functions — the primary public API.
234    reduce_popcount,
235    reduce_popcount_and,
236    reduce_popcount_or,
237    reduce_popcount_xor,
238    scale,
239    spmv_bcoo,
240    // Sparse operations — generic entry points.
241    spmv_csr,
242    spmv_dense_masked,
243    spmv_sellp,
244    sum,
245    tiled_gemm,
246    // Core trait — sealed; implemented for f32 and f64.
247    SimdOps,
248};
249
250/// Target-specific, runtime-dispatched SIMD view wrapper.
251pub enum DispatchedView<'a, T, Align = Unaligned, Mode = Unmasked, Ref = &'a [T]>
252where
253    Align: hermes_simd_core::align::Alignment,
254    Mode: hermes_simd_core::execution::ExecutionMode,
255    Ref: core::ops::Deref<Target = [T]>,
256{
257    /// AVX-512 architecture target.
258    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
259    Avx512(SimdView<'a, T, Avx512, Align, Mode, Ref>),
260    /// AVX2 architecture target.
261    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
262    Avx2(SimdView<'a, T, Avx2, Align, Mode, Ref>),
263    /// NEON architecture target.
264    #[cfg(target_arch = "aarch64")]
265    Neon(SimdView<'a, T, Neon, Align, Mode, Ref>),
266    /// Fallback scalar target.
267    Scalar(SimdView<'a, T, Scalar, Align, Mode, Ref>),
268}
269
270/// Dispatches a shared slice into the best matching `DispatchedView` based on runtime CPU feature detection.
271#[inline]
272#[allow(unreachable_code)]
273pub fn dispatch_view<'a, T, Align>(
274    data: &'a [T],
275) -> Option<DispatchedView<'a, T, Align, Unmasked, &'a [T]>>
276where
277    T: FloatElement,
278    Align: Alignment,
279{
280    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
281    {
282        #[cfg(feature = "std")]
283        {
284            if std::is_x86_feature_detected!("avx512f") {
285                return SimdView::<T, Avx512, Align, Unmasked, &'a [T]>::new(data)
286                    .map(DispatchedView::Avx512);
287            }
288            if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
289                return SimdView::<T, Avx2, Align, Unmasked, &'a [T]>::new(data)
290                    .map(DispatchedView::Avx2);
291            }
292        }
293        #[cfg(not(feature = "std"))]
294        {
295            if cfg!(target_feature = "avx512f") {
296                return SimdView::<T, Avx512, Align, Unmasked, &'a [T]>::new(data)
297                    .map(DispatchedView::Avx512);
298            }
299            if cfg!(target_feature = "avx2") && cfg!(target_feature = "fma") {
300                return SimdView::<T, Avx2, Align, Unmasked, &'a [T]>::new(data)
301                    .map(DispatchedView::Avx2);
302            }
303        }
304    }
305    #[cfg(target_arch = "aarch64")]
306    {
307        return SimdView::<T, Neon, Align, Unmasked, &'a [T]>::new(data).map(DispatchedView::Neon);
308    }
309    SimdView::<T, Scalar, Align, Unmasked, &'a [T]>::new(data).map(DispatchedView::Scalar)
310}
311
312/// Dispatches a mutable slice into the best matching `DispatchedView` based on runtime CPU feature detection.
313#[inline]
314#[allow(unreachable_code)]
315pub fn dispatch_view_mut<'a, T, Align>(
316    data: &'a mut [T],
317) -> Option<DispatchedView<'a, T, Align, Unmasked, &'a mut [T]>>
318where
319    T: FloatElement,
320    Align: Alignment,
321{
322    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
323    {
324        #[cfg(feature = "std")]
325        {
326            if std::is_x86_feature_detected!("avx512f") {
327                return SimdView::<T, Avx512, Align, Unmasked, &'a mut [T]>::new_mut(data)
328                    .map(DispatchedView::Avx512);
329            }
330            if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
331                return SimdView::<T, Avx2, Align, Unmasked, &'a mut [T]>::new_mut(data)
332                    .map(DispatchedView::Avx2);
333            }
334        }
335        #[cfg(not(feature = "std"))]
336        {
337            if cfg!(target_feature = "avx512f") {
338                return SimdView::<T, Avx512, Align, Unmasked, &'a mut [T]>::new_mut(data)
339                    .map(DispatchedView::Avx512);
340            }
341            if cfg!(target_feature = "avx2") && cfg!(target_feature = "fma") {
342                return SimdView::<T, Avx2, Align, Unmasked, &'a mut [T]>::new_mut(data)
343                    .map(DispatchedView::Avx2);
344            }
345        }
346    }
347    #[cfg(target_arch = "aarch64")]
348    {
349        return SimdView::<T, Neon, Align, Unmasked, &'a mut [T]>::new_mut(data)
350            .map(DispatchedView::Neon);
351    }
352    SimdView::<T, Scalar, Align, Unmasked, &'a mut [T]>::new_mut(data).map(DispatchedView::Scalar)
353}
354
355/// Extension trait for `SimdCow` providing vector-register level operations.
356pub trait SimdCowExt<T: SimdScalar, Arch: SimdArch + SimdKernel<T>, Align: Alignment> {
357    /// In-place vector-level transformation.
358    ///
359    /// Promotes `self` to owned if borrowed (zero-copy upgrade), then applies
360    /// the function `f` elementwise to each SIMD vector chunk.
361    fn transform_vectors<F>(&mut self, f: F)
362    where
363        F: FnMut(Vector<T, Arch>) -> Vector<T, Arch>;
364}
365
366impl<'a, T, Arch, Align> SimdCowExt<T, Arch, Align> for SimdCow<'a, T, Arch, Align>
367where
368    T: SimdScalar,
369    Arch: SimdArch + SimdKernel<T>,
370    Align: Alignment,
371{
372    #[inline]
373    fn transform_vectors<F>(&mut self, mut f: F)
374    where
375        F: FnMut(Vector<T, Arch>) -> Vector<T, Arch>,
376    {
377        let owned_vec = self.to_mut();
378        let len = owned_vec.len();
379        let lane_count = Arch::LANE_COUNT;
380        let simd_len = (len / lane_count) * lane_count;
381        let slice = owned_vec.as_mut_slice();
382
383        let mut i = 0;
384        while i < simd_len {
385            unsafe {
386                let ptr = slice.as_mut_ptr().add(i);
387                let vec = if Align::IS_ALIGNED {
388                    Vector::load_aligned(ptr)
389                } else {
390                    Vector::load_unaligned(ptr)
391                };
392                let res = f(vec);
393                if Align::IS_ALIGNED {
394                    res.store_aligned(ptr);
395                } else {
396                    res.store_unaligned(ptr);
397                }
398            }
399            i += lane_count;
400        }
401
402        if simd_len < len {
403            // Tail buffer: sized conservatively at 64 elements (> any current
404            // SIMD lane width: SSE=4, AVX=8, AVX-512=16, SVE≤64).
405            // A debug assertion guards against future wider architectures.
406            debug_assert!(
407                lane_count <= 64,
408                "tail_buf[64] too small: lane_count={lane_count}"
409            );
410            let mut tail_buf = [core::mem::MaybeUninit::<T>::uninit(); 64];
411            let tail_len = len - simd_len;
412            unsafe {
413                for idx in 0..tail_len {
414                    tail_buf[idx].write(slice[simd_len + idx]);
415                }
416                for idx in tail_len..lane_count {
417                    tail_buf[idx].write(T::ZERO);
418                }
419                let vec = Vector::load_unaligned(tail_buf.as_ptr() as *const T);
420                let res = f(vec);
421                res.store_unaligned(tail_buf.as_mut_ptr() as *mut T);
422                for idx in 0..tail_len {
423                    slice[simd_len + idx] = tail_buf[idx].assume_init();
424                }
425            }
426        }
427    }
428}
429
430/// Extension trait for packed Clone-on-Write containers to support zero-copy unpacking directly into `SimdCow`.
431pub trait Packed4CowExt<'a, T: Packable4> {
432    /// Unpack packed elements directly to a `SimdCow` of wider precision with zero intermediate allocations.
433    fn unpack_to_cow<Arch, Align>(&self) -> SimdCow<'static, T::Unpacked, Arch, Align>
434    where
435        Arch: SimdArch,
436        Align: Alignment;
437}
438
439impl<'a, T: Packable4> Packed4CowExt<'a, T> for Packed4Cow<'a, T> {
440    #[inline]
441    fn unpack_to_cow<Arch, Align>(&self) -> SimdCow<'static, T::Unpacked, Arch, Align>
442    where
443        Arch: SimdArch,
444        Align: Alignment,
445    {
446        let len = self.len();
447        let mut dest = AlignedVec::with_capacity(len);
448        unsafe {
449            dest.set_len(len);
450        }
451        let view = self.as_view();
452        let n = view.len().min(dest.len());
453        let even_len = (n / 2) * 2;
454        T::unpack_slice_packed(
455            &view.as_packed_slice()[..even_len / 2],
456            &mut dest[..even_len],
457        );
458        if n % 2 != 0 {
459            if let Some(b) = view.get(n - 1) {
460                dest[n - 1] = T::unpack_single(b);
461            }
462        }
463        SimdCow::Owned(dest)
464    }
465}