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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

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

//! AVX-512 accumulation and activation kernels for WaveNet.

use core::arch::x86_64::*;

/// Applies the "gate" activation (tanh * sigmoid) to audio blocks.
/// Imagine each sound goes through two filters: one that shapes the timbre (tanh)
/// and another that controls the intensity (sigmoid). The result is added to "head_input".
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn gated_activation_and_accumulate_block_avx512(
    head_input: &mut [f32],
    block: &mut [f32],
    ch: usize,
) {
    let num_frames = head_input.len() / ch;
    for f in 0..num_frames {
        let block_offset = f * 2 * ch;
        let head_offset = f * ch;
        let mut c = 0;
        while c + 16 <= ch {
            let z1 = _mm512_loadu_ps(block.as_ptr().add(block_offset + c));
            let z2 = _mm512_loadu_ps(block.as_ptr().add(block_offset + ch + c));

            let (tanh_z1, sig_z2) =
                crate::math::activations::simd_tanh_sigmoid_dual_poly_avx512(z1, z2);
            let activated = _mm512_mul_ps(tanh_z1, sig_z2);

            _mm512_storeu_ps(block.as_mut_ptr().add(block_offset + c), activated);

            let vh = _mm512_loadu_ps(head_input.as_ptr().add(head_offset + c));
            _mm512_storeu_ps(
                head_input.as_mut_ptr().add(head_offset + c),
                _mm512_add_ps(vh, activated),
            );
            c += 16;
        }
        while c < ch {
            let z1 = block[block_offset + c];
            let z2 = block[block_offset + ch + c];
            let activated = z1.tanh() * (1.0 / (1.0 + (-z2).exp()));
            block[block_offset + c] = activated;
            let acc = head_input[head_offset + c] as f64 + activated as f64;
            head_input[head_offset + c] = acc as f32;
            c += 1;
        }
    }
}

/// Applies the "gate" activation (tanh * sigmoid) to audio blocks.
/// Overwrites to "head_input".
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn gated_activation_and_overwrite_block_avx512(
    head_input: &mut [f32],
    block: &mut [f32],
    ch: usize,
) {
    let num_frames = head_input.len() / ch;
    for f in 0..num_frames {
        let block_offset = f * 2 * ch;
        let head_offset = f * ch;
        let mut c = 0;
        while c + 16 <= ch {
            let z1 = _mm512_loadu_ps(block.as_ptr().add(block_offset + c));
            let z2 = _mm512_loadu_ps(block.as_ptr().add(block_offset + ch + c));

            let (tanh_z1, sig_z2) =
                crate::math::activations::simd_tanh_sigmoid_dual_poly_avx512(z1, z2);
            let activated = _mm512_mul_ps(tanh_z1, sig_z2);

            _mm512_storeu_ps(block.as_mut_ptr().add(block_offset + c), activated);
            _mm512_storeu_ps(head_input.as_mut_ptr().add(head_offset + c), activated);
            c += 16;
        }
        while c < ch {
            let z1 = block[block_offset + c];
            let z2 = block[block_offset + ch + c];
            let activated = z1.tanh() * (1.0 / (1.0 + (-z2).exp()));
            block[block_offset + c] = activated;
            head_input[head_offset + c] = activated;
            c += 1;
        }
    }
}

/// Accumulates src into dest using AVX-512.
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn accumulate_head_avx512(dest: &mut [f32], src: &[f32]) {
    let len = dest.len();
    let mut i = 0;
    while i + 16 <= len {
        let vs = _mm512_loadu_ps(src.as_ptr().add(i));
        let vd = _mm512_loadu_ps(dest.as_ptr().add(i));
        _mm512_storeu_ps(dest.as_mut_ptr().add(i), _mm512_add_ps(vd, vs));
        i += 16;
    }
    if i < len {
        let mask = _cvtu32_mask16((1u32 << (len - i)) - 1);
        let vs = _mm512_maskz_loadu_ps(mask, src.as_ptr().add(i));
        let vd = _mm512_maskz_loadu_ps(mask, dest.as_ptr().add(i));
        _mm512_mask_storeu_ps(dest.as_mut_ptr().add(i), mask, _mm512_add_ps(vd, vs));
    }
}

/// Applies tanh in-place on block and accumulates into head_input using AVX-512.
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn tanh_and_accumulate_block_avx512(head_input: &mut [f32], block: &mut [f32]) {
    let len = block.len();
    let mut i = 0;
    while i + 16 <= len {
        let vb = _mm512_loadu_ps(block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_storeu_ps(block.as_mut_ptr().add(i), vt);

        let vh = _mm512_loadu_ps(head_input.as_ptr().add(i));
        _mm512_storeu_ps(head_input.as_mut_ptr().add(i), _mm512_add_ps(vh, vt));
        i += 16;
    }
    if i < len {
        let mask = _cvtu32_mask16((1u32 << (len - i)) - 1);
        let vb = _mm512_maskz_loadu_ps(mask, block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_mask_storeu_ps(block.as_mut_ptr().add(i), mask, vt);

        let vh = _mm512_maskz_loadu_ps(mask, head_input.as_ptr().add(i));
        _mm512_mask_storeu_ps(head_input.as_mut_ptr().add(i), mask, _mm512_add_ps(vh, vt));
    }
}

/// Applies tanh in-place on block and overwrites head_input using AVX-512.
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn tanh_and_overwrite_block_avx512(head_input: &mut [f32], block: &mut [f32]) {
    let len = block.len();
    let mut i = 0;
    while i + 16 <= len {
        let vb = _mm512_loadu_ps(block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_storeu_ps(block.as_mut_ptr().add(i), vt);
        _mm512_storeu_ps(head_input.as_mut_ptr().add(i), vt);
        i += 16;
    }
    if i < len {
        let mask = _cvtu32_mask16((1u32 << (len - i)) - 1);
        let vb = _mm512_maskz_loadu_ps(mask, block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_mask_storeu_ps(block.as_mut_ptr().add(i), mask, vt);
        _mm512_mask_storeu_ps(head_input.as_mut_ptr().add(i), mask, vt);
    }
}

/// Fused Seed + Tanh + Head Accumulate using AVX-512.
///
/// Computes `head_input[i] = seed[i] + tanh(block[i])`.
/// Eliminates the separate `copy_from_slice(seed)` before `tanh_and_accumulate_block`.
#[target_feature(enable = "avx512f,avx512vl")]
pub unsafe fn tanh_and_accumulate_with_seed_avx512(
    head_input: &mut [f32],
    block: &mut [f32],
    seed: &[f32],
) {
    let len = block.len();
    let mut i = 0;
    while i + 16 <= len {
        let vb = _mm512_loadu_ps(block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_storeu_ps(block.as_mut_ptr().add(i), vt);

        let vs = _mm512_loadu_ps(seed.as_ptr().add(i));
        _mm512_storeu_ps(head_input.as_mut_ptr().add(i), _mm512_add_ps(vs, vt));
        i += 16;
    }
    if i < len {
        let mask = _cvtu32_mask16((1u32 << (len - i)) - 1);
        let vb = _mm512_maskz_loadu_ps(mask, block.as_ptr().add(i));
        let vt = crate::math::activations::simd_tanh_poly_avx512(vb);
        _mm512_mask_storeu_ps(block.as_mut_ptr().add(i), mask, vt);

        let vs = _mm512_maskz_loadu_ps(mask, seed.as_ptr().add(i));
        _mm512_mask_storeu_ps(head_input.as_mut_ptr().add(i), mask, _mm512_add_ps(vs, vt));
    }
}