nam_rs/models/wavenet/
post_stack_head.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
4use crate::math::common::AlignedVec;
5use crate::math::common::SimdMath;
6use crate::models::a2::activations::ActivationType;
7
8use super::common::{WAVENET_MAX_NUM_FRAMES, WaveNetLayerState};
9use super::conv1d_dyn::Conv1dDyn;
10use crate::loader::nam_json::model::HeadConfig;
11
12/// Post-stack head sub-object for WaveNet / ConvNet architectures.
13///
14/// Contains a causal Conv1D + activation that processes the signal
15/// after the stack of layer arrays, before the final `head_scale` gain.
16/// Mirrors the `_Head` structure in NAMCore's `convnet.h:108-118`.
17#[derive(Clone)]
18#[repr(align(64))]
19pub struct PostStackHead {
20    /// Causal 1D convolution (dynamic runtime dimensions).
21    pub conv: Conv1dDyn,
22    /// Activation function applied after convolution.
23    pub activation: ActivationType,
24    /// Ring buffer state for causal convolution lookback.
25    pub state: WaveNetLayerState,
26    /// Scratch buffer for convolution output (out_ch * WAVENET_MAX_NUM_FRAMES).
27    scratch: AlignedVec<f32>,
28}
29
30impl PostStackHead {
31    /// Creates a new `PostStackHead` from the parsed `HeadConfig` and the
32    /// input channel count from the last layer array.
33    ///
34    /// Missing fields in `HeadConfig` fall back to sensible defaults:
35    /// - `channels` → `in_channels` (same as the last array's head projection)
36    /// - `out_channels` → 1 (mono output)
37    /// - `kernel_size` → 3
38    /// - `bias` → false
39    /// - `activation` → "Tanh"
40    ///
41    /// Weight and bias arrays are zero-initialized and must be populated
42    /// by the dispatcher via `set_weights` and `set_bias`.
43    pub fn from_config(config: &HeadConfig, in_channels: usize) -> std::io::Result<Self> {
44        let channels = config.channels.unwrap_or(in_channels);
45        let out_channels = config.out_channels.unwrap_or(1);
46        let kernel = config.kernel_size.unwrap_or(3);
47        let do_bias = config.bias.unwrap_or(false);
48        let activation = parse_activation(config.activation.as_deref().unwrap_or("Tanh"));
49
50        let num_blocks = out_channels.div_ceil(4);
51        let weights_len = num_blocks * kernel * channels * 4;
52        let bias_len = out_channels;
53
54        let weights = AlignedVec::new(weights_len, 0.0f32)
55            .map_err(|e| std::io::Error::new(std::io::ErrorKind::OutOfMemory, format!("{e}")))?;
56        let bias = AlignedVec::new(bias_len, 0.0f32)
57            .map_err(|e| std::io::Error::new(std::io::ErrorKind::OutOfMemory, format!("{e}")))?;
58
59        let receptive_field = kernel;
60        let state = WaveNetLayerState::new(channels, receptive_field, 0)?;
61
62        let conv = Conv1dDyn {
63            weights,
64            bias,
65            do_bias,
66            dilation: 1,
67            in_ch: channels,
68            out_ch: out_channels,
69            num_blocks,
70            interleave_width: 4,
71            kernel,
72        };
73
74        let scratch = AlignedVec::new(out_channels * WAVENET_MAX_NUM_FRAMES, 0.0f32)
75            .map_err(|e| std::io::Error::new(std::io::ErrorKind::OutOfMemory, format!("{e}")))?;
76
77        Ok(Self {
78            conv,
79            activation,
80            state,
81            scratch,
82        })
83    }
84
85    /// Returns the receptive field contribution of this head (kernel size).
86    /// Must be added to the global model receptive field for prewarm.
87    pub fn receptive_field(&self) -> usize {
88        self.conv.kernel
89    }
90
91    /// Number of output channels produced by this head.
92    pub fn out_channels(&self) -> usize {
93        self.conv.out_ch
94    }
95
96    /// Number of input channels expected by this head.
97    pub fn in_channels(&self) -> usize {
98        self.conv.in_ch
99    }
100
101    /// Loads convolution weights from a flat f32 slice.
102    pub fn set_weights(&mut self, weights: &[f32]) {
103        let len = self.conv.weights.len().min(weights.len());
104        self.conv.weights[..len].copy_from_slice(&weights[..len]);
105    }
106
107    /// Loads convolution bias from a flat f32 slice, if present.
108    pub fn set_bias(&mut self, bias: &[f32]) {
109        let len = self.conv.bias.len().min(bias.len());
110        self.conv.bias[..len].copy_from_slice(&bias[..len]);
111    }
112
113    /// Public dispatch wrapper that selects the optimal SIMD path.
114    ///
115    /// # Safety
116    /// Input and output slices must have sizes compatible with the head dimensions:
117    /// `input.len() == num_frames * in_ch`, `output.len() == num_frames * out_ch`.
118    /// The ring buffer state must have been properly initialized (via `prewarm` or
119    /// sufficient prior processing) to cover the causal receptive field.
120    #[inline(always)]
121    pub unsafe fn process_block(&mut self, input: &[f32], output: &mut [f32], num_frames: usize) {
122        unsafe {
123            crate::math::common::dispatch_simd!(
124                self,
125                process_block_internal,
126                input,
127                output,
128                num_frames
129            )
130        };
131    }
132
133    /// SIMD-dispatched processing kernel.
134    ///
135    /// Writes `num_frames` of input into the ring buffer, runs the causal
136    /// Conv1D, applies activation, and writes results to output.
137    ///
138    /// Input layout: frame-interleaved `[f0_c0, f0_c1, ..., f1_c0, ...]`.
139    /// Output layout: frame-interleaved `[f0_c0, f0_c1, ..., f1_c0, ...]`.
140    ///
141    /// # Safety
142    /// `input` and `output` must have sizes `num_frames * in_ch` and
143    /// `num_frames * out_ch` respectively. The ring buffer must have been
144    /// properly initialized.
145    #[inline(always)]
146    pub unsafe fn process_block_internal<M: SimdMath>(
147        &mut self,
148        input: &[f32],
149        output: &mut [f32],
150        num_frames: usize,
151    ) {
152        let in_ch = self.conv.in_ch;
153        let out_ch = self.conv.out_ch;
154        let input_len = num_frames * in_ch;
155
156        let buf_start = self.state.buffer_start * in_ch;
157        unsafe {
158            core::ptr::copy_nonoverlapping(
159                input.as_ptr(),
160                self.state.layer_buffer.as_mut_ptr().add(buf_start),
161                input_len,
162            );
163        }
164
165        let scratch_slice = &mut self.scratch[..num_frames * out_ch];
166        unsafe {
167            self.conv.process_block::<M>(
168                &self.state.layer_buffer,
169                scratch_slice,
170                self.state.buffer_start,
171                num_frames,
172                None,
173            );
174        }
175
176        unsafe {
177            self.activation.apply_simd::<M>(scratch_slice);
178        }
179
180        unsafe {
181            core::ptr::copy_nonoverlapping(
182                scratch_slice.as_ptr(),
183                output.as_mut_ptr(),
184                num_frames * out_ch,
185            );
186        }
187
188        self.state.advance_frames(num_frames, in_ch);
189    }
190
191    /// Public prewarm wrapper with SIMD dispatch.
192    #[cold]
193    pub fn prewarm(&mut self) {
194        unsafe {
195            crate::math::common::dispatch_simd!(self, prewarm_internal);
196        }
197    }
198
199    /// Fills the conv state buffer with a single frame of silence replicated
200    /// backward to cover the entire receptive field.
201    ///
202    /// # Safety
203    /// Must be called via `dispatch_simd!` macro. The state buffer must be
204    /// properly allocated and the ring buffer start pointer must be valid.
205    #[inline(always)]
206    pub unsafe fn prewarm_internal<M: SimdMath>(&mut self) {
207        let in_ch = self.conv.in_ch;
208        let out_ch = self.conv.out_ch;
209        let kernel = self.conv.kernel;
210
211        let buf_start = self.state.buffer_start * in_ch;
212
213        self.state.layer_buffer[buf_start..buf_start + in_ch].fill(0.0);
214
215        let start_idx = self.state.buffer_start * in_ch;
216        let src_range = start_idx..start_idx + in_ch;
217        for offset in 1..=kernel {
218            let dst_idx = (self.state.buffer_start - offset) * in_ch;
219            self.state
220                .layer_buffer
221                .copy_within(src_range.clone(), dst_idx);
222        }
223
224        let scratch_slice = &mut self.scratch[..out_ch];
225        unsafe {
226            self.conv.process_single_frame::<M>(
227                &self.state.layer_buffer,
228                scratch_slice,
229                self.state.buffer_start,
230                None,
231            );
232        }
233        unsafe {
234            self.activation.apply_simd::<M>(scratch_slice);
235        }
236
237        self.state.advance_frames(1, in_ch);
238    }
239}
240
241/// Maps an activation function name string to an `ActivationType`.
242///
243/// Supported values match the variant names of `ActivationType`:
244/// `"Tanh"`, `"HardTanh"`, `"FastTanh"`, `"ReLU"`, `"Sigmoid"`,
245/// `"SiLU"`, `"HardSwish"`, `"Softsign"`.
246///
247/// Unrecognized strings fall back to `ActivationType::Tanh`.
248pub fn parse_activation(name: &str) -> ActivationType {
249    match name {
250        "Tanh" => ActivationType::Tanh,
251        "HardTanh" => ActivationType::HardTanh,
252        "FastTanh" => ActivationType::FastTanh,
253        "ReLU" => ActivationType::ReLU,
254        "Sigmoid" => ActivationType::Sigmoid,
255        "SiLU" => ActivationType::SiLU,
256        "HardSwish" => ActivationType::HardSwish,
257        "Softsign" => ActivationType::Softsign,
258        _ => ActivationType::Tanh,
259    }
260}
261
262#[cfg(test)]
263#[path = "post_stack_head_test.rs"]
264mod tests;
nam_rs/models/wavenet/post_stack_head.rs

nam_rs/models/wavenet/
post_stack_head.rs
