nam_rs/models/a2/model/static/
prewarm.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
4//! Pre-warm implementation for the static A2 model.
5//!
6//! Mirrors `A2FastModel::prewarm()` in `a2_fast.cpp`: zeroes all buffers,
7//! then feeds `receptive_field_size` frames of zero input through `process()`
8//! so that layer biases populate the head accumulator ring — matching the
9//! C++ steady-state initial condition.
10//!
11//! ## Root cause of initial A2 Rust×C++ divergence
12//!
13//! The initial Rust port diverged from C++ because `prewarm()` only zeroed buffers
14//! while the C++ `A2FastModel::prewarm()` feeds `_prewarm_samples` frames of silence
15//! through `process()`. Even with zero input, A2 layers produce non-zero activations
16//! (conv bias + mixin × 0 + LeakyReLU), populating the head accumulator ring.
17//! The Rust zero-fill vs C++ silent-process mismatch meant zero initial state in Rust
18//! vs bias-driven steady state in C++ — causing frame-level divergence from the first
19//! output sample. Fixed by making Rust `prewarm()` run `process()` with zeros for
20//! `receptive_field_size` frames after the zero-fill, mirroring the C++ DSP::Reset
21//! → prewarm() flow exactly.
22
23use crate::models::wavenet::common::WAVENET_MAX_NUM_FRAMES;
24
25use super::super::super::params::A2_NUM_LAYERS;
26use super::super::a2_prewarm_common;
27use super::WaveNetA2;
28
29impl<const CH: usize> WaveNetA2<CH> {
30    /// Pre-warms the model by filling the receptive field with silence.
31    #[cold]
32    pub fn prewarm(&mut self) {
33        a2_prewarm_common(
34            A2_NUM_LAYERS,
35            self.receptive_field_size,
36            &mut self.layer_buffers,
37            &self.layer_ring_sizes,
38            &mut self.layer_buffer_starts,
39            &mut self.layer_in,
40            &mut self.head_accum,
41            &mut self.head_write_pos,
42        );
43
44        if self.has_weights() {
45            let prewarm_samples = self.receptive_field_size;
46            let block = WAVENET_MAX_NUM_FRAMES;
47            let zeros = [0.0f32; WAVENET_MAX_NUM_FRAMES];
48            let mut discard = [0.0f32; WAVENET_MAX_NUM_FRAMES];
49            let mut remaining = prewarm_samples;
50            while remaining > 0 {
51                let nf = remaining.min(block);
52                self.process(&zeros[..nf], &mut discard[..nf]);
53                remaining -= nf;
54            }
55        }
56    }
57}
nam_rs/models/a2/model/static/prewarm.rs

nam_rs/models/a2/model/static/
prewarm.rs
