nam_rs/models/wavenet/layer_array.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 super::common::{WaveNetLayerState, WavenetProcessContext};
5use super::dense::DenseLayer;
6use super::layer::WaveNetLayer;
7use crate::math::common::{AlignedVec, SimdMath};
8use core::arch::x86_64::{_MM_HINT_T0, _mm_prefetch};
9
10/// Grouped Multi-Layer WaveNet Unit.
11pub struct WaveNetLayerArray<
12 const IN: usize,
13 const COND: usize,
14 const CH: usize,
15 const K: usize,
16 const HEAD: usize,
17> {
18 /// Vec with structural topology (length defines blocks).
19 pub layers: Vec<WaveNetLayer<COND, CH, K>>,
20 /// RingBuffer states, one for each Layer in the system.
21 pub states: Vec<WaveNetLayerState>,
22 /// Initial `Dense` tensor opening.
23 pub rechannel: DenseLayer<IN, CH>,
24 /// Final tensor closure generating Head projection.
25 pub head_rechannel: DenseLayer<CH, HEAD>,
26
27 /// Pre-allocated temporary output array.
28 /// Array output temporary accumulator.
29 pub array_outputs: AlignedVec<f32>,
30 /// CH-sized intermediate accumulator for layer contributions before the Head projection.
31 pub head_accum: AlignedVec<f32>,
32 /// Allocated global Linear projection memory (HEAD-sized).
33 pub head_outputs: AlignedVec<f32>,
34 /// Dimensional field size (global receptive field) for routing.
35 pub receptive_field_size: usize,
36 /// Shared activation buffer size.
37 pub block_size: usize,
38 /// Temporary accumulator for blocks (pre-allocated).
39 pub block_buffer: AlignedVec<f32>,
40
41 /// Active number of layers for soft-degrade. Set to `layers.len()` by default.
42 pub effective_layers: usize,
43}
44
45impl<const IN: usize, const COND: usize, const CH: usize, const K: usize, const HEAD: usize>
46 WaveNetLayerArray<IN, COND, CH, K, HEAD>
47{
48 /// Sets the effective number of layers for soft-degrade.
49 /// Clamped to [1, layers.len()].
50 #[inline(always)]
51 pub fn set_effective_layers(&mut self, n: usize) {
52 self.effective_layers = n.min(self.layers.len()).max(1);
53 }
54
55 /// Array's central processing. Fully shielded against allocations.
56 ///
57 /// When `prev_head_outputs` is `Some`, it is passed as `seed` in the
58 /// first layer's `WavenetProcessContext` — the layer then fuses seed
59 /// accumulation with tanh in a single SIMD pass. This eliminates the
60 /// dedicated `copy_from_slice` of `num_frames * CH` bytes.
61 ///
62 /// This implements the C++ cascaded head pattern (array N seeds its head
63 /// accumulator with the post-head_rechannel output of array N−1).
64 ///
65 /// # Safety
66 /// State pointers iterate internally without bounds checks.
67 #[inline]
68 pub unsafe fn process_block_internal<M: SimdMath, const PREWARM: bool>(
69 &mut self,
70 layer_inputs: &[f32],
71 condition: &[f32],
72 num_frames: usize,
73 prev_head_outputs: Option<&[f32]>,
74 ) {
75 debug_assert_eq!(self.layers.len(), self.states.len());
76 let states_ptr = self.states.as_mut_ptr();
77
78 unsafe {
79 let state_0 = &mut *states_ptr.add(0);
80 let start = state_0.buffer_start * CH;
81 self.rechannel.process_block::<M>(
82 layer_inputs,
83 &mut state_0.layer_buffer[start..start + num_frames * CH],
84 num_frames,
85 );
86
87 let num_layers = self.effective_layers;
88 let last_layer = num_layers - 1;
89
90 // [STEP 4: Layer Inference Cascade]
91 for (i, layer) in self.layers.iter_mut().enumerate() {
92 if i >= num_layers {
93 break;
94 }
95 let current_state = &mut *states_ptr.add(i);
96
97 if PREWARM {
98 // [STEP 4.1: Static State Propagation (Backfill)]
99 // In prewarm mode, we replicate the current sample to the entire history (Receptive Field).
100 // This ensures the network stabilizes instantly to the stationary value.
101 let start_idx = current_state.buffer_start * CH;
102 let src_range = start_idx..start_idx + CH;
103
104 for offset in 1..=current_state.receptive_field_size {
105 debug_assert!(
106 current_state.buffer_start >= offset,
107 "backfill underflow: bs={}, off={}",
108 current_state.buffer_start,
109 offset
110 );
111 // SAFETY: garantido pelo construtor WaveNetLayerState::new que valida buffer_start >= receptive_field_size
112 let dst_start = current_state.buffer_start - offset;
113 let dst_idx = dst_start * CH;
114 current_state
115 .layer_buffer
116 .copy_within(src_range.clone(), dst_idx);
117 }
118 }
119
120 // Software Prefetch the next state in the cascade.
121 // Bring the cache line of state i+1 (and i+2 if possible) into L1
122 // while the processor resolves the arithmetic pipeline of the current layer.
123 if i + 1 < num_layers {
124 _mm_prefetch::<_MM_HINT_T0>(states_ptr.add(i + 1) as *const i8);
125 }
126 if i + 2 < num_layers {
127 _mm_prefetch::<_MM_HINT_T0>(states_ptr.add(i + 2) as *const i8);
128 }
129
130 if i == last_layer {
131 layer.process_block_internal::<M>(WavenetProcessContext {
132 condition,
133 head_input: &mut self.head_accum[0..num_frames * CH],
134 output: &mut self.array_outputs[0..num_frames * CH],
135 layer_buffer: ¤t_state.layer_buffer[..],
136 buffer_start: current_state.buffer_start,
137 num_frames,
138 block: &mut self.block_buffer[0..num_frames * self.block_size],
139 is_first_layer: i == 0,
140 seed: if i == 0 { prev_head_outputs } else { None },
141 });
142 } else {
143 let next_state = &mut *states_ptr.add(i + 1);
144 let n_start = next_state.buffer_start * CH;
145 let next_layer_buffer =
146 &mut next_state.layer_buffer[n_start..n_start + num_frames * CH];
147
148 layer.process_block_internal::<M>(WavenetProcessContext {
149 condition,
150 head_input: &mut self.head_accum[0..num_frames * CH],
151 output: next_layer_buffer,
152 layer_buffer: ¤t_state.layer_buffer[..],
153 buffer_start: current_state.buffer_start,
154 num_frames,
155 block: &mut self.block_buffer[0..num_frames * self.block_size],
156 is_first_layer: i == 0,
157 seed: if i == 0 { prev_head_outputs } else { None },
158 });
159 }
160
161 if !PREWARM {
162 current_state.advance_frames(num_frames, CH);
163 }
164 }
165
166 // [STEP 5: Dimensional Closure (Head Rechannel)]
167 // The dense matrix funnels the accumulator (sum of skip-connections from all layers,
168 // of size `CH`) into a smaller `HEAD` dimension (e.g., 16 -> 8 or 16 -> 1).
169 // Uses f32 native precision for this critical final projection (tonal fidelity),
170 // while the backbone runs quantized (BF16/F16) for performance.
171 self.head_rechannel.process_block::<M>(
172 &self.head_accum[0..num_frames * CH],
173 &mut self.head_outputs[0..num_frames * HEAD],
174 num_frames,
175 );
176 }
177 }
178
179 /// Processes data in Pre-warm mode to initialize and stabilize temporal memory.
180 ///
181 /// [SCIENTIFIC EXPLANATION]
182 /// Causal convolution neural networks like WaveNet have an internal state that actively
183 /// depends on N past steps (Receptive Field). When loading a fresh model, the allocated
184 /// network memory (Ring Buffers) contains "pure zeros" or computational garbage.
185 /// Pre-warm feeds a continuous inert signal (Absolute Silence) into the network so as to
186 /// fill the entire past window. The resulting cold-start transients "drain"
187 /// silently into oblivion, ensuring the first audio sample when turning on the device
188 /// sounds organic and stable, without pops or clicks.
189 /// # Safety
190 /// Call this via `dispatch_simd!` macro only.
191 #[inline(always)]
192 pub unsafe fn prewarm_internal<M: SimdMath>(
193 &mut self,
194 layer_inputs: &[f32],
195 condition: &[f32],
196 prev_head_outputs: Option<&[f32]>,
197 ) {
198 unsafe {
199 // Unified via shared code: we process 1 frame with the prewarm flag active.
200 // [STEP 4.1] inside `process_block_internal` handles the backfill.
201 self.process_block_internal::<M, true>(layer_inputs, condition, 1, prev_head_outputs);
202 }
203 }
204}