#include "mat_mul.h"

#include <cstdlib>
#include <cstdio>
#include <pthread.h>
#include <immintrin.h>

static float *A, *B, *C;
static int M, N, K;
static int num_threads;
#define BLOCK_SIZE (48)
#define VEC_NUM (1)
#define VEC_SIZE (64)
#define VECDATA_SIZE (VEC_SIZE>>2)
#define MIN(a,b) ((a>b)?b:a)

static void* mat_mul_thread(void *data) {
  // TODO: parallelize & optimize matrix multiplication
  int i,j,k;
  int kk;
  int pid = * (int *) data;

  int slice = (M + num_threads - 1) / num_threads;
  int si = pid * slice;
  int ei = (pid + 1) * slice;
  if (si >= M) return NULL;
  if (ei >= M) ei = M;

  for (kk = 0; kk < K; kk+=BLOCK_SIZE) {
    for (i = si; i < ei; i+=1) {
      for (k = kk; k < MIN(K,kk+BLOCK_SIZE); ++k) {
        float Aik = A[i * K + k];
        unsigned long spC, epC;
        __m512 va = _mm512_set1_ps(Aik);
        __m512 vb[VEC_NUM],vc[VEC_NUM];
        spC = ((unsigned long)&C[i*N] + (VEC_SIZE*VEC_NUM-1)) & (VEC_SIZE*VEC_NUM-1);
        epC = (unsigned long)&C[(i+1)*N] & (VEC_SIZE*VEC_NUM-1);
        
        for (j = 0; (unsigned long)&C[i * N + j] < spC; j=j+1) {
          C[i * N + j] += Aik * B[k * N + j];
        }
        for (; (unsigned long)&C[i * N + j] < epC; j+=VECDATA_SIZE*VEC_NUM) {
          vb[0] = _mm512_load_ps(&B[k*N+j]);
          vc[0] = _mm512_load_ps(&C[i*N+j]);
          vc[0] = _mm512_fmadd_ps(va,vb[0],vc[0]);
          _mm512_store_ps(&C[i*N+j],vc[0]);
          /*
          vb[1] = _mm512_load_ps(&B[k*N+j+VECDATA_SIZE]);
          vc[1] = _mm512_load_ps(&C[i*N+j+VECDATA_SIZE]);
          vc[1] = _mm512_fmadd_ps(va,vb[1],vc[1]);
          _mm512_store_ps(&C[i*N+j+VECDATA_SIZE],vc[1]);*/
        }
        for (; j < N; j=j+1) {
          C[i * N + j] += Aik * B[k * N + j];
        }
      }
    }
  }

  return NULL;
}

void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K, int _num_threads) {
  A = _A, B = _B, C = _C;
  M = _M, N = _N, K = _K;
  num_threads = _num_threads;

  // TODO: create '_num_threads' pthreads
  pthread_t thread[_num_threads];
  int tid[_num_threads];
  for (int i = 0; i < _num_threads; i++) {
    tid[i] = i;
    pthread_create(&thread[i], NULL, mat_mul_thread, &tid[i]);
  }

  for (int i = 0; i < _num_threads; i++)
    pthread_join(thread[i], NULL);
}