nam_rs/math/common/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.
3// SAFETY: Preconditions (alignment, bounds, size) are guaranteed by caller of this SIMD/unsafe function.
4#![warn(clippy::undocumented_unsafe_blocks)]
5
6//! Common foundation for mathematical operations and SIMD.
7//!
8//! This module contains the structural definitions that allow NAM-rs to be
9//! hardware-agnostic while maintaining native performance.
10//!
11//! # Components
12//! - `traits`: The `SimdMath` trait defining the interface for all kernels.
13//! - `dispatch`: Dynamic architecture selection mechanism (AVX2, AVX-512, etc.).
14//! - `avx2_impl` / `avx512`: Concrete kernel implementations for x86-64.
15//! - `scalar_ref`: Fallback implementations for compatibility and testing.
16//! - `aligned`: Structures to guarantee memory alignment (RT-Safety).
17
18pub mod aligned;
19pub mod avx2_impl;
20pub mod half;
21pub mod huge_alloc;
22pub use huge_alloc::HugePageVec;
23pub mod avx512;
24/// Unit tests for the common math infrastructure.
25#[cfg(test)]
26pub mod common_test;
27pub mod dispatch;
28pub mod kahan;
29pub mod ops;
30pub mod scalar_ref;
31pub mod traits;
32pub mod utility;
33
34pub use aligned::Aligned64;
35pub use aligned::AlignedVec;
36pub use avx2_impl::Avx2Math;
37pub use avx512::{Avx512Math, Avx512VnniBf16Math};
38pub use dispatch::{InstructionSet, SIMD_MATH, SimdMathConfig, TEST_ISA_OVERRIDE};
39pub use dispatch::{decode_isa_override, effective_instruction_set, encode_isa_override};
40pub use kahan::{Kahan4F32, KahanF32, kahan_add};
41pub use ops::*;
42pub use scalar_ref::*;
43pub use traits::SimdMath;
44pub use utility::*;
45
46/// Macro for dynamic SIMD dispatch based on global configuration.
47///
48/// Checks `TEST_ISA_OVERRIDE` (u8::MAX = disabled) before consulting
49/// `SIMD_MATH.instruction_set`. This allows integration tests to force a
50/// specific ISA path (AVX2, AVX-512, VNNI+BF16) and measure cross-ISA
51/// determinism.
52#[macro_export]
53macro_rules! dispatch_simd {
54    // Mode 2: Dispatch to specific methods of an object (e.g.: lstm.rs)
55    (@ $target:expr, $m512bf16:ident, $m512:ident, $m256:ident $(, $arg:expr)*) => {
56        {
57            use $crate::math::common::InstructionSet;
58            let __isa = $crate::math::common::effective_instruction_set();
59            match __isa {
60                InstructionSet::Avx512VnniBf16 => $target.$m512bf16($($arg),*),
61                InstructionSet::Avx512 => $target.$m512($($arg),*),
62                InstructionSet::Avx2 => $target.$m256($($arg),*),
63            }
64        }
65    };
66
67    // Mode 1: Dispatch to generic method of an object (e.g.: wavenet.rs)
68    ($target:expr, $method:ident $(, $arg:expr)*) => {
69        {
70            use $crate::math::common::InstructionSet;
71            let __isa = $crate::math::common::effective_instruction_set();
72            match __isa {
73                InstructionSet::Avx512VnniBf16 => {
74                    $target.$method::<$crate::math::common::Avx512VnniBf16Math>($($arg),*)
75                }
76                InstructionSet::Avx512 => {
77                    $target.$method::<$crate::math::common::Avx512Math>($($arg),*)
78                }
79                InstructionSet::Avx2 => {
80                    $target.$method::<$crate::math::common::Avx2Math>($($arg),*)
81                }
82            }
83        }
84    };
85
86    // Mode 3: Static dispatch to SimdMath trait associated functions
87    ($method:ident ($($arg:expr),*)) => {
88        {
89            use $crate::math::common::InstructionSet;
90            use $crate::math::common::traits::SimdMath;
91            let __isa = $crate::math::common::effective_instruction_set();
92            // SAFETY: Inner safety guarantees are upheld by caller invariants or the execution environment.
93            #[allow(clippy::macro_metavars_in_unsafe, clippy::allow_attributes)]
94            unsafe {
95                match __isa {
96                    InstructionSet::Avx512VnniBf16 => {
97                        $crate::math::common::Avx512VnniBf16Math::$method($($arg),*)
98                    }
99                    InstructionSet::Avx512 => {
100                        $crate::math::common::Avx512Math::$method($($arg),*)
101                    }
102                    InstructionSet::Avx2 => {
103                        $crate::math::common::Avx2Math::$method($($arg),*)
104                    }
105                }
106            }
107        }
108    };
109}
110
111pub use dispatch_simd;
nam_rs/math/common/mod.rs

nam_rs/math/common/
mod.rs
