use crate::common::diagnostics::NamErrorCode;
use crate::math::common::AlignedVec;
pub const fn ch_pad<const CH: usize>() -> usize {
CH.next_power_of_two()
}
#[derive(Clone)]
#[repr(align(64))]
pub struct A2Conv1dCh<const CH: usize> {
pub weights: AlignedVec<f32>,
pub bias: AlignedVec<f32>,
pub dilation: usize,
pub kernel: usize,
}
impl<const CH: usize> A2Conv1dCh<CH> {
pub fn new(
raw_weights: &[f32],
out_ch: usize,
in_ch: usize,
kernel: usize,
dilation: usize,
raw_bias: &[f32],
) -> Result<Self, NamErrorCode> {
let ch_pad = ch_pad::<CH>();
debug_assert_eq!(out_ch, CH);
debug_assert_eq!(in_ch, CH);
debug_assert!(kernel == 6 || kernel == 15);
debug_assert_eq!(raw_weights.len(), out_ch * in_ch * kernel);
debug_assert_eq!(raw_bias.len(), CH);
let stride = ch_pad * ch_pad;
let mut weights = AlignedVec::new(kernel * stride, 0.0f32)?;
for out in 0..out_ch {
for inp in 0..in_ch {
for k in 0..kernel {
let src = out * in_ch * kernel + inp * kernel + k;
let dst = k * stride + inp * ch_pad + out;
weights[dst] = raw_weights[src];
}
}
}
let mut bias = AlignedVec::new(ch_pad, 0.0f32)?;
bias[..CH].copy_from_slice(&raw_bias[..CH]);
Ok(Self {
weights,
bias,
dilation,
kernel,
})
}
}