// File: kernel.cpp
#include "loop_130.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_130(struct loop_130_data *data) {
  uint64_t m = data->m;
  uint64_t n = data->n;
  uint64_t k = data->k;
  float32_t *a = data->a;
  float32_t *b = data->b;
  float32_t *c = data->c;

  float32_t *ptr_a, *ptr_b, *ptr_c;
  float32_t *cnd_k = &a[m * k];

  uint64_t m_idx, n_idx;
  svbool_t p_all = svptrue_b32();

  svfloat32x4_t acc_0, acc_1, acc_2, acc_3, acc_4, acc_5;
  svfloat32_t lda_0, lda_1, lda_2, lda_3, lda_4, lda_5;
  svfloat32x4_t ldb;

#define ZERO svdup_f32(0.0f)
#define ZERO_QUAD(x) acc_##x = svcreate4(ZERO, ZERO, ZERO, ZERO)

#define LOADA(y) svld1rq(p_all, &ptr_a[4 * y])

#define GETC(x, y) svreinterpret_u64(svget4(acc_##y, x))
#define UNZIP(x, y, z) \
  svreinterpret_f32(svuzp##z(GETC(2 * x + 0, y), GETC(2 * x + 1, y)))

#if defined(__ARM_FEATURE_SVE2p1)
  svcount_t c_all = svptrue_c32();
#define LOADB_QUAD ldb = svld1_x4(c_all, ptr_b)
#define UNZIP_PAIR(y, z) svcreate2(UNZIP(0, y, z), UNZIP(1, y, z))
#define STORE_PAIR(y, z) \
  svst1_vnum(c_all, &ptr_c[n * (2 * y + z - 1)], 0, UNZIP_PAIR(y, z))
#else
#define LOADB(x) svld1_vnum(p_all, ptr_b, x)
#define LOADB_QUAD ldb = svcreate4(LOADB(0), LOADB(1), LOADB(2), LOADB(3))
#define STORE(x, y, z) \
  svst1_vnum(p_all, &ptr_c[n * (2 * y + z - 1)], x, UNZIP(x, y, z));
#define STORE_PAIR(y, z) STORE(0, y, z) STORE(1, y, z)
#endif

  for (m_idx = 0; m_idx < m; m_idx += 12) {
    for (n_idx = 0; n_idx < n; n_idx += svcntw() * 2) {
      ZERO_QUAD(0);
      ZERO_QUAD(1);
      ZERO_QUAD(2);
      ZERO_QUAD(3);
      ZERO_QUAD(4);
      ZERO_QUAD(5);

      ptr_a = &a[m_idx * 2];
      ptr_b = &b[n_idx * 2];

      while (ptr_a < cnd_k) {
        lda_0 = LOADA(0);
        lda_1 = LOADA(1);
        lda_2 = LOADA(2);
        lda_3 = LOADA(3);
        lda_4 = LOADA(4);
        lda_5 = LOADA(5);

        LOADB_QUAD;

        acc_0 = svcreate4(svmmla(svget4(acc_0, 0), lda_0, svget4(ldb, 0)),
                          svmmla(svget4(acc_0, 1), lda_0, svget4(ldb, 1)),
                          svmmla(svget4(acc_0, 2), lda_0, svget4(ldb, 2)),
                          svmmla(svget4(acc_0, 3), lda_0, svget4(ldb, 3)));
        acc_1 = svcreate4(svmmla(svget4(acc_1, 0), lda_1, svget4(ldb, 0)),
                          svmmla(svget4(acc_1, 1), lda_1, svget4(ldb, 1)),
                          svmmla(svget4(acc_1, 2), lda_1, svget4(ldb, 2)),
                          svmmla(svget4(acc_1, 3), lda_1, svget4(ldb, 3)));
        acc_2 = svcreate4(svmmla(svget4(acc_2, 0), lda_2, svget4(ldb, 0)),
                          svmmla(svget4(acc_2, 1), lda_2, svget4(ldb, 1)),
                          svmmla(svget4(acc_2, 2), lda_2, svget4(ldb, 2)),
                          svmmla(svget4(acc_2, 3), lda_2, svget4(ldb, 3)));
        acc_3 = svcreate4(svmmla(svget4(acc_3, 0), lda_3, svget4(ldb, 0)),
                          svmmla(svget4(acc_3, 1), lda_3, svget4(ldb, 1)),
                          svmmla(svget4(acc_3, 2), lda_3, svget4(ldb, 2)),
                          svmmla(svget4(acc_3, 3), lda_3, svget4(ldb, 3)));
        acc_4 = svcreate4(svmmla(svget4(acc_4, 0), lda_4, svget4(ldb, 0)),
                          svmmla(svget4(acc_4, 1), lda_4, svget4(ldb, 1)),
                          svmmla(svget4(acc_4, 2), lda_4, svget4(ldb, 2)),
                          svmmla(svget4(acc_4, 3), lda_4, svget4(ldb, 3)));
        acc_5 = svcreate4(svmmla(svget4(acc_5, 0), lda_5, svget4(ldb, 0)),
                          svmmla(svget4(acc_5, 1), lda_5, svget4(ldb, 1)),
                          svmmla(svget4(acc_5, 2), lda_5, svget4(ldb, 2)),
                          svmmla(svget4(acc_5, 3), lda_5, svget4(ldb, 3)));

        ptr_a += m * 2;
        ptr_b += n * 2;
      }

      ptr_c = &c[n_idx];
      STORE_PAIR(0, 1);
      STORE_PAIR(0, 2);
      STORE_PAIR(1, 1);
      STORE_PAIR(1, 2);
      STORE_PAIR(2, 1);
      STORE_PAIR(2, 2);
      STORE_PAIR(3, 1);
      STORE_PAIR(3, 2);
      STORE_PAIR(4, 1);
      STORE_PAIR(4, 2);
      STORE_PAIR(5, 1);
      STORE_PAIR(5, 2);
    }
    c += n * 12;
  }
}