// Copyright 2018 The MACE Authors. All Rights Reserved.
//
// 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.

#include <algorithm>
#include <cstring>
#include <vector>

#include "mace/core/tensor.h"
#include "mace/core/runtime/cpu/cpu_runtime.h"
#include "mace/ops/gemm.h"

/**
 * Gemm does fast matrix multiplications with batch.
 * It is optimized for arm64-v8 and armeabi-v7a using neon.
 *
 * We adopt two-level tiling to make better use of l1 cache and register.
 * For register tiling, function like GemmXYZ computes gemm for
 * matrix[X, Y] * matrix[Y, Z] with all data being able to fit in register.
 * For cache tiling, we try to compute one block of multiplication with
 * two input matrices and one output matrix fit in l1 cache.
 */

#if defined(MACE_ENABLE_NEON)
#include <arm_neon.h>
#endif

#if defined(MACE_ENABLE_NEON) && !defined(__aarch64__)
#define vaddvq_f32(v) ((v)[0] + (v)[1] + (v)[2] + (v)[3])
#endif

namespace mace {
namespace ops {

namespace {
inline void GemmBlock(const float *A,
                      const float *B,
                      const index_t height,
                      const index_t K,
                      const index_t width,
                      const index_t stride_a,
                      const index_t stride_b,
                      const index_t stride_c,
                      float *C) {
  for (int i = 0; i < height; ++i) {
    for (int j = 0; j < width; ++j) {
      for (int k = 0; k < K; ++k) {
        C[i * stride_c + j] += A[i * stride_a + k] * B[k * stride_b + j];
      }
    }
  }
}

#if defined(MACE_ENABLE_NEON)
#if defined(__aarch64__)
#define MACE_GEMM_PART_CAL_8(RC, RA, RAN)        \
  c##RC = vfmaq_laneq_f32(c##RC, b0, a##RA, 0);  \
  c##RC = vfmaq_laneq_f32(c##RC, b1, a##RA, 1);  \
  c##RC = vfmaq_laneq_f32(c##RC, b2, a##RA, 2);  \
  c##RC = vfmaq_laneq_f32(c##RC, b3, a##RA, 3);  \
  c##RC = vfmaq_laneq_f32(c##RC, b4, a##RAN, 0); \
  c##RC = vfmaq_laneq_f32(c##RC, b5, a##RAN, 1); \
  c##RC = vfmaq_laneq_f32(c##RC, b6, a##RAN, 2); \
  c##RC = vfmaq_laneq_f32(c##RC, b7, a##RAN, 3);
#else
#define MACE_GEMM_PART_CAL_8(RC, RA, RAN)                      \
  c##RC = vmlaq_lane_f32(c##RC, b0, vget_low_f32(a##RA), 0);   \
  c##RC = vmlaq_lane_f32(c##RC, b1, vget_low_f32(a##RA), 1);   \
  c##RC = vmlaq_lane_f32(c##RC, b2, vget_high_f32(a##RA), 0);  \
  c##RC = vmlaq_lane_f32(c##RC, b3, vget_high_f32(a##RA), 1);  \
  c##RC = vmlaq_lane_f32(c##RC, b4, vget_low_f32(a##RAN), 0);  \
  c##RC = vmlaq_lane_f32(c##RC, b5, vget_low_f32(a##RAN), 1);  \
  c##RC = vmlaq_lane_f32(c##RC, b6, vget_high_f32(a##RAN), 0); \
  c##RC = vmlaq_lane_f32(c##RC, b7, vget_high_f32(a##RAN), 1);
#endif
#endif

#if defined(MACE_ENABLE_NEON)
#if defined(__aarch64__)
#define MACE_GEMM_PART_CAL_4(RC)                \
  c##RC = vfmaq_laneq_f32(c##RC, b0, a##RC, 0); \
  c##RC = vfmaq_laneq_f32(c##RC, b1, a##RC, 1); \
  c##RC = vfmaq_laneq_f32(c##RC, b2, a##RC, 2); \
  c##RC = vfmaq_laneq_f32(c##RC, b3, a##RC, 3);
#else
#define MACE_GEMM_PART_CAL_4(RC)                              \
  c##RC = vmlaq_lane_f32(c##RC, b0, vget_low_f32(a##RC), 0);  \
  c##RC = vmlaq_lane_f32(c##RC, b1, vget_low_f32(a##RC), 1);  \
  c##RC = vmlaq_lane_f32(c##RC, b2, vget_high_f32(a##RC), 0); \
  c##RC = vmlaq_lane_f32(c##RC, b3, vget_high_f32(a##RC), 1);
#endif
#endif

inline void Gemm144(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  MACE_UNUSED(stride_a);
  MACE_UNUSED(stride_c);
  float32x4_t a0;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0;

  a0 = vld1q_f32(a_ptr);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);

  MACE_GEMM_PART_CAL_4(0);

  vst1q_f32(c_ptr, c0);
#else
  GemmBlock(a_ptr, b_ptr, 1, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm244(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0, c1;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 1 * stride_a);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);

  MACE_GEMM_PART_CAL_4(0);
  MACE_GEMM_PART_CAL_4(1);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
#else
  GemmBlock(a_ptr, b_ptr, 2, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm344(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0, c1, c2;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 1 * stride_a);
  a2 = vld1q_f32(a_ptr + 2 * stride_a);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);

  MACE_GEMM_PART_CAL_4(0);
  MACE_GEMM_PART_CAL_4(1);
  MACE_GEMM_PART_CAL_4(2);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
#else
  GemmBlock(a_ptr, b_ptr, 3, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm444(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0, c1, c2, c3;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 1 * stride_a);
  a2 = vld1q_f32(a_ptr + 2 * stride_a);
  a3 = vld1q_f32(a_ptr + 3 * stride_a);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);

  MACE_GEMM_PART_CAL_4(0);
  MACE_GEMM_PART_CAL_4(1);
  MACE_GEMM_PART_CAL_4(2);
  MACE_GEMM_PART_CAL_4(3);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
#else
  GemmBlock(a_ptr, b_ptr, 4, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm544(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0, c1, c2, c3, c4;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 1 * stride_a);
  a2 = vld1q_f32(a_ptr + 2 * stride_a);
  a3 = vld1q_f32(a_ptr + 3 * stride_a);
  a4 = vld1q_f32(a_ptr + 4 * stride_a);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);

  MACE_GEMM_PART_CAL_4(0);
  MACE_GEMM_PART_CAL_4(1);
  MACE_GEMM_PART_CAL_4(2);
  MACE_GEMM_PART_CAL_4(3);
  MACE_GEMM_PART_CAL_4(4);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
#else
  GemmBlock(a_ptr, b_ptr, 5, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm644(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5;
  float32x4_t b0, b1, b2, b3;
  float32x4_t c0, c1, c2, c3, c4, c5;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 1 * stride_a);
  a2 = vld1q_f32(a_ptr + 2 * stride_a);
  a3 = vld1q_f32(a_ptr + 3 * stride_a);
  a4 = vld1q_f32(a_ptr + 4 * stride_a);
  a5 = vld1q_f32(a_ptr + 5 * stride_a);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);
  c5 = vld1q_f32(c_ptr + 5 * stride_c);

  MACE_GEMM_PART_CAL_4(0);
  MACE_GEMM_PART_CAL_4(1);
  MACE_GEMM_PART_CAL_4(2);
  MACE_GEMM_PART_CAL_4(3);
  MACE_GEMM_PART_CAL_4(4);
  MACE_GEMM_PART_CAL_4(5);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
  vst1q_f32(c_ptr + 5 * stride_c, c5);
#else
  GemmBlock(a_ptr, b_ptr, 6, 4, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm884(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14,
      a15;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2, c3, c4, c5, c6, c7;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
  a6 = vld1q_f32(a_ptr + 3 * stride_a);
  a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);
  a8 = vld1q_f32(a_ptr + 4 * stride_a);
  a9 = vld1q_f32(a_ptr + 4 * stride_a + 4);
  a10 = vld1q_f32(a_ptr + 5 * stride_a);
  a11 = vld1q_f32(a_ptr + 5 * stride_a + 4);
  a12 = vld1q_f32(a_ptr + 6 * stride_a);
  a13 = vld1q_f32(a_ptr + 6 * stride_a + 4);
  a14 = vld1q_f32(a_ptr + 7 * stride_a);
  a15 = vld1q_f32(a_ptr + 7 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);
  c5 = vld1q_f32(c_ptr + 5 * stride_c);
  c6 = vld1q_f32(c_ptr + 6 * stride_c);
  c7 = vld1q_f32(c_ptr + 7 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);
  MACE_GEMM_PART_CAL_8(3, 6, 7);
  MACE_GEMM_PART_CAL_8(4, 8, 9);
  MACE_GEMM_PART_CAL_8(5, 10, 11);
  MACE_GEMM_PART_CAL_8(6, 12, 13);
  MACE_GEMM_PART_CAL_8(7, 14, 15);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
  vst1q_f32(c_ptr + 5 * stride_c, c5);
  vst1q_f32(c_ptr + 6 * stride_c, c6);
  vst1q_f32(c_ptr + 7 * stride_c, c7);
#else
  GemmBlock(a_ptr, b_ptr, 8, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm184(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  MACE_UNUSED(stride_a);
  MACE_UNUSED(stride_c);

  float32x4_t a0, a1;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);

  MACE_GEMM_PART_CAL_8(0, 0, 1);

  vst1q_f32(c_ptr, c0);
#else
  GemmBlock(a_ptr, b_ptr, 1, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm284(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
#else
  GemmBlock(a_ptr, b_ptr, 2, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm384(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
#else
  GemmBlock(a_ptr, b_ptr, 3, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm484(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5, a6, a7;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2, c3;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
  a6 = vld1q_f32(a_ptr + 3 * stride_a);
  a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);
  MACE_GEMM_PART_CAL_8(3, 6, 7);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
#else
  GemmBlock(a_ptr, b_ptr, 4, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm584(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5, a6, a7, a8, a9;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2, c3, c4;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
  a6 = vld1q_f32(a_ptr + 3 * stride_a);
  a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);
  a8 = vld1q_f32(a_ptr + 4 * stride_a);
  a9 = vld1q_f32(a_ptr + 4 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);
  MACE_GEMM_PART_CAL_8(3, 6, 7);
  MACE_GEMM_PART_CAL_8(4, 8, 9);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
#else
  GemmBlock(a_ptr, b_ptr, 5, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm684(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2, c3, c4, c5;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
  a6 = vld1q_f32(a_ptr + 3 * stride_a);
  a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);
  a8 = vld1q_f32(a_ptr + 4 * stride_a);
  a9 = vld1q_f32(a_ptr + 4 * stride_a + 4);
  a10 = vld1q_f32(a_ptr + 5 * stride_a);
  a11 = vld1q_f32(a_ptr + 5 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);
  c5 = vld1q_f32(c_ptr + 5 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);
  MACE_GEMM_PART_CAL_8(3, 6, 7);
  MACE_GEMM_PART_CAL_8(4, 8, 9);
  MACE_GEMM_PART_CAL_8(5, 10, 11);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
  vst1q_f32(c_ptr + 5 * stride_c, c5);

#else
  GemmBlock(a_ptr, b_ptr, 6, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void Gemm784(const float *a_ptr,
                    const float *b_ptr,
                    const index_t stride_a,
                    const index_t stride_b,
                    const index_t stride_c,
                    float *c_ptr) {
#if defined(MACE_ENABLE_NEON)
  float32x4_t a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13;
  float32x4_t b0, b1, b2, b3, b4, b5, b6, b7;
  float32x4_t c0, c1, c2, c3, c4, c5, c6;

  a0 = vld1q_f32(a_ptr);
  a1 = vld1q_f32(a_ptr + 4);
  a2 = vld1q_f32(a_ptr + 1 * stride_a);
  a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
  a4 = vld1q_f32(a_ptr + 2 * stride_a);
  a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
  a6 = vld1q_f32(a_ptr + 3 * stride_a);
  a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);
  a8 = vld1q_f32(a_ptr + 4 * stride_a);
  a9 = vld1q_f32(a_ptr + 4 * stride_a + 4);
  a10 = vld1q_f32(a_ptr + 5 * stride_a);
  a11 = vld1q_f32(a_ptr + 5 * stride_a + 4);
  a12 = vld1q_f32(a_ptr + 6 * stride_a);
  a13 = vld1q_f32(a_ptr + 6 * stride_a + 4);

  b0 = vld1q_f32(b_ptr);
  b1 = vld1q_f32(b_ptr + 1 * stride_b);
  b2 = vld1q_f32(b_ptr + 2 * stride_b);
  b3 = vld1q_f32(b_ptr + 3 * stride_b);
  b4 = vld1q_f32(b_ptr + 4 * stride_b);
  b5 = vld1q_f32(b_ptr + 5 * stride_b);
  b6 = vld1q_f32(b_ptr + 6 * stride_b);
  b7 = vld1q_f32(b_ptr + 7 * stride_b);

  c0 = vld1q_f32(c_ptr);
  c1 = vld1q_f32(c_ptr + 1 * stride_c);
  c2 = vld1q_f32(c_ptr + 2 * stride_c);
  c3 = vld1q_f32(c_ptr + 3 * stride_c);
  c4 = vld1q_f32(c_ptr + 4 * stride_c);
  c5 = vld1q_f32(c_ptr + 5 * stride_c);
  c6 = vld1q_f32(c_ptr + 6 * stride_c);

  MACE_GEMM_PART_CAL_8(0, 0, 1);
  MACE_GEMM_PART_CAL_8(1, 2, 3);
  MACE_GEMM_PART_CAL_8(2, 4, 5);
  MACE_GEMM_PART_CAL_8(3, 6, 7);
  MACE_GEMM_PART_CAL_8(4, 8, 9);
  MACE_GEMM_PART_CAL_8(5, 10, 11);
  MACE_GEMM_PART_CAL_8(6, 12, 13);

  vst1q_f32(c_ptr, c0);
  vst1q_f32(c_ptr + 1 * stride_c, c1);
  vst1q_f32(c_ptr + 2 * stride_c, c2);
  vst1q_f32(c_ptr + 3 * stride_c, c3);
  vst1q_f32(c_ptr + 4 * stride_c, c4);
  vst1q_f32(c_ptr + 5 * stride_c, c5);
  vst1q_f32(c_ptr + 6 * stride_c, c6);

#else
  GemmBlock(a_ptr, b_ptr, 7, 8, 4, stride_a, stride_b, stride_c, c_ptr);
#endif
}

inline void GemmTile(const float *A,
                     const float *B,
                     const index_t height,
                     const index_t K,
                     const index_t width,
                     const index_t stride_a,
                     const index_t stride_b,
                     const index_t stride_c,
                     float *C) {
#if defined(MACE_ENABLE_NEON)
  index_t h = 0;
  index_t w = 0;
  index_t k = 0;
#if defined(__aarch64__)
  int reg_height_tile = 8;
  int reg_K_tile = 8;
#else
  int reg_height_tile = 6;
  int reg_K_tile = 4;
#endif

  for (h = 0; h < height - reg_height_tile + 1; h += reg_height_tile) {
    for (k = 0; k < K - reg_K_tile + 1; k += reg_K_tile) {
      const float *a_ptr = A + (h * stride_a + k);
#if defined(__aarch64__) && defined(__clang__)
      int nw = width >> 2;
      if (nw > 0) {
        // load A
        float32x4_t a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13,
            a14, a15;
        a0 = vld1q_f32(a_ptr);
        a1 = vld1q_f32(a_ptr + 4);
        a2 = vld1q_f32(a_ptr + 1 * stride_a);
        a3 = vld1q_f32(a_ptr + 1 * stride_a + 4);
        a4 = vld1q_f32(a_ptr + 2 * stride_a);
        a5 = vld1q_f32(a_ptr + 2 * stride_a + 4);
        a6 = vld1q_f32(a_ptr + 3 * stride_a);
        a7 = vld1q_f32(a_ptr + 3 * stride_a + 4);
        a8 = vld1q_f32(a_ptr + 4 * stride_a);
        a9 = vld1q_f32(a_ptr + 4 * stride_a + 4);
        a10 = vld1q_f32(a_ptr + 5 * stride_a);
        a11 = vld1q_f32(a_ptr + 5 * stride_a + 4);
        a12 = vld1q_f32(a_ptr + 6 * stride_a);
        a13 = vld1q_f32(a_ptr + 6 * stride_a + 4);
        a14 = vld1q_f32(a_ptr + 7 * stride_a);
        a15 = vld1q_f32(a_ptr + 7 * stride_a + 4);

        const float *b_ptr0 = B + k * stride_b;
        const float *b_ptr1 = B + (k + 1) * stride_b;
        const float *b_ptr2 = B + (k + 2) * stride_b;
        const float *b_ptr3 = B + (k + 3) * stride_b;
        const float *b_ptr4 = B + (k + 4) * stride_b;
        const float *b_ptr5 = B + (k + 5) * stride_b;
        const float *b_ptr6 = B + (k + 6) * stride_b;
        const float *b_ptr7 = B + (k + 7) * stride_b;

        float *c_ptr0 = C + h * stride_c;
        float *c_ptr1 = C + (h + 1) * stride_c;
        float *c_ptr2 = C + (h + 2) * stride_c;
        float *c_ptr3 = C + (h + 3) * stride_c;
        float *c_ptr4 = C + (h + 4) * stride_c;
        float *c_ptr5 = C + (h + 5) * stride_c;
        float *c_ptr6 = C + (h + 6) * stride_c;
        float *c_ptr7 = C + (h + 7) * stride_c;

        asm volatile(
            "0:                                 \n"

            "prfm   pldl1keep, [%9, #128]       \n"
            "ld1    {v16.4s}, [%9], #16         \n"

            "prfm   pldl1keep, [%1, #128]       \n"
            "ld1    {v18.4s}, [%1]              \n"

            "prfm   pldl1keep, [%2, #128]       \n"
            "ld1    {v19.4s}, [%2]              \n"

            "prfm   pldl1keep, [%3, #128]       \n"
            "ld1    {v20.4s}, [%3]              \n"
            "prfm   pldl1keep, [%4, #128]       \n"
            "ld1    {v21.4s}, [%4]              \n"
            "prfm   pldl1keep, [%5, #128]       \n"
            "ld1    {v22.4s}, [%5]              \n"
            "prfm   pldl1keep, [%6, #128]       \n"
            "ld1    {v23.4s}, [%6]              \n"
            "prfm   pldl1keep, [%7, #128]       \n"
            "ld1    {v24.4s}, [%7]              \n"
            "prfm   pldl1keep, [%8, #128]       \n"
            "ld1    {v25.4s}, [%8]              \n"
            "prfm   pldl1keep, [%10, #128]      \n"
            "ld1    {v17.4s}, [%10], #16        \n"

            "fmla   v18.4s, v16.4s, %34.s[0]    \n"
            "fmla   v19.4s, v16.4s, %35.s[0]    \n"
            "fmla   v20.4s, v16.4s, %36.s[0]    \n"
            "fmla   v21.4s, v16.4s, %37.s[0]    \n"

            "fmla   v22.4s, v16.4s, %38.s[0]    \n"
            "fmla   v23.4s, v16.4s, %39.s[0]    \n"
            "fmla   v24.4s, v16.4s, %40.s[0]    \n"
            "fmla   v25.4s, v16.4s, %41.s[0]    \n"

            "fmla   v18.4s, v17.4s, %34.s[1]    \n"
            "fmla   v19.4s, v17.4s, %35.s[1]    \n"
            "fmla   v20.4s, v17.4s, %36.s[1]    \n"
            "fmla   v21.4s, v17.4s, %37.s[1]    \n"

            "prfm   pldl1keep, [%11, #128]      \n"
            "ld1    {v16.4s}, [%11], #16        \n"

            "fmla   v22.4s, v17.4s, %38.s[1]    \n"
            "fmla   v23.4s, v17.4s, %39.s[1]    \n"
            "fmla   v24.4s, v17.4s, %40.s[1]    \n"
            "fmla   v25.4s, v17.4s, %41.s[1]    \n"

            "fmla   v18.4s, v16.4s, %34.s[2]    \n"
            "fmla   v19.4s, v16.4s, %35.s[2]    \n"
            "fmla   v20.4s, v16.4s, %36.s[2]    \n"
            "fmla   v21.4s, v16.4s, %37.s[2]    \n"

            "prfm   pldl1keep, [%12, #128]      \n"
            "ld1    {v17.4s}, [%12], #16        \n"

            "fmla   v22.4s, v16.4s, %38.s[2]    \n"
            "fmla   v23.4s, v16.4s, %39.s[2]    \n"
            "fmla   v24.4s, v16.4s, %40.s[2]    \n"
            "fmla   v25.4s, v16.4s, %41.s[2]    \n"

            "fmla   v18.4s, v17.4s, %34.s[3]    \n"
            "fmla   v19.4s, v17.4s, %35.s[3]    \n"
            "fmla   v20.4s, v17.4s, %36.s[3]    \n"
            "fmla   v21.4s, v17.4s, %37.s[3]    \n"

            "prfm   pldl1keep, [%13, #128]      \n"
            "ld1    {v16.4s}, [%13], #16        \n"

            "fmla   v22.4s, v17.4s, %38.s[3]    \n"
            "fmla   v23.4s, v17.4s, %39.s[3]    \n"
            "fmla   v24.4s, v17.4s, %40.s[3]    \n"
            "fmla   v25.4s, v17.4s, %41.s[3]    \n"

            "fmla   v18.4s, v16.4s, %42.s[0]    \n"
            "fmla   v19.4s, v16.4s, %43.s[0]    \n"
            "fmla   v20.4s, v16.4s, %44.s[0]    \n"
            "fmla   v21.4s, v16.4s, %45.s[0]    \n"

            "prfm   pldl1keep, [%14, #128]      \n"
            "ld1    {v17.4s}, [%14], #16        \n"

            "fmla   v22.4s, v16.4s, %46.s[0]    \n"
            "fmla   v23.4s, v16.4s, %47.s[0]    \n"
            "fmla   v24.4s, v16.4s, %48.s[0]    \n"
            "fmla   v25.4s, v16.4s, %49.s[0]    \n"

            "fmla   v18.4s, v17.4s, %42.s[1]    \n"
            "fmla   v19.4s, v17.4s, %43.s[1]    \n"
            "fmla   v20.4s, v17.4s, %44.s[1]    \n"
            "fmla   v21.4s, v17.4s, %45.s[1]    \n"

            "prfm   pldl1keep, [%15, #128]      \n"
            "ld1    {v16.4s}, [%15], #16        \n"

            "fmla   v22.4s, v17.4s, %46.s[1]    \n"
            "fmla   v23.4s, v17.4s, %47.s[1]    \n"
            "fmla   v24.4s, v17.4s, %48.s[1]    \n"
            "fmla   v25.4s, v17.4s, %49.s[1]    \n"

            "fmla   v18.4s, v16.4s, %42.s[2]    \n"
            "fmla   v19.4s, v16.4s, %43.s[2]    \n"
            "fmla   v20.4s, v16.4s, %44.s[2]    \n"
            "fmla   v21.4s, v16.4s, %45.s[2]    \n"

            "prfm   pldl1keep, [%16, #128]      \n"
            "ld1    {v17.4s}, [%16], #16        \n"

            "fmla   v22.4s, v16.4s, %46.s[2]    \n"
            "fmla   v23.4s, v16.4s, %47.s[2]    \n"
            "fmla   v24.4s, v16.4s, %48.s[2]    \n"
            "fmla   v25.4s, v16.4s, %49.s[2]    \n"

            "fmla   v18.4s, v17.4s, %42.s[3]    \n"
            "fmla   v19.4s, v17.4s, %43.s[3]    \n"
            "fmla   v20.4s, v17.4s, %44.s[3]    \n"
            "fmla   v21.4s, v17.4s, %45.s[3]    \n"

            "st1    {v18.4s}, [%1], #16         \n"
            "st1    {v19.4s}, [%2], #16         \n"
            "st1    {v20.4s}, [%3], #16         \n"
            "st1    {v21.4s}, [%4], #16         \n"

            "fmla   v22.4s, v17.4s, %46.s[3]    \n"
            "fmla   v23.4s, v17.4s, %47.s[3]    \n"
            "fmla   v24.4s, v17.4s, %48.s[3]    \n"
            "fmla   v25.4s, v17.4s, %49.s[3]    \n"

            "subs   %w0, %w0, #1                \n"

            "st1    {v22.4s}, [%5], #16         \n"
            "st1    {v23.4s}, [%6], #16         \n"
            "st1    {v24.4s}, [%7], #16         \n"
            "st1    {v25.4s}, [%8], #16         \n"

            "bne    0b                          \n"
            : "=r"(nw),      // 0
              "=r"(c_ptr0),  // 1
              "=r"(c_ptr1),  // 2
              "=r"(c_ptr2),  // 3
              "=r"(c_ptr3),  // 4
              "=r"(c_ptr4),  // 5
              "=r"(c_ptr5),  // 6
              "=r"(c_ptr6),  // 7
              "=r"(c_ptr7),  // 8
              "=r"(b_ptr0),  // 9
              "=r"(b_ptr1),  // 10
              "=r"(b_ptr2),  // 11
              "=r"(b_ptr3),  // 12
              "=r"(b_ptr4),  // 13
              "=r"(b_ptr5),  // 14
              "=r"(b_ptr6),  // 15
              "=r"(b_ptr7)   // 16
            : "0"(nw),       // 17
              "1"(c_ptr0),   // 18
              "2"(c_ptr1),   // 19
              "3"(c_ptr2),   // 20
              "4"(c_ptr3),   // 21
              "5"(c_ptr4),   // 22
              "6"(c_ptr5),   // 23
              "7"(c_ptr6),   // 24
              "8"(c_ptr7),   // 25
              "9"(b_ptr0),   // 26
              "10"(b_ptr1),  // 27
              "11"(b_ptr2),  // 28
              "12"(b_ptr3),  // 29
              "13"(b_ptr4),  // 30
              "14"(b_ptr5),  // 31
              "15"(b_ptr6),  // 32
              "16"(b_ptr7),  // 33
              "w"(a0),       // 34
              "w"(a2),       // 35
              "w"(a4),       // 36
              "w"(a6),       // 37
              "w"(a8),       // 38
              "w"(a10),      // 39
              "w"(a12),      // 40
              "w"(a14),      // 41
              "w"(a1),       // 42
              "w"(a3),       // 43
              "w"(a5),       // 44
              "w"(a7),       // 45
              "w"(a9),       // 46
              "w"(a11),      // 47
              "w"(a13),      // 48
              "w"(a15)       // 49
            : "cc", "memory", "v16", "v17", "v18", "v19", "v20", "v21", "v22",
              "v23", "v24", "v25");

        w = (width >> 2) << 2;
      }
#elif defined(__aarch64__)  // gcc
      for (w = 0; w + 3 < width; w += 4) {
        const float *b_ptr = B + (k * stride_b + w);
        float *c_ptr = C + (h * stride_c + w);
        Gemm884(a_ptr, b_ptr, stride_a, stride_b, stride_c, c_ptr);
      }
#else                       // armv7
      int nw = width >> 2;
      if (nw > 0) {
        float32x4_t a0, a1, a2, a3, a4, a5;

        a0 = vld1q_f32(a_ptr);
        a1 = vld1q_f32(a_ptr + 1 * stride_a);
        a2 = vld1q_f32(a_ptr + 2 * stride_a);
        a3 = vld1q_f32(a_ptr + 3 * stride_a);
        a4 = vld1q_f32(a_ptr + 4 * stride_a);
        a5 = vld1q_f32(a_ptr + 5 * stride_a);

        const float *b_ptr0 = B + k * stride_b;
        const float *b_ptr1 = B + (k + 1) * stride_b;
        const float *b_ptr2 = B + (k + 2) * stride_b;
        const float *b_ptr3 = B + (k + 3) * stride_b;

        float *c_ptr0 = C + h * stride_c;
        float *c_ptr1 = C + (h + 1) * stride_c;
        float *c_ptr2 = C + (h + 2) * stride_c;
        float *c_ptr3 = C + (h + 3) * stride_c;
        float *c_ptr4 = C + (h + 4) * stride_c;
        float *c_ptr5 = C + (h + 5) * stride_c;

        asm volatile(
            "0:                                 \n"

            "pld    [%7, #128]                  \n"
            "vld1.f32    {d12-d13}, [%7]!       \n"
            "pld    [%1, #128]                  \n"
            "vld1.f32    {d16-d17}, [%1]        \n"
            "pld    [%2, #128]                  \n"
            "vld1.f32    {d18-d19}, [%2]        \n"

            "pld    [%3, #128]                  \n"
            "vld1.f32    {d20-d21}, [%3]        \n"
            "pld    [%4, #128]                  \n"
            "vld1.f32    {d22-d23}, [%4]        \n"
            "pld    [%5, #128]                  \n"
            "vld1.f32    {d24-d25}, [%5]        \n"
            "pld    [%6, #128]                  \n"
            "vld1.f32    {d26-d27}, [%6]        \n"

            "pld    [%8, #128]                  \n"
            "vld1.f32    {d14-d15}, [%8]!       \n"

            "vmla.f32   q8, q6, %e22[0]         \n"
            "vmla.f32   q9, q6, %e23[0]         \n"
            "vmla.f32   q10, q6, %e24[0]        \n"
            "vmla.f32   q11, q6, %e25[0]        \n"
            "vmla.f32   q12, q6, %e26[0]        \n"
            "vmla.f32   q13, q6, %e27[0]        \n"

            "pld    [%9, #128]                  \n"
            "vld1.f32    {d12-d13}, [%9]!       \n"

            "vmla.f32   q8, q7, %e22[1]         \n"
            "vmla.f32   q9, q7, %e23[1]         \n"
            "vmla.f32   q10, q7, %e24[1]        \n"
            "vmla.f32   q11, q7, %e25[1]        \n"
            "vmla.f32   q12, q7, %e26[1]        \n"
            "vmla.f32   q13, q7, %e27[1]        \n"

            "pld    [%10, #128]                 \n"
            "vld1.f32    {d14-d15}, [%10]!      \n"

            "vmla.f32   q8, q6, %f22[0]         \n"
            "vmla.f32   q9, q6, %f23[0]         \n"
            "vmla.f32   q10, q6, %f24[0]        \n"
            "vmla.f32   q11, q6, %f25[0]        \n"
            "vmla.f32   q12, q6, %f26[0]        \n"
            "vmla.f32   q13, q6, %f27[0]        \n"

            "vmla.f32   q8, q7, %f22[1]         \n"
            "vmla.f32   q9, q7, %f23[1]         \n"
            "vmla.f32   q10, q7, %f24[1]        \n"
            "vmla.f32   q11, q7, %f25[1]        \n"
            "vmla.f32   q12, q7, %f26[1]        \n"
            "vmla.f32   q13, q7, %f27[1]        \n"

            "vst1.f32   {d16-d17}, [%1]!        \n"
            "vst1.f32   {d18-d19}, [%2]!        \n"
            "vst1.f32   {d20-d21}, [%3]!        \n"
            "vst1.f32   {d22-d23}, [%4]!        \n"
            "vst1.f32   {d24-d25}, [%5]!        \n"
            "vst1.f32   {d26-d27}, [%6]!        \n"

            "subs   %0, #1                      \n"
            "bne    0b                          \n"
            : "=r"(nw),      // 0
              "=r"(c_ptr0),  // 1
              "=r"(c_ptr1),  // 2
              "=r"(c_ptr2),  // 3
              "=r"(c_ptr3),  // 4
              "=r"(c_ptr4),  // 5
              "=r"(c_ptr5),  // 6
              "=r"(b_ptr0),  // 7
              "=r"(b_ptr1),  // 8
              "=r"(b_ptr2),  // 9
              "=r"(b_ptr3)   // 10
            : "0"(nw),       // 11
              "1"(c_ptr0),   // 12
              "2"(c_ptr1),   // 13
              "3"(c_ptr2),   // 14
              "4"(c_ptr3),   // 15
              "5"(c_ptr4),   // 16
              "6"(c_ptr5),   // 17
              "7"(b_ptr0),   // 18
              "8"(b_ptr1),   // 19
              "9"(b_ptr2),   // 20
              "10"(b_ptr3),  // 21
              "w"(a0),       // 22
              "w"(a1),       // 23
              "w"(a2),       // 24
              "w"(a3),       // 25
              "w"(a4),       // 26
              "w"(a5)        // 27
            : "cc", "memory", "q6", "q7", "q8", "q9", "q10", "q11", "q12",
              "q13", "q14", "q15");

        w = (width >> 2) << 2;
      }
#endif
      if (w < width) {
        const float *b_ptr = B + (k * stride_b + w);
        float *c_ptr = C + (h * stride_c + w);
        GemmBlock(a_ptr, b_ptr, reg_height_tile, reg_K_tile, width - w,
                  stride_a, stride_b, stride_c, c_ptr);
      }
    }
    if (k < K) {
      const float *a_ptr = A + (h * stride_a + k);
      const float *b_ptr = B + k * stride_b;
      float *c_ptr = C + h * stride_c;
      GemmBlock(a_ptr, b_ptr, reg_height_tile, K - k, width, stride_a, stride_b,
                stride_c, c_ptr);
    }
  }
  if (h < height) {
    index_t remain_h = height - h;

    auto gemm_fn = Gemm184;
    switch (remain_h) {
      case 1:
      #if defined(__aarch64__)
        gemm_fn = Gemm184;
      #else
        gemm_fn = Gemm144;
      #endif
        break;
      case 2:
      #if defined(__aarch64__)
        gemm_fn = Gemm284;
      #else
        gemm_fn = Gemm244;
      #endif
        break;
      case 3:
      #if defined(__aarch64__)
        gemm_fn = Gemm384;
      #else
        gemm_fn = Gemm344;
      #endif
        break;
      case 4:
      #if defined(__aarch64__)
        gemm_fn = Gemm484;
      #else
        gemm_fn = Gemm444;
      #endif
        break;
      case 5:
      #if defined(__aarch64__)
        gemm_fn = Gemm584;
      #else
        gemm_fn = Gemm544;
      #endif
        break;
      case 6:
      #if defined(__aarch64__)
        gemm_fn = Gemm684;
      #else
        LOG(FATAL) << "remain_h should < 6";
      #endif
        break;
      case 7:
      #if defined(__aarch64__)
        gemm_fn = Gemm784;
      #else
        LOG(FATAL) << "remain_h should < 6";
      #endif
        break;
      default:
        LOG(FATAL) << "remain_h should < 8";
    }

    for (k = 0; k < K - reg_K_tile; k += reg_K_tile) {
      const float *a_ptr = A + (h * stride_a + k);
      index_t w;
      for (w = 0; w + 3 < width; w += 4) {
        const float *b_ptr = B + (k * stride_b + w);
        float *c_ptr = C + (h * stride_c + w);
        gemm_fn(a_ptr, b_ptr, stride_a, stride_b, stride_c, c_ptr);
      }
      if (w < width) {
        const float *b_ptr = B + (k * stride_b + w);
        float *c_ptr = C + (h * stride_c + w);
        GemmBlock(a_ptr, b_ptr, remain_h, reg_K_tile, width - w, stride_a,
                  stride_b, stride_c, c_ptr);
      }
    }
    if (k < K) {
      const float *a_ptr = A + (h * stride_a + k);
      const float *b_ptr = B + k * stride_b;
      float *c_ptr = C + h * stride_c;
      GemmBlock(a_ptr, b_ptr, remain_h, K - k, width, stride_a, stride_b,
                stride_c, c_ptr);
    }
  }
#else
  GemmBlock(A, B, height, K, width, stride_a, stride_b, stride_c, C);
#endif  // MACE_ENABLE_NEON
}

}  // namespace

void Transpose(const float *src,
               index_t height,
               index_t width,
               index_t stride_w,
               float *dst) {
  index_t tile_size = height > 512 || width > 512 ? 64 : 32;
  for (index_t i = 0; i < height; i += tile_size) {
    for (index_t j = 0; j < width; j += tile_size) {
      index_t end_i = std::min(i + tile_size, height);
      index_t end_j = std::min(j + tile_size, width);
      for (index_t tile_i = i; tile_i < end_i; ++tile_i) {
        for (index_t tile_j = j; tile_j < end_j; ++tile_j) {
          dst[tile_j * height + tile_i] = src[tile_i * stride_w + tile_j];
        }
      }
    }
  }
}

// A: height x K, B: K x width, C: height x width
void Gemm(const float *A,
          const float *B,
          const index_t batch,
          const index_t height,
          const index_t K,
          const index_t width,
          float *C,
          const bool transpose_a,
          const bool transpose_b) {
  if (width == 1 && !transpose_a) {
    for (index_t b = 0; b < batch; ++b) {
      Gemv(A + b * height * K, B + b * K, 1, K, height, C + b * height);
    }
    return;
  }
  memset(C, 0, sizeof(float) * batch * height * width);

  std::vector<index_t> block_size_dims {height, width, K};
  index_t thread_count = MaceOpenMPThreadCount;
  MACE_CHECK(thread_count >= 1, "thread should be ge 1");
  // TODO(liyin): apply gcd ?
  if (height % thread_count == 0) {
    block_size_dims[0] = height / thread_count;
  } else if (thread_count == 4 && (height & 1) == 0 && (width & 1) == 0) {
    block_size_dims[0] = height >> 1;
    block_size_dims[1] = width >> 1;
  } else if (width % thread_count == 0) {
    block_size_dims[1] = width / thread_count;
  } else {
    if (height >= thread_count) {
      block_size_dims[0] = height / thread_count;
    } else {
      thread_count = std::min(thread_count, height * width);
      index_t thread_h = height;
      index_t thread_w = RoundUpDiv(thread_count, thread_h);
      block_size_dims[0] = 1;
      block_size_dims[1] = std::max(static_cast<index_t>(1), width / thread_w);
    }
  }

  const index_t block_tile[3] = {height / block_size_dims[0],
                                 width / block_size_dims[1],
                                 K / block_size_dims[2]};
  block_size_dims[0] = height / block_tile[0];
  block_size_dims[1] = width / block_tile[1];
  block_size_dims[2] = K / block_tile[2];

  const index_t remain[3] = {height % block_tile[0],
                             width % block_tile[1],
                             K % block_tile[2]};


#pragma omp parallel for collapse(3)
  for (index_t n = 0; n < batch; ++n) {
    for (index_t bh = 0; bh < block_tile[0]; ++bh) {
      for (index_t bw = 0; bw < block_tile[1]; ++bw) {
        const index_t remain_height = remain[0];
        const index_t remain_width = remain[1];
        const index_t remain_k = remain[2];

        const index_t block_size_height = block_size_dims[0];
        const index_t block_size_width = block_size_dims[1];
        const index_t block_size_k = block_size_dims[2];

        const index_t this_block_size_height =
            block_size_height + (bh < remain_height ? 1 : 0);
        const index_t this_block_size_width =
            block_size_width + (bw < remain_width ? 1 : 0);

        const float *a_base = A + n * height * K;
        const float *b_base = B + n * K * width;
        float *c_base = C + n * height * width;

        const index_t ih_begin =
            bh * block_size_height + (bh < remain_height ? bh : remain_height);
        const index_t
            ih_end = std::min(height, ih_begin + this_block_size_height);
        const index_t iw_begin =
            bw * block_size_width + (bw < remain_width ? bw : remain_width);
        const index_t
            iw_end = std::min(width, iw_begin + this_block_size_width);

        for (index_t bk = 0; bk < block_tile[2]; ++bk) {
          const index_t
              this_block_size_k = block_size_k + (bk < remain_k ? 1 : 0);

          const index_t
              ik_begin = bk * block_size_k + (bk < remain_k ? bk : remain_k);
          const index_t ik_end = std::min(K, ik_begin + this_block_size_k);

          Tensor trans_a(GetCPUAllocator(), DataType::DT_FLOAT);
          Tensor trans_b(GetCPUAllocator(), DataType::DT_FLOAT);
          const float *real_a = nullptr;
          const float *real_b = nullptr;
          float *real_c = c_base + (ih_begin * width + iw_begin);
          index_t stride_a;
          index_t stride_b;
          index_t stride_c = width;

          if (transpose_a) {
            trans_a.Resize({this_block_size_height, this_block_size_k});
            float *trans_a_data = trans_a.mutable_data<float>();
            // A[K, H] -> A[H, K]
            Transpose(a_base + (ik_begin * height + ih_begin),
                      ik_end - ik_begin, ih_end - ih_begin, height,
                      trans_a_data);
            real_a = trans_a_data;
            stride_a = ik_end - ik_begin;
          } else {
            real_a = a_base + (ih_begin * K + ik_begin);
            stride_a = K;
          }

          if (transpose_b) {
            trans_b.Resize({this_block_size_k, this_block_size_width});
            float *trans_b_data = trans_b.mutable_data<float>();
            // B[W, K] -> B[K, W]
            Transpose(b_base + (iw_begin * K + ik_begin), iw_end - iw_begin,
                      ik_end - ik_begin, K, trans_b_data);
            real_b = trans_b_data;
            stride_b = iw_end - iw_begin;
          } else {
            real_b = b_base + (ik_begin * width + iw_begin);
            stride_b = width;
          }

          // inside block:
          // calculate C[bh, bw] += A[bh, bk] * B[bk, bw] for one k
          GemmTile(real_a, real_b, ih_end - ih_begin, ik_end - ik_begin,
                   iw_end - iw_begin, stride_a, stride_b, stride_c, real_c);
        }  // bk
      }    // bw
    }      // bh
  }        // n
}

// A: height x K, B: K x width, C: height x width
void GemmRef(const float *A,
             const float *B,
             const index_t batch,
             const index_t height,
             const index_t K,
             const index_t width,
             float *C,
             const bool transpose_a,
             const bool transpose_b) {
  memset(C, 0, sizeof(float) * batch * height * width);

  Tensor trans_a(GetCPUAllocator(), DataType::DT_FLOAT);
  Tensor trans_b(GetCPUAllocator(), DataType::DT_FLOAT);
  float *trans_a_data = nullptr;
  float *trans_b_data = nullptr;
  if (transpose_a) {
    trans_a.Resize({height, K});
    trans_a_data = trans_a.mutable_data<float>();
  }
  if (transpose_b) {
    trans_b.Resize({K, width});
    trans_b_data = trans_b.mutable_data<float>();
  }

  for (index_t b = 0; b < batch; ++b) {
    const float *real_a = nullptr;
    const float *real_b = nullptr;
    float *real_c = C + b * height * width;
    if (transpose_a) {
      // A[K, H] -> A[H, K]
      Transpose(A + b * height * K, K, height, height, trans_a_data);
      real_a = trans_a_data;
    } else {
      real_a = A + b * height * K;
    }
    if (transpose_b) {
      // B[W, K] -> B[K, W]
      Transpose(B + b * width * K, width, K, K, trans_b_data);
      real_b = trans_b_data;
    } else {
      real_b = B + b * width * K;
    }

    for (index_t i = 0; i < height; ++i) {
      for (index_t j = 0; j < width; ++j) {
        for (index_t k = 0; k < K; ++k) {
          real_c[i * width + j] += real_a[i * K + k] * real_b[k * width + j];
        }
      }
    }
  }
}

void GemvRef(const float *m_ptr,
             const float *v_ptr,
             const index_t batch,
             const index_t width,
             const index_t height,
             float *out_ptr) {
  memset(out_ptr, 0, batch * height * sizeof(float));
#pragma omp parallel for collapse(2)
  for (int b = 0; b < batch; ++b) {
    for (int h = 0; h < height; ++h) {
      for (int w = 0; w < width; ++w) {
        out_ptr[b * height + h] += v_ptr[b * width + w] * m_ptr[h * width + w];
      }
    }
  }
}

void Gemv(const float *m_ptr,
          const float *v_ptr,
          const index_t batch,
          const index_t width,
          const index_t height,
          float *out_ptr) {
#if defined(MACE_ENABLE_NEON)

#pragma omp parallel for collapse(2)
  for (index_t b = 0; b < batch; ++b) {
    for (index_t h = 0; h < height; ++h) {
      const float *m_ptr0 = m_ptr + h * width;
      const float *v_ptr0 = v_ptr + b * width;
      float *out_ptr0 = out_ptr + b * height + h;

      float32x4_t vm0, vm1, vm2, vm3;
      float32x4_t vv0, vv1, vv2, vv3;
      float32x4_t vsum0 = vdupq_n_f32(0.f);
      float32x4_t vsum1 = vdupq_n_f32(0.f);
      float32x4_t vsum2 = vdupq_n_f32(0.f);
      float32x4_t vsum3 = vdupq_n_f32(0.f);

      index_t w;
      for (w = 0; w + 15 < width; w += 16) {
        vm0 = vld1q_f32(m_ptr0);
        vv0 = vld1q_f32(v_ptr0);
        vm1 = vld1q_f32(m_ptr0 + 4);
        vv1 = vld1q_f32(v_ptr0 + 4);
        vm2 = vld1q_f32(m_ptr0 + 8);
        vv2 = vld1q_f32(v_ptr0 + 8);
        vm3 = vld1q_f32(m_ptr0 + 12);
        vv3 = vld1q_f32(v_ptr0 + 12);

        vsum0 = vmlaq_f32(vsum0, vm0, vv0);
        vsum1 = vmlaq_f32(vsum1, vm1, vv1);
        vsum2 = vmlaq_f32(vsum2, vm2, vv2);
        vsum3 = vmlaq_f32(vsum3, vm3, vv3);

        m_ptr0 += 16;
        v_ptr0 += 16;
      }

      for (; w + 7 < width; w += 8) {
        vm0 = vld1q_f32(m_ptr0);
        vv0 = vld1q_f32(v_ptr0);
        vm1 = vld1q_f32(m_ptr0 + 4);
        vv1 = vld1q_f32(v_ptr0 + 4);

        vsum0 = vmlaq_f32(vsum0, vm0, vv0);
        vsum1 = vmlaq_f32(vsum1, vm1, vv1);

        m_ptr0 += 8;
        v_ptr0 += 8;
      }

      for (; w + 3 < width; w += 4) {
        vm0 = vld1q_f32(m_ptr0);
        vv0 = vld1q_f32(v_ptr0);
        vsum0 = vmlaq_f32(vsum0, vm0, vv0);

        m_ptr0 += 4;
        v_ptr0 += 4;
      }
      vsum0 += vsum1;
      vsum2 += vsum3;
      vsum0 += vsum2;
      float sum0 = vaddvq_f32(vsum0);

      // handle remaining w
      for (; w < width; ++w) {
        sum0 += m_ptr0[0] * v_ptr0[0];
        m_ptr0++;
        v_ptr0++;
      }
      *out_ptr0++ = sum0;
    }    // h
  }      // b
#else
  GemvRef(m_ptr, v_ptr, batch, width, height, out_ptr);
#endif
}

}  // namespace ops
}  // namespace mace
