text: /**
   * \file dnn/src/arm_common/matrix_mul/int8/gemv.cpp
   * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
   *
   * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   */
  
  #include "src/arm_common/simd_macro/marm_neon.h"
  #include "src/arm_common/matrix_mul/int8/gemv.h"
  #include "src/common/utils.h"
  #include "megdnn/oprs.h"
  
  #include "midout.h"
  MIDOUT_DECL(megdnn_arm_common_int8_gemv)
  
  using namespace megdnn;
  using namespace arm_common;
  
  
  namespace {
  
  void gemv_naive_n(const int8_t* __restrict A, const int8_t* __restrict B,
                    int32_t* __restrict C, size_t M, size_t N, size_t K,
                    size_t Astride, size_t Bstride, size_t Cstride) {
      megdnn_assert(N == 1 && Bstride == 1);
      size_t m = 0;
      for (; m + 2 <= M; m += 2) {
          int32_t acc0 = 0, acc1 = 0;
          size_t k = 0;
          for (; k + 16 <= K; k += 16) {
              int8x16_t a0 = vld1q_s8(A + m * Astride + k);
              int8x16_t a1 = vld1q_s8(A + (m + 1) * Astride + k);
  
              int8x16_t b0 = vld1q_s8(B + k);
  
              int16x8_t c0 = vmull_s8(vget_low_s8(a0), vget_low_s8(b0));
              c0 = vmlal_high_s8(c0, a0, b0);
  
              int16x8_t c1 = vmull_s8(vget_low_s8(a1), vget_low_s8(b0));
              c1 = vmlal_high_s8(c1, a1, b0);
              acc0 += vaddlvq_s16(c0);
              acc1 += vaddlvq_s16(c1);
          }
  
          for (; k + 8 <= K; k += 8) {
              int8x8_t a0 = vld1_s8(A + m * Astride + k);
              int8x8_t a1 = vld1_s8(A + (m + 1) * Astride + k);
              int8x8_t b0 = vld1_s8(B + k);
  
              int16x8_t c0 = vmull_s8(a0, b0);
  
              int16x8_t c1 = vmull_s8(a1, b0);
              acc0 += vaddlvq_s16(c0);
              acc1 += vaddlvq_s16(c1);
          }
  
          for (; k < K; ++k) {
              acc0 += static_cast<int32_t>(A[m * Astride + k]) * B[k];
              acc1 += static_cast<int32_t>(A[(m + 1) * Astride + k]) * B[k];
          }
          C[m * Cstride] = acc0;
          C[(m + 1) * Cstride] = acc1;
      }
  
      for (; m < M; ++m) {
          int32_t acc0 = 0;
          size_t k = 0;
          for (; k + 16 <= K; k += 16) {
              int8x16_t a0 = vld1q_s8(A + m * Astride + k);
              int8x16_t b0 = vld1q_s8(B + k);
  
              int16x8_t c0 = vmull_s8(vget_low_s8(a0), vget_low_s8(b0));
              c0 = vmlal_high_s8(c0, a0, b0);
  
              acc0 += vaddlvq_s16(c0);
          }
  
          for (; k + 8 <= K; k += 8) {
              int8x8_t a0 = vld1_s8(A + m * Astride + k);
              int8x8_t b0 = vld1_s8(B + k);
  
              int16x8_t c0 = vmull_s8(a0, b0);
              acc0 += vaddlvq_s16(c0);
          }
  
          for (; k < K; ++k) {
              acc0 += static_cast<int32_t>(A[m * Astride + k]) * B[k];
          }
          C[m * Cstride] = acc0;
      }
  }
  
  void gemv_naive_n_mk4(const int8_t* __restrict A, const int8_t* __restrict B,
                        int32_t* __restrict C, size_t M, size_t N, size_t K,
                        size_t Astride, size_t Bstride, size_t Cstride) {
      constexpr size_t PACK_SIZE = 4;
      megdnn_assert(N == 1 && Bstride == 4);
      auto Aptr = A;
      size_t m = 0;
      for (; m < M; m += PACK_SIZE) {
          auto Bptr = B;
          auto Aptr0 = Aptr;
          int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0;
          size_t k = 0;
          for (; k + 16 <= K; k += 16) {
              int8x16x4_t a = vld4q_s8(Aptr0);
              int8x16_t b = vld1q_s8(Bptr);
              int16x8_t c[4];
  
              c[0] = vmull_s8(vget_low_s8(a.val[0]), vget_low_s8(b));
              c[1] = vmull_s8(vget_low_s8(a.val[1]), vget_low_s8(b));
              c[2] = vmull_s8(vget_low_s8(a.val[2]), vget_low_s8(b));
              c[3] = vmull_s8(vget_low_s8(a.val[3]), vget_low_s8(b));
  
              c[0] = vmlal_high_s8(c[0], a.val[0], b);
              c[1] = vmlal_high_s8(c[1], a.val[1], b);
              c[2] = vmlal_high_s8(c[2], a.val[2], b);
              c[3] = vmlal_high_s8(c[3], a.val[3], b);
  
              acc0 += vaddlvq_s16(c[0]);
              acc1 += vaddlvq_s16(c[1]);
              acc2 += vaddlvq_s16(c[2]);
              acc3 += vaddlvq_s16(c[3]);
  
              Bptr += 16;
              Aptr0 += PACK_SIZE * 16;
          }
  
          for (; k + 8 <= K; k += 8) {
              int8x8x4_t a = vld4_s8(Aptr0);
              int8x8_t b = vld1_s8(Bptr);
              int16x8_t c[4];
  
              c[0] = vmull_s8(a.val[0], b);
              c[1] = vmull_s8(a.val[1], b);
              c[2] = vmull_s8(a.val[2], b);
              c[3] = vmull_s8(a.val[3], b);
  
              acc0 += vaddlvq_s16(c[0]);
              acc1 += vaddlvq_s16(c[1]);
              acc2 += vaddlvq_s16(c[2]);
              acc3 += vaddlvq_s16(c[3]);
  
              Bptr += 8;
              Aptr0 += PACK_SIZE * 8;
          }
  
          for (; k < K; ++k) {
              acc0 += static_cast<int32_t>(*(Aptr0 + 0)) * B[k];
              acc1 += static_cast<int32_t>(*(Aptr0 + 1)) * B[k];
              acc2 += static_cast<int32_t>(*(Aptr0 + 2)) * B[k];
              acc3 += static_cast<int32_t>(*(Aptr0 + 3)) * B[k];
              Aptr0 += 4;
          }
  
          C[0] = acc0;
          C[1] = acc1;
          C[2] = acc2;
          C[3] = acc3;
  
          Aptr += Astride;
          C += Cstride;
      }
  }
  
  }  // namespace
  
  #if MGB_ENABLE_DOT
  namespace {
  MEGDNN_ATTRIBUTE_TARGET("dotprod")
  void gemv_naive_n_dot(const int8_t* __restrict A, const int8_t* __restrict B,
                    int32_t* __restrict C, size_t M, size_t N, size_t K,
                    size_t Astride, size_t Bstride, size_t Cstride) {
      megdnn_assert(N == 1 && Bstride == 1);
      size_t m = 0;
      for (; m + 2 <= M; m += 2) {
          int32_t acc[4];
          int32x4_t acc_neon = vdupq_n_s32(0);
          size_t k = 0;
          for (; k + 16 <= K; k += 16) {
              int64x2_t a0 = vreinterpretq_s64_s8(vld1q_s8(A + m * Astride + k));
              int64x2_t a1 =
                      vreinterpretq_s64_s8(vld1q_s8(A + (m + 1) * Astride + k));
              //! the first 8 elements is m, the last 8 elements is m + 1
              int8x16_t a2 = vreinterpretq_s8_s64(vzip1q_s64(a0, a1));
              int8x16_t a3 = vreinterpretq_s8_s64(vzip2q_s64(a0, a1));
  
              int64x2_t b0 = vreinterpretq_s64_s8(vld1q_s8(B + k));
              int8x16_t b2 = vreinterpretq_s8_s64(vzip1q_s64(b0, b0));
              int8x16_t b3 = vreinterpretq_s8_s64(vzip2q_s64(b0, b0));
  
              acc_neon = vdotq_s32(acc_neon, a2, b2);
              acc_neon = vdotq_s32(acc_neon, a3, b3);
          }
          vst1q_s32(acc, acc_neon);
  
          for (; k + 8 <= K; k += 8) {
              int8x8_t a0 = vld1_s8(A + m * Astride + k);
              int8x8_t a1 = vld1_s8(A + (m + 1) * Astride + k);
              int8x8_t b0 = vld1_s8(B + k);
              uint32x2_t zero = vdup_n_s32(0);
              acc[0] += vaddv_s32(vdot_s32(zero, a0, b0));
              zero = vdup_n_s32(0);
              acc[3] += vaddv_s32(vdot_s32(zero, a1, b0));
          }
  
          for (; k < K; ++k) {
              acc[0] += static_cast<int32_t>(A[m * Astride + k]) * B[k];
              acc[3] += static_cast<int32_t>(A[(m + 1) * Astride + k]) * B[k];
          }
          C[m * Cstride] = acc[0] + acc[1];
          C[(m + 1) * Cstride] = acc[2] + acc[3];
      }
  
      for (; m < M; ++m) {
          int32_t acc[4];
          int32x4_t acc_neon = vdupq_n_s32(0);
          size_t k = 0;
          for (; k + 16 <= K; k += 16) {
              int8x16_t a0 = vld1q_s8(A + m * Astride + k);
              int8x16_t b0 = vld1q_s8(B + k);
              acc_neon = vdotq_s32(acc_neon, a0, b0);
          }
          vst1q_s32(acc, acc_neon);
  
          for (; k + 8 <= K; k += 8) {
              int8x8_t a0 = vld1_s8(A + m * Astride + k);
              int8x8_t b0 = vld1_s8(B + k);
              uint32x2_t zero = vdup_n_s32(0);
              acc[0] += vaddv_s32(vdot_s32(zero, a0, b0));
          }
  
          for (; k < K; ++k) {
              acc[0] += static_cast<int32_t>(A[m * Astride + k]) * B[k];
          }
          C[m * Cstride] = acc[0] + acc[1] + acc[2] + acc[3];
      }
  }
  
  MEGDNN_ATTRIBUTE_TARGET("dotprod")
  void gemv_naive_n_mk4_dotprod(const int8_t* __restrict A, const int8_t* __restrict B,
                        int32_t* __restrict C, size_t M, size_t N, size_t K,
                        size_t Astride, size_t Bstride, size_t Cstride) {
      constexpr size_t PACK_SIZE = 4;
      megdnn_assert(N == 1 && Bstride == 4);
  
      auto Aptr = A;
      size_t m = 0;
      for (; m < M; m += PACK_SIZE) {
          auto Bptr = B;
          auto Aptr0 = Aptr;
          int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0;
          size_t k = 0;
          if (k + 16 <= K) {
              int32x4_t acc_neon[4];
              acc_neon[0] = vdupq_n_s32(0);
              acc_neon[1] = vdupq_n_s32(0);
              acc_neon[2] = vdupq_n_s32(0);
              acc_neon[3] = vdupq_n_s32(0);
              for (; k + 16 <= K; k += 16) {
                  int8x16x4_t a = vld4q_s8(Aptr0);
                  int8x16_t b = vld1q_s8(Bptr);
  
                  acc_neon[0] = vdotq_s32(acc_neon[0], a.val[0], b);
                  acc_neon[1] = vdotq_s32(acc_neon[1], a.val[1], b);
                  acc_neon[2] = vdotq_s32(acc_neon[2], a.val[2], b);
                  acc_neon[3] = vdotq_s32(acc_neon[3], a.val[3], b);
  
                  Bptr += 16;
                  Aptr0 += PACK_SIZE * 16;
              }
              acc0 = vaddvq_s32(acc_neon[0]);
              acc1 = vaddvq_s32(acc_neon[1]);
              acc2 = vaddvq_s32(acc_neon[2]);
              acc3 = vaddvq_s32(acc_neon[3]);
          }
  
          if (k + 8 <= K) {
              int32x2_t acc_neon[4];
              acc_neon[0] = vdup_n_s32(0);
              acc_neon[1] = vdup_n_s32(0);
              acc_neon[2] = vdup_n_s32(0);
              acc_neon[3] = vdup_n_s32(0);
  
              int8x8x4_t a = vld4_s8(Aptr0);
              int8x8_t b = vld1_s8(Bptr);
              acc_neon[0] = vdot_s32(acc_neon[0], a.val[0], b);
              acc_neon[1] = vdot_s32(acc_neon[1], a.val[1], b);
              acc_neon[2] = vdot_s32(acc_neon[2], a.val[2], b);
              acc_neon[3] = vdot_s32(acc_neon[3], a.val[3], b);
  
              Bptr += 8;
              Aptr0 += PACK_SIZE * 8;
              k += 8;
  
              acc0 += vaddv_s32(acc_neon[0]);
              acc1 += vaddv_s32(acc_neon[1]);
              acc2 += vaddv_s32(acc_neon[2]);
              acc3 += vaddv_s32(acc_neon[3]);
          }
  
          for (; k < K; ++k) {
              acc0 += static_cast<int32_t>(*(Aptr0 + 0)) * B[k];
              acc1 += static_cast<int32_t>(*(Aptr0 + 1)) * B[k];
              acc2 += static_cast<int32_t>(*(Aptr0 + 2)) * B[k];
              acc3 += static_cast<int32_t>(*(Aptr0 + 3)) * B[k];
              Aptr0 += 4;
          }
  
          C[0] = acc0;
          C[1] = acc1;
          C[2] = acc2;
          C[3] = acc3;
  
          Aptr += Astride;
          C += Cstride;
      }
  }
  
  MEGDNN_ATTRIBUTE_TARGET("dotprod")
  void gemv_naive_n_mk4_dot(const int8_t* __restrict A,
                            const int8_t* __restrict B, int32_t* __restrict C,
                            size_t M, size_t N, size_t K, size_t Astride,
                            size_t Bstride, size_t Cstride) {
      constexpr size_t PACK_SIZE = 4;
      megdnn_assert(N == 1 && Bstride == 4);
  
      auto Aptr = A;
      size_t m = 0;
      for (; m < M; m += PACK_SIZE) {
          auto Bptr = B;
          auto Aptr0 = Aptr;
          size_t k = 0;
          int32x4_t acc_neon;
          acc_neon = vdupq_n_s32(0);
          for (; k + 16 <= K; k += 16) {
              int8x16_t a0 = vld1q_s8(Aptr0);
              int8x16_t a1 = vld1q_s8(Aptr0 + 16);
              int8x16_t a2 = vld1q_s8(Aptr0 + 32);
              int8x16_t a3 = vld1q_s8(Aptr0 + 48);
              int8x16_t b = vld1q_s8(Bptr);
              acc_neon = vdotq_laneq_s32(acc_neon, a0, b, 0);
              acc_neon = vdotq_laneq_s32(acc_neon, a1, b, 1);
              acc_neon = vdotq_laneq_s32(acc_neon, a2, b, 2);
              acc_neon = vdotq_laneq_s32(acc_neon, a3, b, 3);
              Bptr += 16;
              Aptr0 += PACK_SIZE * 16;
          }
  
          if (k + 8 <= K) {
              int8x16_t a0 = vld1q_s8(Aptr0);
              int8x16_t a1 = vld1q_s8(Aptr0 + 16);
              int8x8_t b = vld1_s8(Bptr);
              acc_neon = vdotq_lane_s32(acc_neon, a0, b, 0);
              acc_neon = vdotq_lane_s32(acc_neon, a1, b, 1);
              Bptr += 8;
              Aptr0 += PACK_SIZE * 8;
              k += 8;
          }
  
          if (k + 4 <= K) {
              int8x16_t a = vld1q_s8(Aptr0);
              int32_t tmp = *(reinterpret_cast<const int32_t*>(Bptr));
              int8x8_t b = vdup_n_s32(tmp);
              acc_neon = vdotq_lane_s32(acc_neon, a, b, 0);
          }
  
          vst1q_s32(C, acc_neon);
          Aptr += Astride;
          C += Cstride;
      }
  }
  
  }  // namespace
  #endif
  
  bool arm_common::is_gemv_like_preferred_int8(bool transposeA, bool transposeB,
                                               size_t M, size_t N, size_t K,
                                               size_t LDA, size_t LDB,
                                               size_t LDC) {
      MEGDNN_MARK_USED_VAR(LDA);
      MEGDNN_MARK_USED_VAR(LDB);
      MEGDNN_MARK_USED_VAR(LDC);
      MEGDNN_MARK_USED_VAR(M);
      MEGDNN_MARK_USED_VAR(K);
      if (transposeA)
          return false;
      if (transposeB)
          return false;
  
      return N == 1 && LDB == 1;
  }
  
  void arm_common::gemv_like(const int8_t* __restrict A,
                             const int8_t* __restrict B, int32_t* __restrict C,
                             size_t M, size_t N, size_t K, size_t Astride,
                             size_t Bstride, size_t Cstride) {
      megdnn_assert(N == 1);
      MIDOUT_BEGIN(megdnn_arm_common_int8_gemv,
                   midout_iv("INT8_gemv_like"_hash)) {
  #if MGB_ENABLE_DOT
          if (cpuinfo_has_arm_neon_dot()) {
              return gemv_naive_n_dot(A, B, C, M, N, K, Astride, Bstride,
                                      Cstride);
          } else {
              return gemv_naive_n(A, B, C, M, N, K, Astride, Bstride, Cstride);
          }
  #else
          return gemv_naive_n(A, B, C, M, N, K, Astride, Bstride, Cstride);
  #endif
      }
      MIDOUT_END();
  }
  
  void arm_common::gemv_like_mk4(const int8_t* __restrict A,
                                 const int8_t* __restrict B,
                                 int32_t* __restrict C, size_t M, size_t N,
                                 size_t K, size_t Astride, size_t Bstride,
                                 size_t Cstride) {
      megdnn_assert(N == 1);
      MIDOUT_BEGIN(megdnn_arm_common_int8_gemv,
                   midout_iv("INT8_gemv_like_mk4"_hash)) {
  #if MGB_ENABLE_DOT
          if (cpuinfo_has_arm_neon_dot()) {
              return gemv_naive_n_mk4_dotprod(A, B, C, M, N, K, Astride, Bstride,
                                              Cstride);
          } else {
              return gemv_naive_n_mk4(A, B, C, M, N, K, Astride, Bstride,
                                      Cstride);
          }
  #else
          return gemv_naive_n_mk4(A, B, C, M, N, K, Astride, Bstride, Cstride);
  #endif
      }
      MIDOUT_END();
  }
  
  #if MGB_ENABLE_DOT
  void arm_common::gemv_like_mk4_dot(const int8_t* __restrict A,
                                     const int8_t* __restrict B,
                                     int32_t* __restrict C, size_t M, size_t N,
                                     size_t K, size_t Astride, size_t Bstride,
                                     size_t Cstride) {
      megdnn_assert(N == 1);
      MIDOUT_BEGIN(megdnn_arm_common_int8_gemv,
                   midout_iv("INT8_gemv_like_mk4_dot"_hash)) {
          return gemv_naive_n_mk4_dot(A, B, C, M, N, K, Astride, Bstride,
                                      Cstride);
      }
      MIDOUT_END();
  }
  #endif
  
  
  // vim: syntax=cpp.doxygen
  