nam_rs/math/common/avx512/
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#![allow(
4    // SAFETY: All unsafe blocks in this module are guarded by CPU feature detection
5    // (AVX-512F or AVX-512 VNNI+BF16) at dispatch time via SimdMathConfig::current().
6    // Each function's doc comment documents specific slice/pointer invariants.
7    unsafe_op_in_unsafe_fn,
8    clippy::missing_safety_doc,
9    clippy::too_many_arguments
10)]
11
12//! AVX-512 implementations of the `SimdMath` trait.
13//!
14//! Contains `Avx512Math` and `Avx512VnniBf16Math`.
15//! `Avx512VnniBf16Math` has real implementations (native BF16 dot product via `_mm512_dpbf16_ps`).
16//! Methods delegate to kernel functions in `math::gemm`, `math::wavenet`,
17//! `math::lstm`, `math::dsp`, `math::common::ops`, and `math::common::utility`.
18//!
19//! # Submodules
20//! - `gemv`: GEMV/GEMM kernels, dot products, and 4-gate LSTM.
21//! - `activations`: Activation functions (tanh/sigmoid), accumulation, and LSTM fusions.
22//! - `bf16`: FP32↔BF16 conversions and `store_bf16` wrappers.
23//! - `reduce`: Horizontal sum, energy, and max-diff.
24//! - `dsp`: Convolution, gain, ramp, and WaveNet head-sum.
25
26mod activations;
27mod bf16;
28mod dsp;
29mod gemv;
30mod reduce;
31
32use crate::math::common::InstructionSet;
33use crate::math::common::scalar_ref::*;
34use crate::math::common::traits::SimdMath;
35use core::arch::x86_64::*;
36
37/// SIMD implementation via AVX-512.
38/// This struct groups all mathematical functions optimized for processors that support AVX-512.
39pub struct Avx512Math;
40
41/// Static implementation for AVX-512 with VNNI and BF16 (Brain Float 16) support.
42/// This is the "Ferrari" of audio processing, available on very recent Intel CPUs (e.g.: Sapphire Rapids).
43/// The BF16 format allows the chip to process twice the numbers with almost the same precision as the original f32.
44pub struct Avx512VnniBf16Math;
45
46// ── Avx512Math ──
47
48impl SimdMath for Avx512Math {
49    type V = __m512;
50
51    const ISA: InstructionSet = InstructionSet::Avx512;
52
53    gemv::impl_avx512_gemv!();
54    activations::impl_avx512_activations!();
55    bf16::impl_avx512_bf16!();
56    reduce::impl_avx512_reduce!();
57    dsp::impl_avx512_dsp!();
58}
59
60// ── Avx512VnniBf16Math ──
61
62impl SimdMath for Avx512VnniBf16Math {
63    type V = __m512;
64
65    const ISA: InstructionSet = InstructionSet::Avx512VnniBf16;
66
67    gemv::impl_avx512vnni_bf16_gemv!();
68    activations::impl_avx512vnni_bf16_activations!();
69    bf16::impl_avx512vnni_bf16_bf16!();
70    reduce::impl_avx512vnni_bf16_reduce!();
71    dsp::impl_avx512vnni_bf16_dsp!();
72}
nam_rs/math/common/avx512/mod.rs

nam_rs/math/common/avx512/
mod.rs
