// File: kernel.cpp
#include "loop_219.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(); }

svuint8x4_t svzip4(svuint8_t a, svuint8_t b, svuint8_t c, svuint8_t d)

{
  svuint8_t e = svzip1(a, c);
  svuint8_t f = svzip2(a, c);
  svuint8_t g = svzip1(b, d);
  svuint8_t h = svzip2(b, d);
  return svcreate4(svzip1(e, g), svzip2(e, g), svzip1(f, h), svzip2(f, h));
}

extern "C" void inner_loop_219(struct loop_219_data *data)

{
  uint64_t m = data->m;
  uint64_t n = data->n;
  uint8_t *a = data->a;
  uint8_t *b = data->b;
  uint32_t *c = data->c;

  uint8_t *a_ptr;
  uint32_t *c_ptr;

  uint64_t m_idx, n_idx;
  svbool_t p_ldx = svptrue_pat_b8(SV_VL8);

  svuint32x4_t acc_0, acc_1;
  svuint8x4_t zip_0, zip_1;
  svuint8x2_t lda_0, lda_1, lda_2, lda_3;
  svuint8_t ldx;

#define ZERO svdup_u32(0)
#define ZERO_QUAD(q) acc_##q = svcreate4(ZERO, ZERO, ZERO, ZERO)

#define GETA(q, p) svget2(lda_##q, p)
#define GETB(q, p) svget4(acc_##q, p)

#define ZIP4_QUAD(i) zip_##i = \
  svzip4(GETA(0, i), GETA(1, i), GETA(2, i), GETA(3, i))

#define UDOT(l, q, p) svdot_lane(GETB(q, p), svget4(zip_##q, p), ldx, l)
#define UDOT_LANE(l, q) acc_##q = \
  svcreate4(UDOT(l, q, 0), UDOT(l, q, 1), UDOT(l, q, 2), UDOT(l, q, 3))

#if defined(__ARM_FEATURE_SVE2p1)
  svcount_t c_all = svptrue_c8();
#define LOAD_PAIR(i) lda_##i = svld1_x2(c_all, &a_ptr[m * i])
#define STORE_QUAD(q) svst1_vnum(c_all, c_ptr, q * 4, acc_##q)
#else
  svbool_t p_all = svptrue_b8();
#define LOAD(i, p) svld1_vnum(p_all, &a_ptr[m * i], p)
#define LOAD_PAIR(i) lda_##i = svcreate2(LOAD(i, 0), LOAD(i, 1))
#define STORE(q, p) svst1_vnum(p_all, c_ptr, q * 4 + p, GETB(q, p));
#define STORE_QUAD(q) STORE(q, 0) STORE(q, 1) STORE(q, 2) STORE(q, 3)
#endif

  for (m_idx = 0; m_idx < m; m_idx += svcntb() * 2) {
    ZERO_QUAD(0);
    ZERO_QUAD(1);

    a_ptr = &a[m_idx];
    for (n_idx = 0; n_idx < n; n_idx += 8) {
      ldx = svld1rq(p_ldx, &b[n_idx]);
      LOAD_PAIR(0);
      LOAD_PAIR(1);
      LOAD_PAIR(2);
      LOAD_PAIR(3);
      ZIP4_QUAD(0);
      ZIP4_QUAD(1);
      UDOT_LANE(0, 0);
      UDOT_LANE(0, 1);
      a_ptr += 4 * m;
      LOAD_PAIR(0);
      LOAD_PAIR(1);
      LOAD_PAIR(2);
      LOAD_PAIR(3);
      ZIP4_QUAD(0);
      ZIP4_QUAD(1);
      UDOT_LANE(1, 0);
      UDOT_LANE(1, 1);
      a_ptr += 4 * m;
    }

    c_ptr = &c[m_idx];
    STORE_QUAD(0);
    STORE_QUAD(1);
  }
}