use lele::tensor::TensorView;
use lele::kernels::*;
#[derive(Default)]
pub struct TelcoClassifierWorkspace {
pub buf_0: Vec<f32>,
pub buf_1: Vec<f32>,
pub buf_2: Vec<f32>,
pub buf_3: Vec<f32>,
pub buf_4: Vec<f32>,
pub buf_5: Vec<f32>,
pub buf_6: Vec<f32>,
pub buf_7: Vec<f32>,
pub buf_8: Vec<f32>,
}
impl TelcoClassifierWorkspace {
pub fn new() -> Self {
Self::default()
}
}
pub struct TelcoClassifier<'a> {
data: &'a [u8],
_phantom: std::marker::PhantomData<&'a ()>,
#[cfg(target_arch = "aarch64")]
prepared_weights_cache: std::cell::RefCell<std::collections::HashMap<(usize, usize), std::sync::Arc<lele::kernels::PreparedWeightsArm>>>,
}
impl<'a> TelcoClassifier<'a> {
pub fn new(data: &'a [u8]) -> Self {
Self {
data,
_phantom: std::marker::PhantomData,
#[cfg(target_arch = "aarch64")]
prepared_weights_cache: std::cell::RefCell::new(std::collections::HashMap::new()),
}
}
fn conv1d_relu<'c, 'd>(
&self,
input: lele::tensor::TensorView<'c>,
weight: lele::tensor::TensorView<'c>,
bias: Option<&lele::tensor::TensorView<'c>>,
stride: usize,
dilation: usize,
groups: usize,
padding: usize,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d> {
lele::kernels::conv1d_fused(
&input,
&weight,
bias,
&[dilation as i64],
groups as i64,
&[padding as i64, padding as i64],
&[stride as i64],
true,
output_buf,
)
}
fn layer_norm<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c>,
scale: lele::tensor::TensorView<'c>,
bias: lele::tensor::TensorView<'c>,
epsilon: lele::tensor::TensorView<'c>,
_two: lele::tensor::TensorView<'c>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d> {
let eps = epsilon.data.first().cloned().unwrap_or(1e-5);
lele::kernels::layer_norm(input, &scale, &bias, -1, eps, output_buf)
}
fn linear_quantized<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c, f32>,
weight_int8: lele::tensor::TensorView<'c, f32>,
weight_scale: lele::tensor::TensorView<'c, f32>,
weight_zero: lele::tensor::TensorView<'c, f32>,
bias: lele::tensor::TensorView<'c, f32>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
lele::kernels::fused_quantized_linear(
input, &weight_int8, &weight_scale, &weight_zero, &bias, false, output_buf,
)
}
fn linear_quantized_relu<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c, f32>,
weight_int8: lele::tensor::TensorView<'c, f32>,
weight_scale: lele::tensor::TensorView<'c, f32>,
weight_zero: lele::tensor::TensorView<'c, f32>,
bias: lele::tensor::TensorView<'c, f32>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
lele::kernels::fused_quantized_linear(
input, &weight_int8, &weight_scale, &weight_zero, &bias, true, output_buf,
)
}
#[cfg(target_arch = "aarch64")]
fn linear_quantized_arm<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c, f32>,
weight_offset: usize,
weight_len: usize,
weight_k: usize,
weight_n: usize,
weight_scale: lele::tensor::TensorView<'c, f32>,
weight_zero: lele::tensor::TensorView<'c, f32>,
bias: lele::tensor::TensorView<'c, f32>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
let pw = self.get_prepared_weight(weight_offset, weight_len, weight_k, weight_n);
let zp_b = weight_zero.data.first().map(|&v| v as u8);
lele::kernels::fused_dq_gemm_prepared_arm(
input,
&pw,
zp_b,
&weight_scale,
Some(&bias),
false,
output_buf,
)
}
#[cfg(target_arch = "aarch64")]
fn linear_quantized_relu_arm<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c, f32>,
weight_offset: usize,
weight_len: usize,
weight_k: usize,
weight_n: usize,
weight_scale: lele::tensor::TensorView<'c, f32>,
weight_zero: lele::tensor::TensorView<'c, f32>,
bias: lele::tensor::TensorView<'c, f32>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
let pw = self.get_prepared_weight(weight_offset, weight_len, weight_k, weight_n);
let zp_b = weight_zero.data.first().map(|&v| v as u8);
lele::kernels::fused_dq_gemm_prepared_arm(
input,
&pw,
zp_b,
&weight_scale,
Some(&bias),
true,
output_buf,
)
}
#[cfg(target_arch = "aarch64")]
fn mat_mul_integer_arm<'c, 'd>(
&self,
a: &lele::tensor::TensorView<'c, f32>,
weight_offset: usize,
weight_len: usize,
weight_k: usize,
weight_n: usize,
a_zero_point: Option<&lele::tensor::TensorView<'c, f32>>,
b_zero_point: Option<&lele::tensor::TensorView<'c, f32>>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
let pw = self.get_prepared_weight(weight_offset, weight_len, weight_k, weight_n);
let zp_a = a_zero_point.and_then(|z| z.data.first().cloned());
let zp_b = b_zero_point.and_then(|z| z.data.first()).map(|&v| v as u8);
lele::kernels::mat_mul_integer_prepared_arm(a, &pw, zp_a, zp_b, None, None, false, output_buf)
}
#[inline]
fn linear_quantized_prequant<'c, 'd>(
&self,
input_quantized: &lele::tensor::TensorView<'c, f32>,
input_scale: &lele::tensor::TensorView<'c, f32>,
input_zero_point: &lele::tensor::TensorView<'c, f32>,
weight_int8: lele::tensor::TensorView<'c, f32>,
weight_scale: lele::tensor::TensorView<'c, f32>,
weight_zero: lele::tensor::TensorView<'c, f32>,
bias: lele::tensor::TensorView<'c, f32>,
output_buf: &'d mut Vec<f32>,
scale_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d, f32> {
let combined_scale = lele::kernels::mul(input_scale, &weight_scale, scale_buf);
lele::kernels::mat_mul_integer_with_scale_bias(
input_quantized,
&weight_int8,
Some(input_zero_point),
Some(&weight_zero),
Some(&combined_scale),
Some(&bias),
output_buf,
)
}
fn linear<'c, 'd>(
&self,
input: &lele::tensor::TensorView<'c>,
weight: &lele::tensor::TensorView<'c>,
bias: &lele::tensor::TensorView<'c>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d> {
lele::kernels::matmul_fused_add(input, weight, bias, output_buf)
}
fn embedding_concat<'c, 'd>(
&self,
shape: &lele::tensor::TensorView<'c, i64>,
value: f32,
weight: lele::tensor::TensorView<'c>,
output_buf: &'d mut Vec<f32>,
) -> lele::tensor::TensorView<'d> {
let const_shape: Vec<usize> = shape.data.iter().map(|&x| x as usize).collect();
let const_len: usize = const_shape.iter().product();
output_buf.clear();
output_buf.reserve(weight.data.len() + const_len);
output_buf.extend_from_slice(&weight.data);
output_buf.resize(weight.data.len() + const_len, value);
let mut out_shape = weight.shape.to_vec();
out_shape[0] += const_shape[0];
lele::tensor::TensorView {
data: std::borrow::Cow::Borrowed(output_buf),
shape: std::borrow::Cow::Owned(out_shape),
}
}
fn embedding_concat_i64<'c, 'd>(
&self,
shape: &lele::tensor::TensorView<'c, i64>,
value: i64,
weight: lele::tensor::TensorView<'c, i64>,
output_buf: &'d mut Vec<i64>,
) -> lele::tensor::TensorView<'d, i64> {
let const_shape: Vec<usize> = shape.data.iter().map(|&x| x as usize).collect();
let const_len: usize = const_shape.iter().product();
output_buf.clear();
output_buf.reserve(weight.data.len() + const_len);
output_buf.extend_from_slice(&weight.data);
output_buf.resize(weight.data.len() + const_len, value);
let mut out_shape = weight.shape.to_vec();
out_shape[0] += const_shape[0];
lele::tensor::TensorView {
data: std::borrow::Cow::Borrowed(output_buf),
shape: std::borrow::Cow::Owned(out_shape),
}
}
#[inline(never)]
fn run_chunk_0<'w>(&self, ws: &'w mut TelcoClassifierWorkspace, waveform: TensorView<'w, f32>) -> TensorView<'static, f32> {
let view = lele::kernels::reshape(&waveform, &[-1, 96000]);
let view_1 = lele::kernels::reshape(&view, &[1, -1, 96000]);
let pad = lele::kernels::pad(&view_1, &[0, 0, 256, 0, 0, 256], None, "reflect", &mut ws.buf_1);
let view_2 = lele::kernels::reshape(&pad, &[-1, 96512]);
let val_38 = lele::kernels::stft_power_spectrum(&view_2, 512, 160, 512, Some(&self.weight_f32(676192, 2048, &[512])), &mut ws.buf_2);
let matmul = lele::kernels::matmul(&val_38, &self.weight_f32(610368, 65792, &[257, 64]), &mut ws.buf_6);
let transpose_1 = lele::kernels::transpose(&matmul, &[0, 2, 1], &mut ws.buf_1);
let clamp = lele::kernels::clip(&transpose_1, Some(&self.weight_f32(678240, 4, &[])), None, &mut ws.buf_2);
let val_49 = lele::kernels::log(&clamp, &mut ws.buf_3);
let log10 = lele::kernels::div(&val_49, &self.weight_f32(678256, 4, &[]), &mut ws.buf_4);
let mul_39 = lele::kernels::mul(&log10, &self.weight_f32(699568, 4, &[]), &mut ws.buf_5);
let view_5 = lele::kernels::reshape(&mul_39, &[1, -1, 601]);
let amax = lele::kernels::reduce_max(&view_5, &[-3, -2, -1], false, &mut ws.buf_0);
let sub_18 = lele::kernels::sub(&amax, &self.weight_f32(699648, 4, &[]), &mut ws.buf_1);
let view_6 = lele::kernels::reshape(&sub_18, &[-1, 1, 1, 1]);
let maximum = lele::kernels::max(&view_5, &view_6, &mut ws.buf_2);
let view_7 = lele::kernels::reshape(&maximum, &[1, -1, 601]);
let mean = lele::kernels::reduce_mean(&view_7, &[-2, -1], true, &mut ws.buf_3);
let val_77 = lele::kernels::reduce_mean(&view_7, &[1, 2], true, &mut ws.buf_4);
let val_78 = lele::kernels::sub(&view_7, &val_77, &mut ws.buf_0);
let val_79 = lele::kernels::mul(&val_78, &val_78, &mut ws.buf_6);
let val_80 = lele::kernels::reduce_mean(&val_79, &[1, 2], false, &mut ws.buf_1);
let val_85 = lele::kernels::mul(&val_80, &self.weight_f32(678320, 4, &[]), &mut ws.buf_5);
let var = lele::kernels::div(&val_85, &self.weight_f32(678336, 4, &[]), &mut ws.buf_7);
let mut buf_val_89 = Vec::<i64>::new();
let val_89 = lele::kernels::shape(&var);
let mut buf_val_91 = Vec::<i64>::new();
let val_91 = lele::kernels::reshape(&val_89, &[-1]);
let mut buf_val_94 = Vec::<i64>::new();
let val_94 = lele::kernels::concat(&[&val_91, &self.weight_i64(699424, 8, &[1]), &self.weight_i64(699424, 8, &[1])], 0, &mut buf_val_94);
let val_95 = lele::kernels::reshape(&var, &val_94.data[..]);
let sqrt = lele::kernels::sqrt(&val_95, &mut ws.buf_1);
let clamp_1 = lele::kernels::clip(&sqrt, Some(&self.weight_f32(678352, 4, &[])), None, &mut ws.buf_4);
let sub_24 = lele::kernels::sub(&view_7, &mean, &mut ws.buf_5);
let div = lele::kernels::div(&sub_24, &clamp_1, &mut ws.buf_6);
let unsqueeze = lele::kernels::unsqueeze(&div, &[1]);
let relu = lele::kernels::conv2d_fused(&unsqueeze, &self.weight_f32(0, 1152, &[32, 1, 3, 3]), Some(&self.weight_f32(698016, 128, &[32])), &[1, 1], 1, &[1, 1, 1, 1], &[1, 1], true, &mut ws.buf_1);
let max_pool2d = lele::kernels::max_pool2d(&relu, &[2, 2], &[2, 2], &[0, 0, 0, 0], &[1, 1], false, &mut ws.buf_4);
let relu_1 = lele::kernels::conv2d_fused(&max_pool2d, &self.weight_f32(7744, 73728, &[64, 32, 3, 3]), Some(&self.weight_f32(698144, 256, &[64])), &[1, 1], 1, &[1, 1, 1, 1], &[1, 1], true, &mut ws.buf_3);
let max_pool2d_1 = lele::kernels::max_pool2d(&relu_1, &[2, 2], &[2, 2], &[0, 0, 0, 0], &[1, 1], false, &mut ws.buf_5);
let conv2d_2 = lele::kernels::conv2d(&max_pool2d_1, &self.weight_f32(1152, 2304, &[64, 1, 3, 3]), None, &[1, 1], 64, &[1, 1, 1, 1], &[1, 1], &mut ws.buf_6);
let relu_2 = lele::kernels::conv2d_fused(&conv2d_2, &self.weight_f32(81472, 32768, &[128, 64, 1, 1]), Some(&self.weight_f32(698400, 512, &[128])), &[1, 1], 1, &[0, 0, 0, 0], &[1, 1], true, &mut ws.buf_1);
let max_pool2d_2 = lele::kernels::max_pool2d(&relu_2, &[2, 2], &[2, 2], &[0, 0, 0, 0], &[1, 1], false, &mut ws.buf_4);
let conv2d_4 = lele::kernels::conv2d(&max_pool2d_2, &self.weight_f32(114240, 4608, &[128, 1, 3, 3]), None, &[1, 1], 128, &[1, 1, 1, 1], &[1, 1], &mut ws.buf_2);
let relu_3 = lele::kernels::conv2d_fused(&conv2d_4, &self.weight_f32(118848, 65536, &[128, 128, 1, 1]), Some(&self.weight_f32(698912, 512, &[128])), &[1, 1], 1, &[0, 0, 0, 0], &[1, 1], true, &mut ws.buf_5);
let max_pool2d_3 = lele::kernels::max_pool2d(&relu_3, &[8, 1], &[8, 1], &[0, 0, 0, 0], &[1, 1], false, &mut ws.buf_6);
let permute = lele::kernels::transpose(&max_pool2d_3, &[0, 3, 1, 2], &mut ws.buf_0);
let view_8 = lele::kernels::reshape(&permute, &[1, -1, 128]);
let val_171 = lele::kernels::transpose(&view_8, &[1, 0, 2], &mut ws.buf_1);
let mut buf_val_209_h = Vec::<f32>::new();
let mut buf_val_209 = Vec::<f32>::new();
let (val_209, _) = lele::kernels::gru(&val_171, &self.weight_f32(1093328, 196608, &[1, 384, 128]), &self.weight_f32(1289936, 196608, &[1, 384, 128]), Some(&self.weight_f32(678400, 3072, &[1, 768])), None, false, &mut buf_val_209, &mut buf_val_209_h);
let val_210 = lele::kernels::transpose(&val_209, &[0, 2, 1, 3], &mut ws.buf_5);
let val_223 = lele::kernels::reshape(&val_210, &[75, 1, 128]);
let getitem_12 = lele::kernels::transpose(&val_223, &[1, 0, 2], &mut ws.buf_6);
let linear = self.linear(&getitem_12, &self.weight_f32(681504, 16384, &[128, 32]), &self.weight_f32(3456, 128, &[32]), &mut ws.buf_1);
let tanh = lele::kernels::tanh_kernel(&linear, &mut ws.buf_2);
let linear_1 = self.linear(&tanh, &self.weight_f32(697888, 128, &[32, 1]), &self.weight_f32(3584, 4, &[1]), &mut ws.buf_4);
let softmax = lele::kernels::softmax(&linear_1, 1, &mut ws.buf_5);
let mul_92 = lele::kernels::mul(&getitem_12, &softmax, &mut ws.buf_7);
let sum_1 = lele::kernels::reduce_sum(&mul_92, &[1], false, &mut ws.buf_0);
let layer_norm = lele::kernels::layer_norm(&sum_1, &self.weight_f32(3600, 512, &[128]), &self.weight_f32(4112, 512, &[128]), -1, 0.00001, &mut ws.buf_1);
let linear_2 = lele::kernels::gemm(&layer_norm, &self.weight_f32(577600, 32768, &[64, 128]), Some(&self.weight_f32(4624, 256, &[64])), 1.0, 1.0, false, true, &mut ws.buf_2);
let relu_4 = lele::kernels::relu(&linear_2, &mut ws.buf_3);
let linear_3 = lele::kernels::gemm(&relu_4, &self.weight_f32(4880, 2816, &[11, 64]), Some(&self.weight_f32(7696, 44, &[11])), 1.0, 1.0, false, true, &mut ws.buf_4);
let probabilities = lele::kernels::softmax(&linear_3, -1, &mut ws.buf_5);
probabilities.to_owned()
}
#[cfg(target_arch = "aarch64")]
fn get_prepared_weight(&self, offset: usize, len: usize, k: usize, n: usize) -> std::sync::Arc<lele::kernels::PreparedWeightsArm> {
let key = (offset, len);
{
let cache = self.prepared_weights_cache.borrow();
if let Some(pw) = cache.get(&key) {
return pw.clone();
}
}
let raw_bytes = &self.data[offset..offset+len];
let pw = std::sync::Arc::new(lele::kernels::prepare_weights_arm(raw_bytes, k, n));
self.prepared_weights_cache.borrow_mut().insert(key, pw.clone());
pw
}
fn weight_f32(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'a, f32> {
TensorView::from_bytes_f32(&self.data[offset..offset+len], shape)
}
fn weight_i64(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, i64> {
TensorView::from_bytes_i64(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_i32_i64(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, i64> {
TensorView::from_bytes_i32_as_i64(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_i32(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, i32> {
TensorView::from_bytes_i32(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_i64_f32(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, f32> {
TensorView::from_bytes_i64_as_f32(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_i32_f32(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, f32> {
TensorView::from_bytes_i32_as_f32(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_u8(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, f32> {
TensorView::from_bytes_u8(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_i8(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, f32> {
TensorView::from_bytes_i8(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_f16(&self, offset: usize, len: usize, shape: &'a [usize]) -> TensorView<'static, f32> {
TensorView::from_bytes_f16(&self.data[offset..offset+len], shape.to_vec())
}
fn weight_u8_raw(&self, offset: usize, len: usize) -> &'a [u8] {
&self.data[offset..offset+len]
}
pub fn forward(&self, waveform: TensorView<'a>) -> TensorView<'static> {
let mut ws = TelcoClassifierWorkspace::new();
let res = self.forward_with_workspace(&mut ws, waveform);
res.to_owned()
}
pub fn forward_with_workspace<'w>(&self, ws: &'w mut TelcoClassifierWorkspace, waveform: TensorView<'w>) -> TensorView<'w> {
let (probabilities) = self.run_chunk_0(ws, waveform);
probabilities
}
}