#include "convolution.h"
#include <mpi.h>
#include <stdio.h>
#include "util.h"
#include <immintrin.h>

static float *input, *output, *filter;
static int N, C, H, W;
static int K, R, S;
static int OH, OW;
static int pad;
static int dilation;
static int stride;
static int mpi_rank, mpi_world_size;


int num_threads = 120;
void convolution(
    float *_input, float *_output, float *_filter, 
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
  input = _input;
  output = _output;
  filter = _filter;

  int size[2];
  MPI_Request request;
  MPI_Status status;
  if ( mpi_world_size ==2) size[1] = _N/2; //TBC
  else size[1] = 0;
  size[0] = N - size[1];

  
  OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
  OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;
  
  if (mpi_rank == 0 && mpi_world_size ==2) 
  {
    MPI_Isend(&input[size[0]*C*H*W], size[1]*C*H*W, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &request);
    MPI_Isend(filter, K*C*R*S, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &request);
  }
  else if(mpi_world_size ==2)
  {
    alloc_tensor(&input, size[1], C, H, W);
    alloc_tensor(&output, size[1], K, OH, OW);
    alloc_tensor(&filter, K, C, R, S);
    MPI_Recv(input, size[1]*C*H*W, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &status);
    MPI_Recv(filter, K*C*R*S, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &status);
  }

 int performance;
 performance = (dilation ==1 && pad == 0 && stride ==1) & (R==16) & (S==16);
performance = 0;
//#pragma omp parallel for num_threads(num_threads) collapse(3) schedule(dynamic) 
//#pragma omp parallel for num_threads(num_threads) collapse(3) schedule(dynamic) 

float ii[16];
float ff[16];
__m512 i0;
__m512 f0;
__m512 s0;
float tmp;


#pragma omp parallel for num_threads(num_threads) collapse(3) schedule(dynamic) 
  for (int n = 0; n < size[mpi_rank]; ++n) {
    for (int k = 0; k < K; ++k) {
      for (int oh = 0; oh < OH; ++oh) {
        for (int ow = 0; ow < OW; ++ow) {
          float o = 0.f;
          for (int c = 0; c < C; ++c) {
            for (int r = 0; r < R; ++r) {
	      if(performance)
              {
                 //for (int s = 0; s < 16; ++s) {
                  int h = oh  + r ;
                  //int w = ow  + s ;
                 // if (h < 0 || h >= H || w < 0 || w >= W)
		 // {
			  //printf("continue........\n");
		//	  continue;
		 // }
//                  ii[0]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 0];
//                  ff[0]= filter[k * C * R * S + c * R * S + r * S + 0];
//                  ii[1]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 1];
//                  ff[1]= filter[k * C * R * S + c * R * S + r * S + 1];
//                  ii[2]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 2];
//                  ff[2]= filter[k * C * R * S + c * R * S + r * S + 2];
//                  ii[3]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 3];
//                  ff[3]= filter[k * C * R * S + c * R * S + r * S + 3];
//		  
//                  ii[4]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 4];
//                  ff[4]= filter[k * C * R * S + c * R * S + r * S + 4];
//                  ii[5]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 5];
//                  ff[5]= filter[k * C * R * S + c * R * S + r * S + 5];
//                  ii[6]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 6];
//                  ff[6]= filter[k * C * R * S + c * R * S + r * S + 6];
//                  ii[7]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 7];
//                  ff[7]= filter[k * C * R * S + c * R * S + r * S + 7];
//
//
//                  ii[8]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 8];
//                  ff[8]= filter[k * C * R * S + c * R * S + r * S + 8];
//                  ii[9]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 9];
//                  ff[9]= filter[k * C * R * S + c * R * S + r * S + 9];
//                  ii[10]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 10];
//                  ff[10]= filter[k * C * R * S + c * R * S + r * S + 10];
//                  ii[11]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 11];
//                  ff[11]= filter[k * C * R * S + c * R * S + r * S + 11];
//
//                  ii[12]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 12];
//                  ff[12]= filter[k * C * R * S + c * R * S + r * S + 12];
//                  ii[13]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 13];
//                  ff[13]= filter[k * C * R * S + c * R * S + r * S + 13];
//                  ii[14]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 14];
//                  ff[14]= filter[k * C * R * S + c * R * S + r * S + 14];
//                  ii[15]= input[n * C * H * W + c * H * W + (oh+r) * W + ow + 15];
//                  ff[15]= filter[k * C * R * S + c * R * S + r * S + 15];
//		  
//		  o += ii[0]*ff[0] + ii[1]*ff[1] + ii[2]*ff[2] + ii[3]*ff[3]
//		  + ii[4]*ff[4] + ii[5]*ff[5] + ii[6]*ff[6] + ii[7]*ff[7]
//		  + ii[8]*ff[8] + ii[9]*ff[9] + ii[10]*ff[10] + ii[11]*ff[11]
//		  + ii[12]*ff[12] + ii[13]*ff[13] + ii[14]*ff[14] + ii[15]*ff[15];
//
		  i0 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]);
		  //i1 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16);
		  //i2 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*2);
		  //i3 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*3);
		  //i4 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*4);
		  //i5 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*5);
		  //i6 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*6);
		  //i7 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*7);
		  //i8 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*8);
		  //i9 =  _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*9);
		  //i10 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*10);
		  //i11 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*11);
		  //i12 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*12);
		  //i13 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*13);
		  //i14 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*14);
		  //i15 = _mm512_load_ps(&input[n * C * H * W + c * H * W + h * W + ow]+16*15);

		  f0 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S ]);
		  //f1 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16);
		  //f2 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*2);
		  //f3 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*3);
		  //f4 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*4);
		  //f5 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*5);
		  //f6 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*6);
		  //f7 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*7);
		  //f8 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*8);
		  //f9 =  _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*9);
		  //f10 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*10);
		  //f11 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*11);
		  //f12 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*12);
		  //f13 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*13);
		  //f14 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*14);
		  //f15 = _mm512_load_ps(&filter[k * C * R * S + c * R * S + r * S + s]+16*15);
		  //
		  s0 +=  _mm512_fmadd_ps(i0, f0, s0);
		  //s1 =  _mm512_fmadd_ps(i1, f1, s1);
		  //s2 =  _mm512_fmadd_ps(i2, f2, s2);
		  //s3 =  _mm512_fmadd_ps(i3, f3, s3);
		  //s4 =  _mm512_fmadd_ps(i4, f4, s4);
		  //s5 =  _mm512_fmadd_ps(i5, f5, s5);
		  //s6 =  _mm512_fmadd_ps(i6, f6, s6);
		  //s7 =  _mm512_fmadd_ps(i7, f7, s7);
		  //s8 =  _mm512_fmadd_ps(i8, f8, s8);
		  //s9 =  _mm512_fmadd_ps(i9, f9, s9);
		  //s10 = _mm512_fmadd_ps(i10, f10, s10);
		  //s11 = _mm512_fmadd_ps(i11, f11, s11);
		  //s12 = _mm512_fmadd_ps(i12, f12, s12);
		  //s13 = _mm512_fmadd_ps(i13, f13, s13);
		  //s14 = _mm512_fmadd_ps(i14, f14, s14);
		  //s15 = _mm512_fmadd_ps(i15, f15, s15);
		  tmp =s0[0] + s0[1] + s0[2] + s0[3] + s0[4] + s0[5] + s0[6] + s0[7] 
			  + s0[8] + s0[9] + s0[10] + s0[11] + s0[12] + s0[13] +  s0[14] + s0[15];
		  o += tmp;
	      }
	      else
	      {
                 for (int s = 0; s < S; ++s) {
                  int h = oh * stride - pad + r * dilation;
                  int w = ow * stride - pad + s * dilation;
                  if (h < 0 || h >= H || w < 0 || w >= W) continue;
                  float i = input[n * C * H * W + c * H * W + h * W + w];
                  float f = filter[k * C * R * S + c * R * S + r * S + s];
                  o += i * f;
		 }
              }
            }//r
          }//c
          output[n * K * OH * OW + k * OH * OW + oh * OW + ow] = o;
        }//ow
      }//oh
    }//k
  }//n

  if (mpi_rank == 0 && mpi_world_size ==2) 
  {
    MPI_Recv(&output[size[0]*K*OH*OW], size[1]*K*OH*OW, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &status);
  }
  else if(mpi_world_size ==2)
  {
    MPI_Isend(output, size[1]*K*OH*OW, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &request);
  }
}

void convolution_init(
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
  N = _N; C = _C; H = _H; W = _W;
  K = _K; R = _R; S = _S;
  pad = _pad;
  dilation = _dilation;
  stride = _stride;

  MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
  MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size);
}

void convolution_final(
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
}