nam_rs/models/wavenet/model.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::WAVENET_MAX_NUM_FRAMES;
5use super::layer_array::WaveNetLayerArray;
6use crate::math::common::SimdMath;
7
8/// Complete WaveNet Model containing Two Heterogeneous Layer Array Blocks.
9///
10/// **Scientific Reference:** van den Oord, A., et al. (2016). *"WaveNet: A Generative Model for Raw Audio."* DeepMind.
11///
12/// `CH` = Array1 channels (layer 0 of the JSON, e.g., 16 for Standard)
13/// `K` = kernel size (always 3)
14/// `HEAD` = Array1 head_size = Array2 channels (e.g., 8 for Standard)
15///
16/// Array2 uses `HEAD` channels and projects to 1 output (`HEAD2=1`),
17/// following the C++ pattern: `WaveNetLayerArrayT<CH, 1, 1, HEAD, K, Dilations, true>`.
18pub struct WaveNetModel<const CH: usize, const K: usize, const HEAD: usize> {
19 /// Inner array 01: IN=1, COND=1, CH channels, HEAD outputs, no HeadBias.
20 pub array1: WaveNetLayerArray<1, 1, CH, K, HEAD>,
21 /// Inner array 02: IN=CH, COND=1, HEAD channels, 1 output, with HeadBias.
22 pub array2: WaveNetLayerArray<CH, 1, HEAD, K, 1>,
23 /// Final voltage compensation scale (Target Output Scale).
24 pub head_scale: f32,
25 /// Largest circular buffer required at the Kernel's temporal root.
26 pub receptive_field_size: usize,
27 /// Whether to execute prewarm during `reset()`. Default: `true`.
28 pub prewarm_on_reset: bool,
29}
30
31impl<const CH: usize, const K: usize, const HEAD: usize> WaveNetModel<CH, K, HEAD> {
32 /// Sets the effective number of layers on both arrays for soft-degrade.
33 #[inline(always)]
34 pub fn set_effective_layers(&mut self, n: usize) {
35 self.array1.set_effective_layers(n);
36 self.array2.set_effective_layers(n);
37 }
38
39 /// Backs up the `buffer_start` pointers of both layer arrays into a slice.
40 #[inline(always)]
41 pub fn backup_buffer_starts(&self, starts: &mut [usize], offset: &mut usize) {
42 for state in &self.array1.states {
43 if *offset < starts.len() {
44 starts[*offset] = state.buffer_start;
45 *offset += 1;
46 }
47 }
48 for state in &self.array2.states {
49 if *offset < starts.len() {
50 starts[*offset] = state.buffer_start;
51 *offset += 1;
52 }
53 }
54 }
55
56 /// Restores the `buffer_start` pointers of both layer arrays from a slice.
57 #[inline(always)]
58 pub fn restore_buffer_starts(&mut self, starts: &[usize], offset: &mut usize) {
59 for state in &mut self.array1.states {
60 if *offset < starts.len() {
61 state.buffer_start = starts[*offset];
62 *offset += 1;
63 }
64 }
65 for state in &mut self.array2.states {
66 if *offset < starts.len() {
67 state.buffer_start = starts[*offset];
68 *offset += 1;
69 }
70 }
71 }
72
73 /// Resolves the full forward pass and produces waveform samples in zero allocation (DSP).
74 ///
75 /// Combines the outputs of both arrays: `sum(head1) + sum(head2)` × `head_scale`.
76 ///
77 /// **For Scientists and Devs:** The `dispatch_simd!` macro monomorphizes
78 /// this function via the `SimdMath` trait, matching the global hardware
79 /// configuration once at the call site. The compiler generates branchless
80 /// assembly optimized for your processor (AVX2, AVX-512, or AVX-512 VNNI BF16)
81 /// without v-table or dynamic dispatch overhead.
82 pub fn process(&mut self, input: &[f32], output: &mut [f32]) {
83 unsafe { crate::math::common::dispatch_simd!(self, process_internal, input, output) };
84 }
85
86 #[inline(always)]
87 /// Fast, generic routine that implements the neural network (WaveNet).
88 /// The `<M: SimdMath>` constraint forces the compiler to generate assembly focused on
89 /// large registers (256-bit or 512-bit) without branches (branchless).
90 unsafe fn process_internal<M: SimdMath>(&mut self, input: &[f32], output: &mut [f32]) {
91 let total_frames = input.len();
92 if total_frames == 0 {
93 return;
94 }
95
96 let mut pos = 0;
97 // [PROCESSING IN CHUNKS (BLOCKS)]
98 // To maintain zero-allocation invariants (no temporary RAM vector allocations)
99 // and respect the restricted L1/L2 Cache hierarchy, we limit processing
100 // to `WAVENET_MAX_NUM_FRAMES` (typically 64 samples) at a time.
101 // This loop iterates until it consumes the entire buffer (e.g., 256, 512, 1024 frames).
102 while pos < total_frames {
103 let num_frames = (total_frames - pos).min(WAVENET_MAX_NUM_FRAMES);
104 let in_slice = &input[pos..pos + num_frames];
105
106 unsafe {
107 // [STEP 1: Array1 Forward]
108 // Conditioning and Input (1D: 1 channel) -> formatted as blocks of IN frames.
109 // In the standard NAM topology, this Array performs convolutions using huge dilations
110 // (e.g., from 1 to 512, 1 to 512 successively) to capture amplifier sub-bass.
111 // Its output enters `array1.array_outputs` and the skips enter `array1.head_outputs`.
112 self.array1
113 .process_block_internal::<M, false>(in_slice, in_slice, num_frames, None);
114
115 // [STEP 2: Array2 Forward (Cascaded Head)]
116 // C++ parity: array2 seeds its head_accum with array1's post-head_rechannel
117 // output — all layers then accumulate on top of this seed. The final output
118 // is head_scale × array2.head_outputs (only the last array's head).
119 let array1_head_out = &self.array1.head_outputs[0..num_frames * HEAD];
120 let array1_outputs = &self.array1.array_outputs[0..num_frames * CH];
121 self.array2.process_block_internal::<M, false>(
122 array1_outputs,
123 in_slice,
124 num_frames,
125 Some(array1_head_out),
126 );
127 }
128
129 // [STEP 3: Final Scale]
130 // C++ reference: output = head_scale × last_array.head_outputs
131 let array2_head = &self.array2.head_outputs[0..num_frames];
132 let out_slice = &mut output[pos..pos + num_frames];
133 unsafe {
134 core::ptr::copy_nonoverlapping(
135 array2_head.as_ptr(),
136 out_slice.as_mut_ptr(),
137 num_frames,
138 );
139 }
140 unsafe {
141 M::apply_gain(out_slice, self.head_scale);
142 }
143 pos += num_frames;
144 }
145 }
146
147 /// Stabilizes the model by processing silence (Zero Input) for pre-warm.
148 ///
149 /// Dispatches via `dispatch_simd!` macro — a single static `match` on the
150 /// globally detected `SIMD_MATH.instruction_set` — to run the AVX2, AVX-512,
151 /// or AVX-512 VNNI BF16 kernel without runtime feature detection per call.
152 #[cold]
153 pub fn prewarm(&mut self) {
154 unsafe {
155 crate::math::common::dispatch_simd!(self, prewarm_internal);
156 }
157 }
158
159 /// Prewarm strictly optimized for AVX-512 architecture.
160 ///
161 /// # Safety
162 /// Requires a supported processor (AVX-512).
163 #[target_feature(enable = "avx512f,avx512vl")]
164 #[cold]
165 pub unsafe fn prewarm_avx512(&mut self) {
166 unsafe { self.prewarm_internal::<crate::math::common::Avx512Math>() };
167 }
168
169 /// Prewarm strictly optimized for AVX2 architecture.
170 ///
171 /// # Safety
172 /// Requires an x86-64-v3 (AVX2) processor.
173 #[cold]
174 pub unsafe fn prewarm_avx2(&mut self) {
175 unsafe { self.prewarm_internal::<crate::math::common::Avx2Math>() };
176 }
177
178 /// # Safety
179 /// Call this via `dispatch_simd!` macro only.
180 #[inline(always)]
181 #[cold]
182 unsafe fn prewarm_internal<M: SimdMath>(&mut self) {
183 let condition = [0.0f32];
184 let layer_inputs_1 = [0.0f32];
185
186 unsafe {
187 self.array1
188 .prewarm_internal::<M>(&layer_inputs_1, &condition, None);
189 }
190 let array1_outputs = &self.array1.array_outputs[0..CH];
191 let array1_head_out = &self.array1.head_outputs[0..HEAD];
192 unsafe {
193 self.array2
194 .prewarm_internal::<M>(array1_outputs, &condition, Some(array1_head_out));
195 }
196 }
197}