/**
 * Copyright 2020-2021 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef NNACL_INT8_MATMUL_H_
#define NNACL_INT8_MATMUL_H_

#include <string.h>
#include "nnacl_c/op_base.h"
#include "nnacl_c/matmul_parameter.h"

#ifdef __cplusplus
extern "C" {
#endif
/* 4x16 16x4 -> 4x4 */
/* sdot 4x4 4x16 -> 4x16 */
/* matmul */
void MatMulInt8_16x4(const int8_t *a, const int8_t *b, int32_t *dst, int row_4, int col_4, int deep_16,
                     const int32_t *input_sum, const int32_t *bias);
void RowMajor2Row16x4MajorInt8(const int8_t *src_ptr, int8_t *dst_ptr, int row, int col);
void RowMajor2Col16x4MajorInt8(const int8_t *src, int8_t *dst, int row, int col);
void RowMajor2Col4x16MajorInt8(const int8_t *src, int8_t *dst, int row, int col);
void RowMajor2Col4x16MajorPartInt8(const int8_t *src, int8_t *dst, int row, int col, int cur_oc);
void PackInput2Col4x4AndInputSumPert(const int8_t *src_input, int8_t *packed_input, int32_t *input_sum, int row,
                                     int col, int row_stride, int32_t filter_zp);
void CalcInputSums(const int8_t *input, int row, int col, int weight_zp, int32_t *dst, DataOrder order);
void CalcWeightBiasSums(const int8_t *weight, int row, int col, int input_zp, const int32_t *weight_zp_ptr,
                        const int32_t *bias, int32_t *dst, DataOrder order, bool filter_per_channel);
void CalcPartWeightBiasSums(const int8_t *weight, int row, int stride, int cur_col, int input_zp,
                            const int32_t *weight_zp_ptr, const int32_t *bias, int32_t *dst, DataOrder order,
                            bool filter_per_channel);
void MatmulInt8Opt(const int8_t *a, const int8_t *b, int8_t *dst, int row, int col, int deep16, const int32_t *a_sums,
                   const int32_t *bias, int act_min, int act_max, int out_zp, const int32_t *multiplier,
                   const int32_t *left_shift, const int32_t *right_shift, size_t stride, size_t filter_peroc,
                   const int32_t *filter_zp);
/* 8x4 4x8 -> 8x8 */
/* optimize conv */
void RowMajor2Row8x4MajorInt8(const int8_t *src_ptr, int8_t *dst_ptr, int row, int col);
void MatMulInt8_8x8_r(const int8_t *a, const int8_t *b, int8_t *dst, size_t row, size_t col, size_t deep_4,
                      size_t stride, const int32_t *input_sum, const int32_t *bias, const int32_t *left_shift,
                      const int32_t *right_shift, const int32_t *multiplier, int32_t output_zp, int32_t mini,
                      int32_t maxi, size_t per_channel);

/* 4x16 16x2 -> 4x2 */
/* arm32 conv1x1 */
void RowMajor2Row2x16MajorInt8(const int8_t *src_ptr, int8_t *dst_ptr, int row, int col);
void RowMajor2Col16x2MajorInt8(const int8_t *src_ptr, int8_t *dst_ptr, int row, int col);
void MatMulInt8_4x2_r(const int8_t *a, const int8_t *b, int8_t *dst, size_t row, size_t col, size_t deep_16,
                      size_t stride, const int32_t *input_sum, const int32_t *bias, const int32_t *left_shift,
                      const int32_t *right_shift, const int32_t *multiplier, int32_t output_zp, int32_t mini,
                      int32_t maxi, bool peroc);

/* 4x4 4x16 -> 4x16 */
/* optimize conv1x1 */
void RowMajor2Row4x16MajorInt8(const int8_t *src_ptr, int8_t *dst_ptr, int row, int col);
void PackInput4x4AndInputSumPert(const int8_t *src_input, int8_t *packed_input, int32_t *input_sum,
                                 size_t input_channel, size_t plane_size, int32_t filter_zp);
void MatMulInt8_4x16_r(const int8_t *a, const int8_t *b, int8_t *dst, size_t row, size_t col, size_t deep_4,
                       size_t stride, const int32_t *input_sum, const int32_t *bias, const int32_t *left_shift,
                       const int32_t *right_shift, const int32_t *multiplier, int32_t output_zp, int32_t mini,
                       int32_t maxi, size_t per_channel, const int32_t *filter_zp);
void MatmulInt8LowMemory(const int8_t *a, const int8_t *b, int8_t *dst, int row, int col, int deep16,
                         const int32_t *a_sums, const int32_t *bias, int mini, int maxi, int out_zp,
                         const int32_t *multiplier, const int32_t *left_shift, const int32_t *right_shift,
                         size_t stride, size_t filter_peroc, const int32_t *filter_zp, bool transB, bool transA);

#ifdef ENABLE_ARM64
void MatmulInt8Neon64(const int8_t *a, const int8_t *b, int8_t *dst, int row4, int col4, int deep16,
                      const int32_t *a_sums, const int32_t *bias, int act_min, int act_max, int out_zp,
                      int32_t *multiplier, int32_t *left_shift, int32_t *right_shift, int row, int col, int stride,
                      int filter_peroc);

void MatMulR4Int8Neon64(const int8_t *a, const int8_t *b, int32_t *dst, int row4, int col4, int deep16,
                        const int32_t *input_sum, const int32_t *bias);
#endif
#ifdef ENABLE_ARM32
void MatmulInt8Neon32(const int8_t *a, const int8_t *b, int8_t *dst, int row, int col, int deep16,
                      const int32_t *input_sums, const int32_t *weight_bias, int act_min, int act_max, int out_zp,
                      int32_t *multiplier, int32_t *left_shift, int32_t *right_shift, int stride, int per_channel);
#endif
#ifdef __cplusplus
}
#endif

#endif  // LITE_SRC_BACKEND_ARM_NNACL_INT8_MATMUL_H_
