NeuralAmpModeler-rs 3.0.0

An opinionated, high-performance Neural Amp Modeler (NAM) client and core implementation in Rust for Linux/PipeWire and CLAP plugins.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.

//! 2-layer LSTM model with pipelining and SIMD dispatch.

use super::layer::LstmLayer;

macro_rules! define_lstm2_process_pipelined {
    (
        $fn_name:ident,
        $target_meta:meta,
        $layer_proc:ident,
    ) => {
        #[$target_meta]
        unsafe fn $fn_name(&mut self, input: &[f32], output: &mut [f32]) {
            unsafe {
                let len = input.len();
                if len >= 1 {
                    self.layer1.$layer_proc(&[input[0]]);
                    let mut prev_h1 = [0.0; H];
                    prev_h1.copy_from_slice(self.layer1.get_hidden_state());

                    for i in 1..len {
                        self.layer1.$layer_proc(&[input[i]]);
                        self.layer2.$layer_proc(&prev_h1);
                        let h_f32 = self.layer2.get_hidden_state();
                        output[i - 1] =
                            $crate::math::common::scalar_ref::dot_product_f32_native_kahan(
                                h_f32,
                                &self.head_weights_f32,
                            ) + self.head_bias;
                        prev_h1.copy_from_slice(self.layer1.get_hidden_state());
                    }

                    self.layer2.$layer_proc(&prev_h1);
                    let h_f32 = self.layer2.get_hidden_state();
                    output[len - 1] =
                        $crate::math::common::scalar_ref::dot_product_f32_native_kahan(
                            h_f32,
                            &self.head_weights_f32,
                        ) + self.head_bias;
                }
            }
        }
    };
}

/// 2-layer LSTM model.
pub struct LstmModel2<const H: usize, const H1_IH: usize, const H2_IH: usize, const H_H4: usize> {
    /// Model layer 1.
    pub layer1: LstmLayer<1, H, H1_IH, H_H4>,
    /// Model layer 2.
    pub layer2: LstmLayer<H, H, H2_IH, H_H4>,
    /// Output head weights.
    pub head_weights: [f32; H],
    /// Output head weights.
    pub head_weights_f32: [f32; H],
    /// Output head bias.
    pub head_bias: f32,
    /// Whether to execute prewarm during `reset()`. Default: `true`.
    pub prewarm_on_reset: bool,
    /// Expected sample rate (Hz) for prewarm calculation. Default: `48000.0`.
    pub expected_sample_rate: f64,
}

impl<const H: usize, const H1_IH: usize, const H2_IH: usize, const H_H4: usize>
    LstmModel2<H, H1_IH, H2_IH, H_H4>
{
    /// Creates a new 2-layer LSTM model.
    pub fn new() -> Self {
        Self {
            layer1: LstmLayer::new(),
            layer2: LstmLayer::new(),
            head_weights: [0.0f32; H],
            head_weights_f32: [0.0f32; H],
            head_bias: 0.0,
            prewarm_on_reset: true,
            expected_sample_rate: 48000.0,
        }
    }
    define_lstm2_process_pipelined!(
        process_avx2,
        target_feature(enable = "avx2,fma,f16c"),
        process_sample_avx2,
    );

    define_lstm2_process_pipelined!(
        process_avx512,
        target_feature(enable = "avx512f,avx512vl"),
        process_sample_avx512,
    );
    /// Processes an audio block through the model (SIMD dispatch).
    pub fn process(&mut self, input: &[f32], output: &mut [f32]) {
        unsafe {
            crate::math::common::dispatch_simd!(
                @self,
                process_avx512,
                process_avx512,
                process_avx2,
                input,
                output
            );
        }
    }
    /// Scalar processing (fallback).
    ///
    /// # Note
    /// Exclusively for parity tests. Extremely slow.
    pub fn process_scalar(&mut self, input: &[f32], output: &mut [f32]) {
        for i in 0..input.len() {
            self.layer1.process_sample_scalar(&[input[i]]);
            self.layer2
                .process_sample_scalar(self.layer1.get_hidden_state());
            let hidden2 = self.layer2.get_hidden_state();
            let dot = crate::math::common::scalar_ref::dot_product_f32_native_kahan(
                hidden2,
                &self.head_weights_f32,
            );
            output[i] = dot + self.head_bias;
        }
    }
    /// Resets the internal states.
    pub fn reset_states(&mut self) {
        self.layer1.reset_states();
        self.layer2.reset_states();
    }
}

impl<const H: usize, const H1_IH: usize, const H2_IH: usize, const H_H4: usize> Default
    for LstmModel2<H, H1_IH, H2_IH, H_H4>
{
    fn default() -> Self {
        Self::new()
    }
}