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
// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.
use super::common::{WAVENET_MAX_NUM_FRAMES, WavenetProcessContext};
use super::conv1d_dyn::Conv1dDyn;
use super::dense_dyn::DenseLayerDyn;
use crate::common::diagnostics::NamErrorCode;
use crate::math::common::{AlignedVec, SimdMath};
/// Complete Convolutional Cell (WaveNet Layer) with runtime dimensions.
#[derive(Clone)]
pub struct WaveNetLayerDyn {
/// This layer's parametric dilated Causal 1D convolution mesh.
pub conv1d: Conv1dDyn,
/// Conditional injection mixing network.
pub input_mixin: DenseLayerDyn,
/// 1x1 decompression linear affine transform of the layer.
pub one_by_one: DenseLayerDyn,
/// Pre-allocated scratch buffer for conditioning mixin output
/// (size: `ch * WAVENET_MAX_NUM_FRAMES`).
pub scratch_mixin: AlignedVec<f32>,
/// Pre-allocated scratch buffer for Conv1D + mixin intermediate results
/// (size: `ch * WAVENET_MAX_NUM_FRAMES`).
pub scratch_conv: AlignedVec<f32>,
}
impl WaveNetLayerDyn {
/// Creates a new dynamic WaveNet layer.
///
/// Scratch buffers are pre-allocated with size `ch * WAVENET_MAX_NUM_FRAMES`,
/// ensuring zero allocation on the hot-path and supporting CH > 16
/// (vs the const-generic stack buffer `[f32; 1024]` which is limited to CH ≤ 16).
pub fn new(
ch: usize,
conv1d: Conv1dDyn,
input_mixin: DenseLayerDyn,
one_by_one: DenseLayerDyn,
) -> Result<Self, NamErrorCode> {
let scratch_size = ch * WAVENET_MAX_NUM_FRAMES;
Ok(Self {
conv1d,
input_mixin,
one_by_one,
scratch_mixin: AlignedVec::new(scratch_size, 0.0f32)?,
scratch_conv: AlignedVec::new(scratch_size, 0.0f32)?,
})
}
/// Processes a full WaveNet layer with runtime-dimensional scratch buffers.
///
/// Uses dual-frame Conv1D tiling (process_block) for SIMD amortization.
/// Scratch buffers are pre-allocated `AlignedVec<f32>` on the heap,
/// replacing the const-generic stack buffer `[f32; 1024]` to support CH > 16.
///
/// Uses full-precision f32 weights (no quantization) and exact `f32::tanh`.
///
/// # Safety
/// Math dispatch via pointer to inlined intrinsic functions.
#[inline(always)]
pub unsafe fn process_block_internal<M: SimdMath>(&mut self, ctx: WavenetProcessContext<'_>) {
let ch = self.conv1d.in_ch;
let scratch_len = ctx.num_frames * ch;
debug_assert!(
scratch_len <= self.scratch_mixin.len(),
"process_block_internal: num_frames*ch ({}) exceeds scratch buffer ({})",
scratch_len,
self.scratch_mixin.len()
);
unsafe {
let mixin_out = &mut self.scratch_mixin[..scratch_len];
self.input_mixin
.process_block::<M>(ctx.condition, mixin_out, ctx.num_frames);
let conv_slice = &mut self.scratch_conv[..scratch_len];
self.conv1d.process_block::<M>(
ctx.layer_buffer,
conv_slice,
ctx.buffer_start,
ctx.num_frames,
Some(mixin_out),
);
if let Some(s) = ctx.seed {
M::tanh_and_accumulate_with_seed(ctx.head_input, conv_slice, s);
} else if ctx.is_first_layer {
M::tanh_and_overwrite_block(ctx.head_input, conv_slice);
} else {
M::tanh_and_accumulate_block(ctx.head_input, conv_slice);
}
let lb_offset = ctx.buffer_start * ch;
let residual_slice = ctx
.layer_buffer
.get_unchecked(lb_offset..lb_offset + scratch_len);
self.one_by_one.process_residual_batch::<M>(
conv_slice,
residual_slice,
ctx.output,
ctx.num_frames,
);
}
}
}