#include "cl_viv_vx_ext.h"

_viv_uniform VXC_512Bits uniU8SubZp_lo_2x8;
_viv_uniform VXC_512Bits uniU8SubZp_hi_2x8;
_viv_uniform VXC_512Bits uniU8Conv1d_part0_8x2;
_viv_uniform VXC_512Bits uniU8Conv1d_part1_8x2;
_viv_uniform VXC_512Bits uniU8Conv1d_part2_8x2;
_viv_uniform VXC_512Bits uniU8Conv1d_part3_8x2;
_viv_uniform VXC_512Bits uniSumOrderUchar_2x8;

_viv_uniform int kernel_cnt_x16;
_viv_uniform int weight_ZP;
_viv_uniform float output_ZP;
_viv_uniform float scaleOut;
_viv_uniform int input_height;
_viv_uniform int input_width;
_viv_uniform int output_width;

__kernel void conv1d_U8U8I32toU8_K1024_SMALL(
     __read_only image2d_array_t   input,
     __read_only image2d_array_t   weight,
     __read_only image2d_t         bias,
    __write_only image2d_array_t   output,
                 int               stride,
                 int               pad_front,
                 int               pad_end,
                 int               dilation,
                 int               overflow_policy)
{
    int  start_x = get_global_id(0) - pad_front;
    int4 coord   = (int4)(start_x, get_global_id(1), 0, get_global_id(0));
    int4 coord_w = (int4)(0, 0, get_global_id(1), 0);
    float4 sum0, sum1, dst;
    vxc_short8 coef;
    vxc_short8 w_zp = (short)weight_ZP;
    vxc_uchar16 input_val = 0, weight_val = 0;
    int temp = 0, i, j;

    temp = read_imagei(bias, coord.yz).x;
    sum0 = convert_float(temp);
    sum1 = sum0;

    for (i = 0; i < input_height; i++)
    {
        for (j = 0; j < kernel_cnt_x16; j++)
        {
            VXC_ReadImage2DArray(weight_val, weight, coord_w, VXC_5BITOFFSET_XY(0, 0), \
                                 VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            VXC_ReadImage(input_val, input, coord.xz, VXC_5BITOFFSET_XY(0, 0), \
                                 VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            VXC_DP2x8(coef, weight_val, w_zp, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniU8SubZp_lo_2x8);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part0_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part1_8x2);
            sum0 += dst;
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part2_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part3_8x2);
            sum1 += dst;
            VXC_ReadImage(input_val, input, coord.xz, VXC_5BITOFFSET_XY(8, 0), \
                                 VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            VXC_DP2x8(coef, weight_val, w_zp, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniU8SubZp_hi_2x8);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part0_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part1_8x2);
            sum0 += dst;
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part2_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part3_8x2);
            sum1 += dst;
            coord_w.x += 16;
            coord.x += 16;
        }
        coord_w.x = 0;
        coord_w.y++;
        coord.z++;
        coord.x = start_x;
    }

    sum0 = sum0 * scaleOut + output_ZP;
    sum1 = sum1 * scaleOut + output_ZP;
    uchar4 result0, result1;
    _viv_asm(CONV_SAT_RTE, result0, sum0);
    _viv_asm(CONV_SAT_RTE, result1, sum1);
    vxc_uchar8 result;
    VXC_DP2x8(result, result0, result1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniSumOrderUchar_2x8);
    VXC_WriteImage(output, coord.wy, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
}

__kernel void conv1d_U8U8I32toU8_K1024_LARGE(
     __read_only image2d_array_t   input,
     __read_only image2d_array_t   weight,
     __read_only image2d_t         bias,
    __write_only image2d_array_t   output,
                 int               stride,
                 int               pad_front,
                 int               pad_end,
                 int               dilation,
                 int               overflow_policy)
{
    int  start_x = get_global_id(0);
    int  w_left  = output_width - start_x;
    int  out_x   = w_left < 8 ? get_global_id(0) - w_left : get_global_id(0);
    int4 coord   = (int4)(start_x, get_global_id(1), 0, out_x);
    int4 coord_w = (int4)(0, 0, get_global_id(1), 0);
    float4 sum0, sum1, dst;
    vxc_short8 coef;
    vxc_short8 w_zp = (short)weight_ZP;
    vxc_uchar16 input_val = 0, weight_val = 0;
    int temp = 0, i, j;
    Tensor src_tensor = create_tensor_from_image2d_array(input, 1);
    uchar *src_ptr_base = (uchar *)src_tensor.ptr;
    uchar *src_ptr;
    Tensor dst_tensor = create_tensor_from_image2d_array(output, 1);
    uchar *dst_ptr = (uchar *)dst_tensor.ptr;

    temp = read_imagei(bias, coord.yz).x;
    sum0 = convert_float(temp);
    sum1 = sum0;

    for (i = 0; i < input_height; i++)
    {
        src_ptr = src_ptr_base + (coord.x + coord.z * src_tensor.stride_y);
        for (j = 0; j < kernel_cnt_x16; j++)
        {
            VXC_ReadImage2DArray(weight_val, weight, coord_w, VXC_5BITOFFSET_XY(0, 0), \
                                 VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
            VXC_Vload16(input_val, src_ptr, 0);
            VXC_DP2x8(coef, weight_val, w_zp, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniU8SubZp_lo_2x8);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part0_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part1_8x2);
            sum0 += dst;
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part2_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part3_8x2);
            sum1 += dst;
            src_ptr += 8;
            VXC_Vload16(input_val, src_ptr, 0);
            VXC_DP2x8(coef, weight_val, w_zp, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniU8SubZp_hi_2x8);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part0_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part1_8x2);
            sum0 += dst;
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(0, 1, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part2_8x2);
            VXC_DP8x2(dst, input_val, coef, VXC_MODIFIER(2, 3, 0, VXC_RM_TowardZero, 0), uniU8Conv1d_part3_8x2);
            sum1 += dst;
            coord_w.x += 16;
            coord.x += 16;
            src_ptr += 8;
        }
        coord_w.x = 0;
        coord_w.y++;
        coord.z++;
        coord.x = start_x;
    }

    sum0 = sum0 * scaleOut + output_ZP;
    sum1 = sum1 * scaleOut + output_ZP;
    uchar4 result0, result1;
    _viv_asm(CONV_SAT_RTE, result0, sum0);
    _viv_asm(CONV_SAT_RTE, result1, sum1);
    vxc_uchar8 result;
    VXC_DP2x8(result, result0, result1, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniSumOrderUchar_2x8);
    dst_ptr = dst_ptr + (coord.w + coord.y * dst_tensor.stride_y);
    VXC_Vstore8(dst_ptr, 0, result);
}

