active-call 0.3.84

// Auto-generated by lele_compiler
use lele::kernels::*;
use lele::tensor::TensorView;

#[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>,
}
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,
        )
    }

    /// ARM-optimized quantized linear + ReLU with pre-packed weights.
    /// Uses fused DynQuant+GEMM: eliminates f32 intermediate buffer, per-call u8
    /// allocation, and separate f32→u8 conversion pass.
    #[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,
        )
    }

    /// ARM-optimized MatMulInteger with pre-packed weight cache.
    /// Used for unfused MatMulInteger nodes where B is a static model weight.
    /// Eliminates: per-call B packing, B u8→f32→u8 roundtrip, heap alloc.
    #[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,
        )
    }

    // Helper for pre-quantized inputs (used in attention where input is already quantized)
    #[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);

        // FUSED: MatMul + Scale + Bias in one operation
        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> {
        // ConstantOfShape + Concat pattern
        // shape defines the shape of the constant tensor filled with `value`
        // Then concatenate weight and the constant along axis 0
        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> {
        // ConstantOfShape + Concat pattern (i64)
        // shape defines the shape of the constant tensor filled with `value`
        // Then concatenate weight and the constant along axis 0
        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_1 = lele::kernels::reshape(&waveform, &[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_37 = lele::kernels::stft(
            &view_2,
            512,
            160,
            512,
            Some(&self.weight_f32(676480, 2048, &[512])),
            &mut ws.buf_2,
        );
        let stft = lele::kernels::transpose(&val_37, &[0, 2, 1, 3], &mut ws.buf_3);
        let abs_1 = lele::kernels::reduce_l2(&stft, &[-1], false, &mut ws.buf_4);
        let pow_1 = lele::kernels::pow(&abs_1, &self.weight_f32(700432, 4, &[]), &mut ws.buf_5);
        let transpose = lele::kernels::transpose(&pow_1, &[0, 2, 1], &mut ws.buf_0);
        let matmul = lele::kernels::matmul(
            &transpose,
            &self.weight_f32(610624, 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(678528, 4, &[])),
            None,
            &mut ws.buf_2,
        );
        let val_48 = lele::kernels::log(&clamp, &mut ws.buf_3);
        let log10 = lele::kernels::div(&val_48, &self.weight_f32(678544, 4, &[]), &mut ws.buf_4);
        let mul = lele::kernels::mul(&log10, &self.weight_f32(700448, 4, &[]), &mut ws.buf_5);
        let view_5 = lele::kernels::reshape(&mul, &[1, 1, 64, 601]);
        let amax = lele::kernels::reduce_max(&view_5, &[-3, -2, -1], false, &mut ws.buf_0);
        let sub_1 = lele::kernels::sub(&amax, &self.weight_f32(700496, 4, &[]), &mut ws.buf_1);
        let view_6 = lele::kernels::reshape(&sub_1, &[-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, 64, 601]);
        let mean = lele::kernels::reduce_mean(&view_7, &[-2, -1], true, &mut ws.buf_3);
        let val_76 = lele::kernels::reduce_mean(&view_7, &[1, 2], true, &mut ws.buf_4);
        let val_77 = lele::kernels::sub(&view_7, &val_76, &mut ws.buf_0);
        let val_78 = lele::kernels::mul(&val_77, &val_77, &mut ws.buf_6);
        let val_79 = lele::kernels::reduce_mean(&val_78, &[1, 2], false, &mut ws.buf_1);
        let val_84 = lele::kernels::mul(&val_79, &self.weight_f32(678672, 4, &[]), &mut ws.buf_5);
        let var = lele::kernels::div(&val_84, &self.weight_f32(678688, 4, &[]), &mut ws.buf_7);
        let val_94 = lele::kernels::reshape(&var, &[1, 1, 1]);
        let sqrt = lele::kernels::sqrt(&val_94, &mut ws.buf_4);
        let clamp_1 = lele::kernels::clip(
            &sqrt,
            Some(&self.weight_f32(678736, 4, &[])),
            None,
            &mut ws.buf_0,
        );
        let sub_2 = lele::kernels::sub(&view_7, &mean, &mut ws.buf_6);
        let div = lele::kernels::div(&sub_2, &clamp_1, &mut ws.buf_1);
        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(698944, 128, &[32])),
            &[1, 1],
            1,
            &[1, 1, 1, 1],
            &[1, 1],
            true,
            &mut ws.buf_7,
        );
        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(8000, 73728, &[64, 32, 3, 3]),
            Some(&self.weight_f32(699072, 256, &[64])),
            &[1, 1],
            1,
            &[1, 1, 1, 1],
            &[1, 1],
            true,
            &mut ws.buf_0,
        );
        let max_pool2d_1 = lele::kernels::max_pool2d(
            &relu_1,
            &[2, 2],
            &[2, 2],
            &[0, 0, 0, 0],
            &[1, 1],
            false,
            &mut ws.buf_3,
        );
        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_1,
        );
        let relu_2 = lele::kernels::conv2d_fused(
            &conv2d_2,
            &self.weight_f32(81728, 32768, &[128, 64, 1, 1]),
            Some(&self.weight_f32(699328, 512, &[128])),
            &[1, 1],
            1,
            &[0, 0, 0, 0],
            &[1, 1],
            true,
            &mut ws.buf_6,
        );
        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(114496, 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(119104, 65536, &[128, 128, 1, 1]),
            Some(&self.weight_f32(699840, 512, &[128])),
            &[1, 1],
            1,
            &[0, 0, 0, 0],
            &[1, 1],
            true,
            &mut ws.buf_3,
        );
        let max_pool2d_3 = lele::kernels::max_pool2d(
            &relu_3,
            &[8, 1],
            &[8, 1],
            &[0, 0, 0, 0],
            &[1, 1],
            false,
            &mut ws.buf_1,
        );
        let permute = lele::kernels::transpose(&max_pool2d_3, &[0, 3, 1, 2], &mut ws.buf_5);
        let view_8 = lele::kernels::reshape(&permute, &[1, 75, 128]);
        let val_168 = lele::kernels::transpose(&view_8, &[1, 0, 2], &mut ws.buf_4);
        let mut buf_val_206_h = Vec::<f32>::new();
        let mut buf_val_206 = Vec::<f32>::new();
        let (val_206, _) = lele::kernels::gru(
            &val_168,
            &self.weight_f32(1093792, 196608, &[1, 384, 128]),
            &self.weight_f32(1290400, 196608, &[1, 384, 128]),
            Some(&self.weight_f32(679328, 3072, &[1, 768])),
            Some(&self.weight_f32(678784, 512, &[1, 1, 128])),
            false,
            &mut buf_val_206,
            &mut buf_val_206_h,
        );
        let val_207 = lele::kernels::transpose(&val_206, &[0, 2, 1, 3], &mut ws.buf_0);
        let val_220 = lele::kernels::reshape(&val_207, &[75, 1, 128]);
        let getitem_12 = lele::kernels::transpose(&val_220, &[1, 0, 2], &mut ws.buf_1);
        let linear = self.linear(
            &getitem_12,
            &self.weight_f32(682432, 16384, &[128, 32]),
            &self.weight_f32(3456, 128, &[32]),
            &mut ws.buf_5,
        );
        let tanh = lele::kernels::tanh_kernel(&linear, &mut ws.buf_4);
        let linear_1 = self.linear(
            &tanh,
            &self.weight_f32(698816, 128, &[32, 1]),
            &self.weight_f32(3584, 4, &[1]),
            &mut ws.buf_6,
        );
        let softmax = lele::kernels::softmax(&linear_1, 1, &mut ws.buf_0);
        let mul_1 = lele::kernels::mul(&getitem_12, &softmax, &mut ws.buf_7);
        let sum_1 = lele::kernels::reduce_sum(&mul_1, &[1], false, &mut ws.buf_3);
        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_5,
        );
        let linear_2 = lele::kernels::gemm(
            &layer_norm,
            &self.weight_f32(577856, 32768, &[64, 128]),
            Some(&self.weight_f32(4624, 256, &[64])),
            1.0,
            1.0,
            false,
            true,
            &mut ws.buf_4,
        );
        let relu_4 = lele::kernels::relu(&linear_2, &mut ws.buf_2);
        let linear_3 = lele::kernels::gemm(
            &relu_4,
            &self.weight_f32(4880, 3072, &[12, 64]),
            Some(&self.weight_f32(7952, 48, &[12])),
            1.0,
            1.0,
            false,
            true,
            &mut ws.buf_6,
        );
        let probabilities = lele::kernels::softmax(&linear_3, -1, &mut ws.buf_0);
        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]
    }

    // Inference Entry Point
    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
    }
}