// File: kernel.cpp
#include "loop_135.h"
#include <stdint.h>
#include <arm_sve.h>
#define restrict __restrict
#define SC_SVE_ATTR
#define FOR_COND(P, S, I, N) svptest_first(svptrue_b##S(), P = svwhilelt_b##S(I, N))
#define FOR_LOOP(T, I, M, N, P, S, W) for (T I = M; FOR_COND(P, S, I, N); I += svcnt##W())
#define FOR_LOOP_8(T, I, M, N, P)  FOR_LOOP(T, I, M, N, P, 8, b)
#define FOR_LOOP_16(T, I, M, N, P) FOR_LOOP(T, I, M, N, P, 16, h)
#define FOR_LOOP_32(T, I, M, N, P) FOR_LOOP(T, I, M, N, P, 32, w)
#define FOR_LOOP_64(T, I, M, N, P) FOR_LOOP(T, I, M, N, P, 64, d)
static inline uint32_t get_sve_vl(void) {
  uint64_t vl = 0; asm volatile("incb %[vl], all, mul #8" : [vl] "+r"(vl) : :);
  return (uint32_t)vl;
}
static inline uint32_t get_vl(void) { return get_sve_vl(); }

extern "C" void inner_loop_135(struct loop_135_data *data) {
  const uint64_t m = data->m;
  const uint64_t n = data->n;
  const uint64_t k = data->k;
  const int8_t* const a = data->a;
  const int8_t* const b = data->b;
  int32_t* const c = data->c;
  const svbool_t pg = svptrue_b8();

  // M loop
  for(uint64_t m_idx = 0; m_idx < m; m_idx += 12) {

    // N loop
    for(uint64_t n_idx = 0; n_idx < n; n_idx += 2 * svcntw()) {
      // 24 Accumulators
      svint32_t acc0_0 = svdup_s32(0);
      svint32_t acc0_1 = svdup_s32(0);
      svint32_t acc0_2 = svdup_s32(0);
      svint32_t acc0_3 = svdup_s32(0);
      svint32_t acc1_0 = svdup_s32(0);
      svint32_t acc1_1 = svdup_s32(0);
      svint32_t acc1_2 = svdup_s32(0);
      svint32_t acc1_3 = svdup_s32(0);
      svint32_t acc2_0 = svdup_s32(0);
      svint32_t acc2_1 = svdup_s32(0);
      svint32_t acc2_2 = svdup_s32(0);
      svint32_t acc2_3 = svdup_s32(0);
      svint32_t acc3_0 = svdup_s32(0);
      svint32_t acc3_1 = svdup_s32(0);
      svint32_t acc3_2 = svdup_s32(0);
      svint32_t acc3_3 = svdup_s32(0);
      svint32_t acc4_0 = svdup_s32(0);
      svint32_t acc4_1 = svdup_s32(0);
      svint32_t acc4_2 = svdup_s32(0);
      svint32_t acc4_3 = svdup_s32(0);
      svint32_t acc5_0 = svdup_s32(0);
      svint32_t acc5_1 = svdup_s32(0);
      svint32_t acc5_2 = svdup_s32(0);
      svint32_t acc5_3 = svdup_s32(0);

      // K loop
      for(uint64_t k_idx = 0; k_idx < k; k_idx += 8) {
        const int8_t* const a_ptr = a + k_idx * m + 8 * m_idx;
        const int8_t* const b_ptr = b + k_idx * n + 8 * n_idx;

	svint8_t vb_0 = svld1_vnum_s8(pg, b_ptr, 0);
	svint8_t vb_1 = svld1_vnum_s8(pg, b_ptr, 1);
	svint8_t vb_2 = svld1_vnum_s8(pg, b_ptr, 2);
	svint8_t vb_3 = svld1_vnum_s8(pg, b_ptr, 3);

        svint8_t va_0 = svld1rq_s8(pg, a_ptr);
	acc0_0 = svmmla_s32(acc0_0, va_0, vb_0);
	acc0_1 = svmmla_s32(acc0_1, va_0, vb_1);
	acc0_2 = svmmla_s32(acc0_2, va_0, vb_2);
	acc0_3 = svmmla_s32(acc0_3, va_0, vb_3);

        svint8_t va_1 = svld1rq_s8(pg, a_ptr + 16);
	acc1_0 = svmmla_s32(acc1_0, va_1, vb_0);
	acc1_1 = svmmla_s32(acc1_1, va_1, vb_1);
	acc1_2 = svmmla_s32(acc1_2, va_1, vb_2);
	acc1_3 = svmmla_s32(acc1_3, va_1, vb_3);

        svint8_t va_2 = svld1rq_s8(pg, a_ptr + 32);
	acc2_0 = svmmla_s32(acc2_0, va_2, vb_0);
	acc2_1 = svmmla_s32(acc2_1, va_2, vb_1);
	acc2_2 = svmmla_s32(acc2_2, va_2, vb_2);
	acc2_3 = svmmla_s32(acc2_3, va_2, vb_3);

        svint8_t va_3 = svld1rq_s8(pg, a_ptr + 48);
	acc3_0 = svmmla_s32(acc3_0, va_3, vb_0);
	acc3_1 = svmmla_s32(acc3_1, va_3, vb_1);
	acc3_2 = svmmla_s32(acc3_2, va_3, vb_2);
	acc3_3 = svmmla_s32(acc3_3, va_3, vb_3);

        svint8_t va_4 = svld1rq_s8(pg, a_ptr + 64);
	acc4_0 = svmmla_s32(acc4_0, va_4, vb_0);
	acc4_1 = svmmla_s32(acc4_1, va_4, vb_1);
	acc4_2 = svmmla_s32(acc4_2, va_4, vb_2);
	acc4_3 = svmmla_s32(acc4_3, va_4, vb_3);

        svint8_t va_5 = svld1rq_s8(pg, a_ptr + 80);
	acc5_0 = svmmla_s32(acc5_0, va_5, vb_0);
	acc5_1 = svmmla_s32(acc5_1, va_5, vb_1);
	acc5_2 = svmmla_s32(acc5_2, va_5, vb_2);
	acc5_3 = svmmla_s32(acc5_3, va_5, vb_3);
      }

      // Store
      int32_t* const c0ptr = c + m_idx * n + n_idx;
      int32_t* const c1ptr = c0ptr + svcntw();

      svint64_t res_0l = svuzp1_s64(svreinterpret_s64_s32(acc0_0), svreinterpret_s64_s32(acc0_1));
      svint64_t res_0h = svuzp1_s64(svreinterpret_s64_s32(acc0_2), svreinterpret_s64_s32(acc0_3));
      svint64_t res_1l = svuzp2_s64(svreinterpret_s64_s32(acc0_0), svreinterpret_s64_s32(acc0_1));
      svint64_t res_1h = svuzp2_s64(svreinterpret_s64_s32(acc0_2), svreinterpret_s64_s32(acc0_3));
      svint64_t res_2l = svuzp1_s64(svreinterpret_s64_s32(acc1_0), svreinterpret_s64_s32(acc1_1));
      svint64_t res_2h = svuzp1_s64(svreinterpret_s64_s32(acc1_2), svreinterpret_s64_s32(acc1_3));
      svint64_t res_3l = svuzp2_s64(svreinterpret_s64_s32(acc1_0), svreinterpret_s64_s32(acc1_1));
      svint64_t res_3h = svuzp2_s64(svreinterpret_s64_s32(acc1_2), svreinterpret_s64_s32(acc1_3));

      svst1_s32(pg, c0ptr, svreinterpret_s32_s64(res_0l));
      svst1_s32(pg, c1ptr, svreinterpret_s32_s64(res_0h));
      svst1_s32(pg, c0ptr + n, svreinterpret_s32_s64(res_1l));
      svst1_s32(pg, c1ptr + n, svreinterpret_s32_s64(res_1h));
      svst1_s32(pg, c0ptr + 2 * n, svreinterpret_s32_s64(res_2l));
      svst1_s32(pg, c1ptr + 2 * n, svreinterpret_s32_s64(res_2h));
      svst1_s32(pg, c0ptr + 3 * n, svreinterpret_s32_s64(res_3l));
      svst1_s32(pg, c1ptr + 3 * n, svreinterpret_s32_s64(res_3h));

      svint64_t res_4l = svuzp1_s64(svreinterpret_s64_s32(acc2_0), svreinterpret_s64_s32(acc2_1));
      svint64_t res_4h = svuzp1_s64(svreinterpret_s64_s32(acc2_2), svreinterpret_s64_s32(acc2_3));
      svint64_t res_5l = svuzp2_s64(svreinterpret_s64_s32(acc2_0), svreinterpret_s64_s32(acc2_1));
      svint64_t res_5h = svuzp2_s64(svreinterpret_s64_s32(acc2_2), svreinterpret_s64_s32(acc2_3));
      svint64_t res_6l = svuzp1_s64(svreinterpret_s64_s32(acc3_0), svreinterpret_s64_s32(acc3_1));
      svint64_t res_6h = svuzp1_s64(svreinterpret_s64_s32(acc3_2), svreinterpret_s64_s32(acc3_3));
      svint64_t res_7l = svuzp2_s64(svreinterpret_s64_s32(acc3_0), svreinterpret_s64_s32(acc3_1));
      svint64_t res_7h = svuzp2_s64(svreinterpret_s64_s32(acc3_2), svreinterpret_s64_s32(acc3_3));

      svst1_s32(pg, c0ptr + 4 * n, svreinterpret_s32_s64(res_4l));
      svst1_s32(pg, c1ptr + 4 * n, svreinterpret_s32_s64(res_4h));
      svst1_s32(pg, c0ptr + 5 * n, svreinterpret_s32_s64(res_5l));
      svst1_s32(pg, c1ptr + 5 * n, svreinterpret_s32_s64(res_5h));
      svst1_s32(pg, c0ptr + 6 * n, svreinterpret_s32_s64(res_6l));
      svst1_s32(pg, c1ptr + 6 * n, svreinterpret_s32_s64(res_6h));
      svst1_s32(pg, c0ptr + 7 * n, svreinterpret_s32_s64(res_7l));
      svst1_s32(pg, c1ptr + 7 * n, svreinterpret_s32_s64(res_7h));

      svint64_t res_8l = svuzp1_s64(svreinterpret_s64_s32(acc4_0), svreinterpret_s64_s32(acc4_1));
      svint64_t res_8h = svuzp1_s64(svreinterpret_s64_s32(acc4_2), svreinterpret_s64_s32(acc4_3));
      svint64_t res_9l = svuzp2_s64(svreinterpret_s64_s32(acc4_0), svreinterpret_s64_s32(acc4_1));
      svint64_t res_9h = svuzp2_s64(svreinterpret_s64_s32(acc4_2), svreinterpret_s64_s32(acc4_3));
      svint64_t res_10l = svuzp1_s64(svreinterpret_s64_s32(acc5_0), svreinterpret_s64_s32(acc5_1));
      svint64_t res_10h = svuzp1_s64(svreinterpret_s64_s32(acc5_2), svreinterpret_s64_s32(acc5_3));
      svint64_t res_11l = svuzp2_s64(svreinterpret_s64_s32(acc5_0), svreinterpret_s64_s32(acc5_1));
      svint64_t res_11h = svuzp2_s64(svreinterpret_s64_s32(acc5_2), svreinterpret_s64_s32(acc5_3));

      svst1_s32(pg, c0ptr + 8 * n, svreinterpret_s32_s64(res_8l));
      svst1_s32(pg, c1ptr + 8 * n, svreinterpret_s32_s64(res_8h));
      svst1_s32(pg, c0ptr + 9 * n, svreinterpret_s32_s64(res_9l));
      svst1_s32(pg, c1ptr + 9 * n, svreinterpret_s32_s64(res_9h));
      svst1_s32(pg, c0ptr + 10 * n, svreinterpret_s32_s64(res_10l));
      svst1_s32(pg, c1ptr + 10 * n, svreinterpret_s32_s64(res_10h));
      svst1_s32(pg, c0ptr + 11 * n, svreinterpret_s32_s64(res_11l));
      svst1_s32(pg, c1ptr + 11 * n, svreinterpret_s32_s64(res_11h));
    }
  }

}