Skip to main content

nam_rs/math/common/
traits.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
4//! [Definition of the abstract interface for SIMD operations.
5
6use super::InstructionSet;
7
8/// Abstraction trait for static dispatch of SIMD mathematical operations.
9///
10/// # Safety
11///
12/// All implementations of this trait use x86-64 SIMD intrinsics that require
13/// specific CPU features (AVX2+FMA minimum, enforced project-wide via `x86-64-v3`
14/// in `.cargo/config.toml`). Dynamic dispatch only selects higher extensions
15/// (AVX-512, AVX-512VNNI+BF16) at runtime via `SimdMathConfig::current()`.
16///
17/// Each method documents its own preconditions below. In general:
18/// - **All** pointer/slice arguments must be valid (non-null, within allocation,
19///   properly aligned for the target ISA).
20/// - **Convolution** functions additionally require coefficient pointers aligned to
21///   32 bytes (AVX2) or 64 bytes (AVX-512); all other operations use unaligned
22///   loads and stores.
23/// - **Length** invariants (e.g., `weights.len() >= input_len * output_len`) are
24///   documented per method. Callers must uphold them.
25/// - **Shared mutability** must not occur: no slice aliasing between `&mut` and
26///   any other access during the call.
27/// - Input values must be finite IEEE 754 `f32`; behavior with NaN/Inf is
28///   implementation-defined.
29/// - BF16-typed slices (`&[u16]`) must contain valid BF16 bit patterns.
30///
31/// # Operation Groups
32///
33/// The operations of this trait are organized into the following groups:
34/// - **(A) Dot Products**: Scalar/4x/dual-frame dot products (e.g., `dot_product`).
35/// - **(B) GEMV/GEMM Fused**: Fused matrix-vector/matrix-matrix kernels (e.g., `fused_add_gemv`).
36/// - **(C) Activations**: Tanh/sigmoid activation functions and gated fusions (e.g., `tanh_slice`).
37/// - **(D) Conversions**: f32 ↔ bf16 conversion utilities (e.g., `f32_to_bf16`).
38/// - **(E) LSTM Gates**: Specific kernels for LSTM cell gates (e.g., `fused_lstm_gates_dyn`).
39/// - **(F) Complex MAC**: Complex multiply-accumulate spectral kernels.
40/// - **(G) FFT Butterfly**: SIMD Radix-2 DIT FFT butterfly stages.
41/// - **(H) BatchNorm**: Frame-major batch normalization affine transform.
42pub trait SimdMath {
43    /// SIMD register type used (e.g.: __m256 or __m512).
44    type V: Copy;
45
46    /// Indicates whether this implementation uses weights and signals in BF16 format.
47    const IS_BF16: bool = false;
48
49    /// ISA that this implementation targets (compile-time constant for monomorphization).
50    const ISA: InstructionSet;
51
52    // --- (A) Dot Products ---
53
54    /// Computes the dot product between two f32 vectors.
55    ///
56    /// Effective length is `min(a.len(), b.len())`. No alignment required.
57    ///
58    /// # Safety
59    /// `a` and `b` must be valid slices.
60    unsafe fn dot_product(a: &[f32], b: &[f32]) -> f32;
61
62    /// Computes the dot product between two BF16 vectors.
63    ///
64    /// Effective length is `min(a.len(), b.len())`. No alignment required.
65    ///
66    /// # Safety
67    /// `a` and `b` must be valid slices with valid BF16 bit patterns.
68    unsafe fn dot_product_bf16(a: &[u16], b: &[u16]) -> f32;
69
70    /// Computes 4 simultaneous BF16 dot products (interleaved) with f32 input.
71    ///
72    /// No alignment required. The output is `[a·w0, a·w1, a·w2, a·w3]`
73    /// where each `w_k` is the k-th column of the interleaved weight matrix.
74    ///
75    /// # Safety
76    /// `weights.len() >= state.len()`. Both slices must be valid
77    /// and accessible for reading.
78    unsafe fn dot_product_4x_interleaved(weights: &[[u16; 4]], state: &[f32]) -> [f32; 4];
79
80    /// Computes 4 simultaneous BF16 dot products (interleaved) for 2 parallel frames.
81    /// Returns `(results_f0, results_f1)`.
82    ///
83    /// No alignment required.
84    ///
85    /// # Safety
86    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
87    /// All slices must be valid and accessible for reading.
88    unsafe fn dot_product_4x_interleaved_dual_frame(
89        weights: &[[u16; 4]],
90        state_f0: &[f32],
91        state_f1: &[f32],
92    ) -> ([f32; 4], [f32; 4]);
93
94    /// Computes 4 simultaneous dot products with native f32 weights.
95    ///
96    /// No alignment required.
97    ///
98    /// # Safety
99    /// `weights.len() >= state.len()`. Each element of `weights` must be
100    /// a valid `[f32; 4]` row. Both slices must be accessible for reading.
101    unsafe fn dot_product_4x_f32(weights: &[[f32; 4]], state: &[f32]) -> [f32; 4];
102
103    /// Computes 4 simultaneous dot products with native f32 weights
104    /// for 2 parallel frames.
105    /// Returns `(results_f0, results_f1)`.
106    ///
107    /// No alignment required.
108    ///
109    /// # Safety
110    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
111    /// All slices must be valid and accessible for reading.
112    unsafe fn dot_product_4x_f32_dual(
113        weights: &[[f32; 4]],
114        state_f0: &[f32],
115        state_f1: &[f32],
116    ) -> ([f32; 4], [f32; 4]);
117
118    /// Computes 8 simultaneous dot products with native f32 weights.
119    ///
120    /// No alignment required.
121    ///
122    /// # Safety
123    /// `weights.len() >= state.len()`. Each element of `weights` must be
124    /// a valid `[f32; 8]` row. Both slices must be accessible for reading.
125    unsafe fn dot_product_8x_f32(weights: &[[f32; 8]], state: &[f32]) -> [f32; 8];
126
127    /// Computes 8 simultaneous dot products with native f32 weights
128    /// for 2 parallel frames.
129    /// Returns `(results_f0, results_f1)`.
130    ///
131    /// No alignment required.
132    ///
133    /// # Safety
134    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
135    /// All slices must be valid and accessible for reading.
136    unsafe fn dot_product_8x_f32_dual(
137        weights: &[[f32; 8]],
138        state_f0: &[f32],
139        state_f1: &[f32],
140    ) -> ([f32; 8], [f32; 8]);
141
142    /// Computes 16 simultaneous dot products with native f32 weights.
143    ///
144    /// No alignment required.
145    ///
146    /// # Safety
147    /// `weights.len() >= state.len()`. Each element of `weights` must be
148    /// a valid `[f32; 16]` row. Both slices must be accessible for reading.
149    unsafe fn dot_product_16x_f32(weights: &[[f32; 16]], state: &[f32]) -> [f32; 16];
150
151    /// Computes 16 simultaneous dot products with native f32 weights
152    /// for 2 parallel frames.
153    /// Returns `(results_f0, results_f1)`.
154    ///
155    /// No alignment required.
156    ///
157    /// # Safety
158    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
159    /// All slices must be valid and accessible for reading.
160    unsafe fn dot_product_16x_f32_dual(
161        weights: &[[f32; 16]],
162        state_f0: &[f32],
163        state_f1: &[f32],
164    ) -> ([f32; 16], [f32; 16]);
165
166    /// Fused accumulate variant: computes 4 simultaneous dot products with native f32
167    /// weights and adds the `init` accumulator (bias + mixin).
168    ///
169    /// Equivalent to `dot_product_4x_f32(weights, state)` followed by element-wise
170    /// addition of `init`. Fusing both avoids an extra pass over the output array.
171    /// No alignment required.
172    ///
173    /// # Safety
174    /// `weights.len() >= state.len()`.
175    unsafe fn dot_product_4x_f32_accumulate(
176        weights: &[[f32; 4]],
177        state: &[f32],
178        init: &[f32; 4],
179    ) -> [f32; 4];
180
181    /// Fused accumulate dual-frame variant: computes 4 simultaneous dot products for
182    /// 2 parallel frames and adds the respective `init_f0` / `init_f1` accumulators.
183    /// Returns `(results_f0, results_f1)`.
184    ///
185    /// No alignment required.
186    ///
187    /// # Safety
188    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
189    unsafe fn dot_product_4x_f32_dual_accumulate(
190        weights: &[[f32; 4]],
191        state_f0: &[f32],
192        state_f1: &[f32],
193        init_f0: &[f32; 4],
194        init_f1: &[f32; 4],
195    ) -> ([f32; 4], [f32; 4]);
196
197    /// Fused accumulate variant: computes 8 simultaneous dot products with native f32
198    /// weights and adds the `init` accumulator.
199    ///
200    /// No alignment required.
201    ///
202    /// # Safety
203    /// `weights.len() >= state.len()`.
204    unsafe fn dot_product_8x_f32_accumulate(
205        weights: &[[f32; 8]],
206        state: &[f32],
207        init: &[f32; 8],
208    ) -> [f32; 8];
209
210    /// Fused accumulate dual-frame variant: computes 8 simultaneous dot products for
211    /// 2 parallel frames and adds the respective `init_f0` / `init_f1` accumulators.
212    /// Returns `(results_f0, results_f1)`.
213    ///
214    /// No alignment required.
215    ///
216    /// # Safety
217    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
218    unsafe fn dot_product_8x_f32_dual_accumulate(
219        weights: &[[f32; 8]],
220        state_f0: &[f32],
221        state_f1: &[f32],
222        init_f0: &[f32; 8],
223        init_f1: &[f32; 8],
224    ) -> ([f32; 8], [f32; 8]);
225
226    /// Fused accumulate variant: computes 16 simultaneous dot products with native f32
227    /// weights and adds the `init` accumulator.
228    ///
229    /// No alignment required.
230    ///
231    /// # Safety
232    /// `weights.len() >= state.len()`.
233    unsafe fn dot_product_16x_f32_accumulate(
234        weights: &[[f32; 16]],
235        state: &[f32],
236        init: &[f32; 16],
237    ) -> [f32; 16];
238
239    /// Fused accumulate dual-frame variant: computes 16 simultaneous dot products for
240    /// 2 parallel frames and adds the respective `init_f0` / `init_f1` accumulators.
241    /// Returns `(results_f0, results_f1)`.
242    ///
243    /// No alignment required.
244    ///
245    /// # Safety
246    /// `weights.len() >= max(state_f0.len(), state_f1.len())`.
247    unsafe fn dot_product_16x_f32_dual_accumulate(
248        weights: &[[f32; 16]],
249        state_f0: &[f32],
250        state_f1: &[f32],
251        init_f0: &[f32; 16],
252        init_f1: &[f32; 16],
253    ) -> ([f32; 16], [f32; 16]);
254
255    /// Computes 4 simultaneous BF16 dot products with separate weight vectors.
256    ///
257    /// No alignment required.
258    ///
259    /// # Safety
260    /// Each of `w0`, `w1`, `w2`, `w3` must have length `>= in_frame.len()`.
261    /// All slices must be valid and contain valid BF16 bit patterns.
262    unsafe fn dot_product_bf16_4x(
263        w0: &[u16],
264        w1: &[u16],
265        w2: &[u16],
266        w3: &[u16],
267        in_frame: &[u16],
268    ) -> [f32; 4];
269
270    /// Horizontal sum of `N` consecutive f32 values starting at `ptr`.
271    ///
272    /// No alignment required.
273    ///
274    /// # Safety
275    /// `ptr` must point to at least `N` valid, initialized `f32` elements.
276    unsafe fn horizontal_sum<const N: usize>(ptr: *const f32) -> f32;
277
278    // --- (B) GEMV/GEMM Fused ---
279
280    /// Fused add + GEMV kernel (f32 weights).
281    ///
282    /// Computes `out += weights^T * in_frame [+ bias]` in a single pass.
283    /// No alignment required.
284    ///
285    /// # Safety
286    /// Let `in_len = in_frame.len()`, `out_len = out_frame.len()`.
287    /// `weights.len() >= in_len * out_len`.
288    /// `bias.len() >= out_len` (when `do_bias` is `true`).
289    /// All slices must be valid; `out_frame` must not alias `in_frame`.
290    unsafe fn fused_add_gemv(
291        in_frame: &[f32],
292        weights: &[f32],
293        bias: &[f32],
294        out_frame: &mut [f32],
295        do_bias: bool,
296    );
297
298    /// Fused add + batch GEMM kernel (f32 weights).
299    ///
300    /// No alignment required. Frames are batched in groups of 4 (AVX2)
301    /// or 8 (AVX-512) for throughput.
302    ///
303    /// # Safety
304    /// `in_frames.len()` and `out_frames.len()` must divide exactly by `num_frames`.
305    /// Let `in_len = in_frames.len() / num_frames`, `out_len = out_frames.len() / num_frames`.
306    /// `weights.len() >= in_len * out_len`.
307    /// `bias.len() >= out_len` (when `do_bias` is `true`).
308    /// `num_frames > 0`. All slices must be valid and non-aliasing.
309    unsafe fn fused_add_gemm_batch(
310        in_frames: &[f32],
311        weights: &[f32],
312        bias: &[f32],
313        out_frames: &mut [f32],
314        num_frames: usize,
315        do_bias: bool,
316    );
317
318    /// Fused residual batch GEMM kernel (f32 weights).
319    ///
320    /// Computes `out = weights^T * in + residual [+ bias]`.
321    /// No alignment required.
322    ///
323    /// # Safety
324    /// Same preconditions as `fused_add_gemm_batch`, plus:
325    /// `residual.len()` must divide exactly by `num_frames`,
326    /// with `residual.len() / num_frames == out_len`.
327    unsafe fn fused_gemm_residual_batch(
328        in_frames: &[f32],
329        weights: &[f32],
330        bias: &[f32],
331        residual: &[f32],
332        out_frames: &mut [f32],
333        num_frames: usize,
334        do_bias: bool,
335    );
336
337    /// Fused residual batch GEMM kernel with native f32 weights.
338    ///
339    /// Used where the 1x1 projection operates on full-precision f32 weights.
340    /// Fuses GEMV + bias + residual addition into a single SIMD pass.
341    /// No alignment required.
342    ///
343    /// # Safety
344    /// Same preconditions as `fused_gemm_residual_batch`.
345    unsafe fn fused_gemm_residual_batch_f32(
346        in_frames: &[f32],
347        weights: &[f32],
348        bias: &[f32],
349        residual: &[f32],
350        out_frames: &mut [f32],
351        num_frames: usize,
352        do_bias: bool,
353    );
354
355    /// GEMV kernel with overwrite (f16c-quantized weights).
356    ///
357    /// Computes `out = weights^T * in [+ bias]` (no accumulation).
358    /// No alignment required.
359    ///
360    /// # Safety
361    /// Let `in_len = in_frame.len()`, `out_len = out_frame.len()`.
362    /// `weights.len() >= in_len * out_len`.
363    /// `bias.len() >= out_len` (when `do_bias` is `true`).
364    /// All slices must be valid; `out_frame` must not alias `in_frame`.
365    unsafe fn gemv_overwrite(
366        in_frame: &[f32],
367        weights: &[f32],
368        bias: &[f32],
369        out_frame: &mut [f32],
370        do_bias: bool,
371    );
372
373    /// GEMV overwrite in batch (f32 weights).
374    ///
375    /// Delegates to per-frame `gemv_overwrite`. No alignment required.
376    ///
377    /// # Safety
378    /// Same preconditions as `fused_add_gemm_batch`.
379    unsafe fn gemv_overwrite_batch(
380        in_frames: &[f32],
381        weights: &[f32],
382        bias: &[f32],
383        out_frames: &mut [f32],
384        num_frames: usize,
385        do_bias: bool,
386    );
387
388    /// GEMV kernel with overwrite (BF16 input and weights).
389    ///
390    /// No alignment required.
391    ///
392    /// # Safety
393    /// Let `in_len = in_frame.len()`, `out_len = out_frame.len()`.
394    /// `weights.len() >= in_len * out_len`.
395    /// `bias.len() >= out_len` (when `do_bias` is `true`).
396    /// All slices must be valid and contain valid BF16 bit patterns.
397    unsafe fn gemv_overwrite_bf16(
398        in_frame: &[u16],
399        weights: &[u16],
400        bias: &[f32],
401        out_frame: &mut [f32],
402        do_bias: bool,
403    );
404
405    /// GEMV overwrite in batch using native f32 weights (always adds bias).
406    ///
407    /// Used for mixed-precision head projection where the final stage requires
408    /// full FP32 precision while the backbone runs quantized.
409    /// No alignment required. Shape-dependent dispatch selects an optimized
410    /// path based on `out_len` (1, ≤4, or ≥8).
411    ///
412    /// # Safety
413    /// `in_frames.len()` and `out_frames.len()` must divide exactly by `num_frames`.
414    /// Let `in_len = in_frames.len() / num_frames`, `out_len = out_frames.len() / num_frames`.
415    /// `weights.len() == in_len * out_len`.
416    /// `bias.len() >= out_len`.
417    /// `num_frames > 0`. All slices must be valid and non-aliasing.
418    unsafe fn gemv_with_bias_f32(
419        in_frames: &[f32],
420        weights: &[f32],
421        bias: &[f32],
422        out_frames: &mut [f32],
423        num_frames: usize,
424    );
425
426    /// GEMV overwrite in batch using native f32 weights (no bias).
427    ///
428    /// Used for mixed-precision head projection where the final stage requires
429    /// full FP32 precision while the backbone runs quantized.
430    /// Overwrites without adding bias. No alignment required.
431    ///
432    /// # Safety
433    /// `in_frames.len()` and `out_frames.len()` must divide exactly by `num_frames`.
434    /// Let `in_len = in_frames.len() / num_frames`, `out_len = out_frames.len() / num_frames`.
435    /// `weights.len() == in_len * out_len`.
436    /// `num_frames > 0`. All slices must be valid and non-aliasing.
437    unsafe fn gemv_no_bias_f32(
438        in_frames: &[f32],
439        weights: &[f32],
440        out_frames: &mut [f32],
441        num_frames: usize,
442    );
443
444    /// GEMV with overwrite for 4 simultaneous LSTM gates (f16c-quantized).
445    ///
446    /// Computes `gates = W^T * in_frame [+ bias]` where the weight matrix
447    /// contains concatenated gate weights: `[W_input | W_forget | W_cell | W_output]`.
448    /// No alignment required.
449    ///
450    /// # Safety
451    /// `weights.len() == in_frame.len() * hidden_size * 4`.
452    /// `out_gates.len() >= hidden_size * 4`.
453    /// `bias.len() >= hidden_size * 4` (when `do_bias` is `true`).
454    /// All slices must be valid and non-aliasing.
455    unsafe fn gemv_overwrite_4gate(
456        in_frame: &[f32],
457        weights: &[u16],
458        bias: &[f32],
459        out_gates: &mut [f32],
460        hidden_size: usize,
461        do_bias: bool,
462    );
463
464    /// GEMV with overwrite for 4 simultaneous LSTM gates (BF16 input and weights).
465    ///
466    /// No alignment required.
467    ///
468    /// # Safety
469    /// Same preconditions as `gemv_overwrite_4gate`, plus all `u16` slices
470    /// must contain valid BF16 bit patterns.
471    unsafe fn gemv_overwrite_bf16_4gate(
472        in_frame: &[u16],
473        weights: &[u16],
474        bias: &[f32],
475        out_gates: &mut [f32],
476        hidden_size: usize,
477        do_bias: bool,
478    );
479
480    // --- (C) Activations ---
481
482    /// Accumulates `src` into `dest`: `dest[i] += src[i]`.
483    ///
484    /// No alignment required.
485    ///
486    /// # Safety
487    /// `dest.len() == src.len()`. Both slices must be valid
488    /// and must not alias.
489    unsafe fn accumulate_head(dest: &mut [f32], src: &[f32]);
490
491    /// Fused Tanh + Head Accumulate.
492    ///
493    /// Computes `head_input[i] += tanh(block[i])`.
494    /// No alignment required.
495    ///
496    /// # Safety
497    /// `head_input.len() == block.len()`. Both slices must be valid
498    /// and must not alias.
499    unsafe fn tanh_and_accumulate_block(head_input: &mut [f32], block: &mut [f32]);
500
501    /// Fused Gated Activation + Head Accumulate.
502    ///
503    /// For each channel group of size `ch`, applies `tanh` to the gated portion
504    /// and accumulates into the head:
505    /// `head_input[i*ch + j] += tanh(block[i*ch + j])`.
506    /// No alignment required.
507    ///
508    /// # Safety
509    /// `head_input.len() == block.len()` and both divide exactly by `ch`.
510    /// Both slices must be valid and must not alias. `ch > 0`.
511    unsafe fn gated_activation_and_accumulate_block(
512        head_input: &mut [f32],
513        block: &mut [f32],
514        ch: usize,
515    );
516
517    /// Fused Tanh + Head Overwrite.
518    ///
519    /// Computes `head_input[i] = tanh(block[i])`.
520    /// No alignment required.
521    ///
522    /// # Safety
523    /// `head_input.len() == block.len()`. Both slices must be valid
524    /// and must not alias.
525    unsafe fn tanh_and_overwrite_block(head_input: &mut [f32], block: &mut [f32]);
526
527    /// Fused Seed + Tanh + Head Accumulate.
528    ///
529    /// Computes `head_input[i] = seed[i] + tanh(block[i])`.
530    /// Eliminates the separate `copy_from_slice(seed)` before
531    /// `tanh_and_accumulate_block` on the first layer of a cascaded array.
532    /// No alignment required.
533    ///
534    /// # Safety
535    /// `head_input.len() == block.len() == seed.len()`.
536    /// All slices must be valid and must not alias.
537    unsafe fn tanh_and_accumulate_with_seed(
538        head_input: &mut [f32],
539        block: &mut [f32],
540        seed: &[f32],
541    );
542
543    /// Fused Gated Activation + Head Overwrite.
544    ///
545    /// For each channel group of size `ch`, applies `tanh` to the gated portion
546    /// and overwrites the head:
547    /// `head_input[i*ch + j] = tanh(block[i*ch + j])`.
548    /// No alignment required.
549    ///
550    /// # Safety
551    /// `head_input.len() == block.len()` and both divide exactly by `ch`.
552    /// Both slices must be valid and must not alias. `ch > 0`.
553    unsafe fn gated_activation_and_overwrite_block(
554        head_input: &mut [f32],
555        block: &mut [f32],
556        ch: usize,
557    );
558
559    /// Computes `max(energy(l), energy(r))` as max mean-square energy.
560    ///
561    /// No alignment required.
562    ///
563    /// # Safety
564    /// `l` and `r` must be valid slices.
565    unsafe fn compute_energy_stereo(l: &[f32], r: &[f32]) -> f32;
566
567    /// Computes the mean-square energy: `(1/N) * Σ x_i²`.
568    ///
569    /// No alignment required.
570    ///
571    /// # Safety
572    /// `data` must be a valid slice.
573    unsafe fn compute_energy(data: &[f32]) -> f32;
574
575    /// Computes `max(|a[i] - b[i]|)`.
576    ///
577    /// No alignment required.
578    ///
579    /// # Safety
580    /// `a.len() == b.len()`. Both slices must be valid.
581    unsafe fn compute_max_diff(a: &[f32], b: &[f32]) -> f32;
582
583    /// Computes `(max(|left|), max(|right|))`.
584    ///
585    /// No alignment required.
586    ///
587    /// # Safety
588    /// `left.len() == right.len()`. Both slices must be valid.
589    unsafe fn compute_peak_abs_stereo(left: &[f32], right: &[f32]) -> (f32, f32);
590
591    /// Computes `max(|x[i]|)` for a single channel.
592    ///
593    /// No alignment required.
594    ///
595    /// # Safety
596    /// `data` must be a valid slice.
597    unsafe fn compute_peak_abs_mono(data: &[f32]) -> f32;
598
599    /// Applies Tanh element-wise to `slice`.
600    ///
601    /// No alignment required.
602    ///
603    /// # Safety
604    /// `slice` must be a valid mutable slice.
605    unsafe fn tanh_slice(slice: &mut [f32]);
606
607    /// Applies Sigmoid element-wise to `slice`.
608    ///
609    /// No alignment required.
610    ///
611    /// # Safety
612    /// `slice` must be a valid mutable slice.
613    unsafe fn sigmoid_slice(slice: &mut [f32]);
614
615    /// Applies high-fidelity Tanh element-wise to `slice`.
616    ///
617    /// Uses polynomial exp-based approximation (~2.4e-7 error vs `f32::tanh`).
618    /// Higher cost, lower error, lower aliasing.  No alignment required.
619    ///
620    /// # Safety
621    /// `slice` must be a valid mutable slice.
622    unsafe fn tanh_slice_hf(slice: &mut [f32]) {
623        // SAFETY: default implementation delegates to the standard tanh_slice,
624        // which has its own documented safety invariants (valid mutable slice).
625        unsafe {
626            Self::tanh_slice(slice);
627        }
628    }
629
630    /// Applies high-fidelity Sigmoid element-wise to `slice`.
631    ///
632    /// Uses polynomial exp-based approximation (~2.1e-7 error vs `f32::exp`).
633    /// No alignment required.
634    ///
635    /// # Safety
636    /// `slice` must be a valid mutable slice.
637    unsafe fn sigmoid_slice_hf(slice: &mut [f32]) {
638        // SAFETY: default implementation delegates to the standard sigmoid_slice,
639        // which has its own documented safety invariants (valid mutable slice).
640        unsafe {
641            Self::sigmoid_slice(slice);
642        }
643    }
644
645    /// Applies ReLU element-wise to `slice`.
646    ///
647    /// No alignment required.
648    ///
649    /// # Safety
650    /// `slice` must be a valid mutable slice.
651    unsafe fn relu_slice(slice: &mut [f32]);
652
653    /// Applies PReLU element-wise: `slice[i] = max(0, slice[i]) + slopes[i] * min(0, slice[i])`.
654    ///
655    /// No alignment required.
656    ///
657    /// # Safety
658    /// `slice` and `slopes` must be valid slices with `slopes.len() >= slice.len()`.
659    unsafe fn prelu_slice(slice: &mut [f32], slopes: &[f32]);
660
661    /// Applies Softsign element-wise to `slice`.
662    ///
663    /// No alignment required.
664    ///
665    /// # Safety
666    /// `slice` must be a valid mutable slice.
667    unsafe fn softsign_slice(slice: &mut [f32]);
668
669    /// Applies SiLU element-wise to `slice`.
670    ///
671    /// No alignment required.
672    ///
673    /// # Safety
674    /// `slice` must be a valid mutable slice.
675    unsafe fn silu_slice(slice: &mut [f32]);
676
677    /// Applies HardTanh element-wise to `slice`.
678    ///
679    /// No alignment required.
680    ///
681    /// # Safety
682    /// `slice` must be a valid mutable slice.
683    unsafe fn hard_tanh_slice(slice: &mut [f32]);
684
685    /// Applies HardSwish element-wise to `slice`.
686    ///
687    /// No alignment required.
688    ///
689    /// # Safety
690    /// `slice` must be a valid mutable slice.
691    unsafe fn hard_swish_slice(slice: &mut [f32]);
692
693    /// Applies FastTanh (rational approximation) element-wise to `slice`.
694    ///
695    /// No alignment required.
696    ///
697    /// # Safety
698    /// `slice` must be a valid mutable slice.
699    unsafe fn fast_tanh_slice(slice: &mut [f32]);
700
701    /// Applies LeakyHardTanh element-wise.
702    ///
703    /// Uses piecewise-linear mapping with configurable saturation and leak slopes.
704    /// No alignment required.
705    ///
706    /// # Safety
707    /// `slice` must be a valid mutable slice.
708    unsafe fn leaky_hard_tanh_slice(
709        slice: &mut [f32],
710        min_val: f32,
711        max_val: f32,
712        min_slope: f32,
713        max_slope: f32,
714    );
715
716    /// Tanh activation on a block (shorthand for `tanh_slice`).
717    ///
718    /// No alignment required.
719    ///
720    /// # Safety
721    /// `buf` must be a valid mutable slice.
722    unsafe fn activation_tanh_block(buf: &mut [f32]);
723
724    // --- (D) Conversions ---
725
726    /// Conversion from F32 to BF16: `dest[i] = bf16(src[i])`.
727    ///
728    /// No alignment required.
729    ///
730    /// # Safety
731    /// `src.len() == dest.len()`. Both slices must be valid
732    /// and must not alias.
733    unsafe fn f32_to_bf16(src: &[f32], dest: &mut [u16]);
734
735    /// Stores the contents of a SIMD register as BF16 (truncated).
736    ///
737    /// Writes to memory via unaligned 128-bit store. No alignment required.
738    ///
739    /// # Safety
740    /// `ptr` must be valid for at least enough `u16` elements to hold
741    /// the register contents (8 elements for AVX2 `__m256`, 16 for AVX-512 `__m512`).
742    unsafe fn store_bf16(ptr: *mut u16, v: Self::V);
743
744    // --- (E) LSTM Gates ---
745
746    /// Fused kernel for dynamic LSTM gate processing.
747    ///
748    /// Combines sigmoid/tanh activation of the 4 gates with cell and hidden
749    /// state update in a single SIMD pass. No alignment required.
750    ///
751    /// # Safety
752    /// `gates.len() >= hidden_size * 4`.
753    /// `cell_state.len() == hidden_size`.
754    /// `cell_error.len() == hidden_size`.
755    /// `hidden_state.len() == hidden_size`.
756    /// All slices must be valid and non-aliasing. `hidden_size > 0`.
757    unsafe fn fused_lstm_gates_dyn(
758        gates: &mut [f32],
759        cell_state: &mut [f32],
760        cell_error: &mut [f32],
761        hidden_state: &mut [f32],
762        hidden_size: usize,
763    );
764
765    /// Stereo convolution (used in the resampler).
766    ///
767    /// Computes the dot product between a coefficient bank and two input buffers (L/R).
768    /// Coefficients are loaded with **aligned** SIMD loads for performance.
769    ///
770    /// # Safety
771    /// `coeffs` must be aligned to 32 bytes (AVX2 `__m256`) or 64 bytes (AVX-512 `__m512`)
772    /// depending on the concrete implementation.
773    /// `coeffs`, `input_l`, and `input_r` must each point to at least `taps` valid `f32` elements.
774    /// Input pointers use unaligned loads — no alignment required.
775    unsafe fn convolve_stereo(
776        coeffs: *const f32,
777        input_l: *const f32,
778        input_r: *const f32,
779        taps: usize,
780    ) -> (f32, f32);
781
782    /// Dual stereo convolution (reuses input loads).
783    ///
784    /// Computes the dot product between two coefficient banks and two input buffers (L/R).
785    /// Both coefficient banks are loaded with **aligned** SIMD loads.
786    ///
787    /// # Safety
788    /// `coeffs0` and `coeffs1` must be aligned to 32 bytes (AVX2) or 64 bytes (AVX-512).
789    /// `coeffs0`, `coeffs1`, `input_l`, and `input_r` must each point to at least `taps`
790    /// valid `f32` elements. Input pointers use unaligned loads — no alignment required.
791    unsafe fn convolve_stereo_dual(
792        coeffs0: *const f32,
793        coeffs1: *const f32,
794        input_l: *const f32,
795        input_r: *const f32,
796        taps: usize,
797    ) -> ((f32, f32), (f32, f32));
798
799    /// Mono convolution (used in the resampler).
800    ///
801    /// Coefficients are loaded with **aligned** SIMD loads for performance.
802    ///
803    /// # Safety
804    /// `coeffs` must be aligned to 32 bytes (AVX2 `__m256`) or 64 bytes (AVX-512 `__m512`).
805    /// `coeffs` and `input` must each point to at least `taps` valid `f32` elements.
806    /// Input uses unaligned loads — no alignment required.
807    unsafe fn convolve_mono(coeffs: *const f32, input: *const f32, taps: usize) -> f32;
808
809    /// Dual mono convolution (reuses input loads).
810    ///
811    /// Both coefficient banks are loaded with **aligned** SIMD loads.
812    ///
813    /// # Safety
814    /// `coeffs0` and `coeffs1` must be aligned to 32 bytes (AVX2) or 64 bytes (AVX-512).
815    /// `coeffs0`, `coeffs1`, and `input` must each point to at least `taps`
816    /// valid `f32` elements. Input uses unaligned loads — no alignment required.
817    unsafe fn convolve_mono_dual(
818        coeffs0: *const f32,
819        coeffs1: *const f32,
820        input: *const f32,
821        taps: usize,
822    ) -> (f32, f32);
823
824    /// Applies gain and detects clipping in mono.
825    ///
826    /// Computes `data[i] *= gain` and returns `true` if any `|data[i]| > 1.0`.
827    /// No alignment required.
828    ///
829    /// # Safety
830    /// `data` must be a valid mutable slice.
831    unsafe fn apply_gain_and_detect_clipping_mono(data: &mut [f32], gain: f32) -> bool;
832
833    /// Applies gain and detects clipping in stereo.
834    ///
835    /// No alignment required.
836    ///
837    /// # Safety
838    /// `left` and `right` must be valid mutable slices.
839    unsafe fn apply_gain_and_detect_clipping_stereo(
840        left: &mut [f32],
841        right: &mut [f32],
842        gain: f32,
843    ) -> bool;
844
845    /// Applies constant gain in stereo without clipping detection.
846    ///
847    /// No alignment required.
848    ///
849    /// # Safety
850    /// `left` and `right` must be valid mutable slices.
851    unsafe fn apply_gain_stereo(left: &mut [f32], right: &mut [f32], gain: f32);
852
853    /// Applies constant gain to a mono buffer.
854    ///
855    /// No alignment required.
856    ///
857    /// # Safety
858    /// `data` must be a valid mutable slice.
859    unsafe fn apply_gain(data: &mut [f32], gain: f32);
860
861    /// Applies a linear gain ramp to a mono buffer.
862    ///
863    /// Computes `data[i] *= start + i * step`.
864    /// No alignment required.
865    ///
866    /// # Safety
867    /// `data` must be a valid mutable slice.
868    unsafe fn apply_ramp(data: &mut [f32], start: f32, step: f32);
869
870    /// Crossfade blend: `out[i] = out[i] * (1 - t) + pending[i] * t`.
871    ///
872    /// Computed as `fma(pending[i] - out[i], t, out[i])` for single-rounding precision.
873    /// No alignment required.
874    ///
875    /// # Safety
876    /// `out` and `pending` must be valid slices.
877    /// Effective length is `min(out.len(), pending.len())`.
878    unsafe fn crossfade_blend_mono(out: &mut [f32], pending: &[f32], t: f32);
879
880    /// Applies a linear gain ramp in stereo.
881    ///
882    /// No alignment required.
883    ///
884    /// # Safety
885    /// `left` and `right` must be valid mutable slices.
886    unsafe fn apply_ramp_stereo(left: &mut [f32], right: &mut [f32], start: f32, step: f32);
887
888    /// Fused gain + dither: `data[i] = data[i] * gain + offset` in a single pass.
889    ///
890    /// Eliminates the overhead of two separate passes (`apply_gain` then
891    /// `apply_dither_add`) on the same buffer. Uses FMA for single-rounding
892    /// precision where the ISA supports it.
893    ///
894    /// No alignment required.
895    ///
896    /// # Safety
897    /// `data` must be a valid mutable slice.
898    unsafe fn apply_gain_then_dither(data: &mut [f32], gain: f32, offset: f32);
899
900    /// Adds a broadcast constant (dither offset) to every element of a mono buffer.
901    ///
902    /// No alignment required.
903    ///
904    /// # Safety
905    /// `data` must be a valid mutable slice.
906    unsafe fn apply_dither_add(data: &mut [f32], offset: f32);
907
908    // --- (F) Complex MAC ---
909
910    /// Complex multiply-accumulate (overwrite): spectral kernel.
911    ///
912    /// Computes the element-wise complex multiplication of two vectors
913    /// `(h_re, h_im)` and `(x_re, x_im)` writing the result to `(out_re, out_im)`.
914    ///
915    /// `out_re[i] = h_re[i] * x_re[i] - h_im[i] * x_im[i]`
916    /// `out_im[i] = h_re[i] * x_im[i] + h_im[i] * x_re[i]`
917    ///
918    /// No alignment required.
919    ///
920    /// # Safety
921    /// `n = h_re.len() == h_im.len() == x_re.len() == x_im.len() == out_re.len() == out_im.len()`.
922    /// All slices must be valid and must not alias.
923    unsafe fn complex_mac_overwrite(
924        h_re: &[f32],
925        h_im: &[f32],
926        x_re: &[f32],
927        x_im: &[f32],
928        out_re: &mut [f32],
929        out_im: &mut [f32],
930    );
931
932    /// Complex multiply-accumulate (accumulate): spectral kernel.
933    ///
934    /// Computes the element-wise complex multiplication of two vectors
935    /// `(h_re, h_im)` and `(x_re, x_im)` adding the result to `(acc_re, acc_im)`.
936    ///
937    /// `acc_re[i] += h_re[i] * x_re[i] - h_im[i] * x_im[i]`
938    /// `acc_im[i] += h_re[i] * x_im[i] + h_im[i] * x_re[i]`
939    ///
940    /// No alignment required.
941    ///
942    /// # Safety
943    /// `n = h_re.len() == h_im.len() == x_re.len() == x_im.len() == acc_re.len() == acc_im.len()`.
944    /// All slices must be valid and must not alias.
945    unsafe fn complex_mac_accumulate(
946        h_re: &[f32],
947        h_im: &[f32],
948        x_re: &[f32],
949        x_im: &[f32],
950        acc_re: &mut [f32],
951        acc_im: &mut [f32],
952    );
953
954    // --- (G) FFT Butterfly ---
955
956    /// SIMD Radix-2 DIT FFT butterfly stage for one group.
957    ///
958    /// Processes `half` butterflies for a single group starting at
959    /// `group_start`. Each butterfly kernel computes:
960    ///
961    /// ```text
962    /// t_re = w_re * re_bot - w_im * im_bot
963    /// t_im = w_re * im_bot + w_im * re_bot
964    /// re[top] = re[top] + t_re    re[bot] = re[top] - t_re
965    /// im[top] = im[top] + t_im    im[bot] = im[top] - t_im
966    /// ```
967    ///
968    /// where `top = group_start + j`, `bot = group_start + j + half`,
969    /// and the twiddle factors `(w_re, w_im)` for `j = 0..half` are
970    /// stored contiguously in `tw_re`/`tw_im`.
971    ///
972    /// When `inverse` is `true`, the twiddle imaginary part is
973    /// conjugated (`w_im = -w_im`) before multiplication.
974    ///
975    /// # Safety
976    /// - `re` and `im` point to `n` valid `f32` values.
977    /// - `tw_re`/`tw_im` point to `half` valid `f32` values each.
978    /// - `group_start + 2 * half <= n`.
979    /// - CPU supports the required SIMD ISA (verified by dispatch).
980    /// - Source and destination may alias within the same array
981    ///   (in-place butterfly).
982    unsafe fn fft_butterfly_stage(
983        re: *mut f32,
984        im: *mut f32,
985        half: usize,
986        tw_re: *const f32,
987        tw_im: *const f32,
988        group_start: usize,
989        inverse: bool,
990    );
991
992    // --- (H) BatchNorm ---
993
994    /// Frame-major batch normalization affine transform.
995    ///
996    /// Computes `data[f * n_ch + c] = data[f * n_ch + c] * scale[c] + offset[c]`
997    /// for all frames and channels in a single pass.
998    /// Processes frames in the outer loop for contiguous channel access.
999    ///
1000    /// # Safety
1001    /// `data.len() == num_frames * n_ch`. `scale` and `offset` must each have
1002    /// at least `n_ch` elements. `num_frames > 0`.
1003    unsafe fn batch_norm_process(
1004        data: &mut [f32],
1005        scale: &[f32],
1006        offset: &[f32],
1007        n_ch: usize,
1008        num_frames: usize,
1009    );
1010}