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

#include "util.h"

#define MAX_NUM_THREAD  40 // 0-39

#define CEIL_DIV(x,y) ( ((x) + (y) - 1) / (y) )
#define CEIL(x,y)     ( CEIL_DIV((x),(y)) * (y) )
#define MIN(a,b)      ( ((a) < (b)) ? (a) : (b) )

#define likely(x)       __builtin_expect((x),1)
#define unlikely(x)     __builtin_expect((x),0)

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;

static void conv_core(float *in, float *flt, float *out)
{
  for (int c = 0; c < C; ++c)
  {
    float *inForCh = &in[c * H * W];
    float *fltForCh = &flt[c * R * S];

#pragma omp parallel num_threads(MAX_NUM_THREAD) 
#pragma omp for collapse(2) schedule(static) nowait
    for (int oh = 0; oh < OH; oh++)
    {
      for (int ow = 0; ow < OW; ow++)
      {
        float o = 0.0f;

        for (int r = 0; r < R; r++)
        {
          for (int s = 0; s < S; s++)
          {
            int h = oh * stride - pad + r * dilation;
            if (unlikely(h < 0 || h >= H))
              continue;

            int w = ow * stride - pad + s * dilation;
            if (unlikely(w < 0 || w >= W))
              continue;

            float i = inForCh[h * W + w];
            float f = fltForCh[r * S + s];

            o += i * f;
          }
        }

        out[oh * OW + ow] += o;
      }
    }
  }
}

#define MPI_ASYNC 1

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)
{
#if MPI_ASYNC == 1
  MPI_Request reqs[3] = { MPI_REQUEST_NULL, MPI_REQUEST_NULL, MPI_REQUEST_NULL };
#endif

  ////////////////////////////////////////////////////////////////////////////////
  // scatter
  if (mpi_rank == 0)
  {
    input = _input;
    output = _output;
    filter = _filter;
  }

#if MPI_ASYNC == 1
  MPI_Ibcast(filter, K * C * R * S, MPI_FLOAT, 0, MPI_COMM_WORLD, &reqs[0]);
#else
  MPI_Bcast(filter, K * C * R * S, MPI_FLOAT, 0, MPI_COMM_WORLD);
#endif

  if (mpi_rank == 0)
  {
    const int slicedN = CEIL_DIV(_N, mpi_world_size);
    const int sizeOfN = slicedN * C * H * W;
    float *inputForNodes = input + (N * C * H * W);

    for (int i = 1; i < mpi_world_size; i++)
    {
#if MPI_ASYNC == 1
      MPI_Isend(inputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 1, MPI_COMM_WORLD, &reqs[1]);
#else
      MPI_Send(inputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 1, MPI_COMM_WORLD);
#endif
    }
  }
  else
  {
#if MPI_ASYNC == 1
    MPI_Irecv(input, N * C * H * W, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &reqs[1]);
#else
    MPI_Recv(input, N * C * H * W, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
#endif

    zero_tensor(output, N, K, OH, OW);
  }

#if MPI_ASYNC == 1
  MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE);
#endif

  ////////////////////////////////////////////////////////////////////////////////
  // calc
  for (int n = 0; n < N; ++n)
  {
    for (int k = 0; k < K; ++k)
    {
      conv_core(
        &input[n * C * H * W],
        &filter[k * C * R * S],
        &output[n * K * OH * OW + k * OH * OW]);
    }
  }
  
  ////////////////////////////////////////////////////////////////////////////////
  // gather
  if (mpi_rank == 0)
  {
    const int slicedN = CEIL_DIV(_N, mpi_world_size);
    const int sizeOfN = slicedN * K * OH * OW;
    float *outputForNodes = output + (N * K * OH * OW);

    for (int i = 1; i < mpi_world_size; i++)
    {
#if MPI_ASYNC == 1
      MPI_Irecv(outputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 2, MPI_COMM_WORLD, &reqs[2]);
#else
      MPI_Recv(outputForNodes + sizeOfN * (i - 1), sizeOfN, MPI_FLOAT, i, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
#endif
    }
  }
  else
  {
#if MPI_ASYNC == 1
    MPI_Isend(output, N * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD, &reqs[2]);
#else
    MPI_Send(output, N * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD);
#endif
  }

#if MPI_ASYNC == 1
  MPI_Wait(&reqs[2], MPI_STATUS_IGNORE);
#endif
}

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

  OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
  OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;

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

  N = CEIL_DIV(_N, mpi_world_size);

  if (mpi_rank == 0)
  {
    N = _N - N * (mpi_world_size - 1);
  }
  else
  {
    alloc_tensor(&input, N, C, H, W);
    alloc_tensor(&filter, K, C, R, S);
    alloc_tensor(&output, N, K, OH, OW);
  }
}

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

  if (mpi_rank != 0)
  {
    free(input);
    free(filter);
    free(output);
  }
}