#include <cstdio>
#include <cuda_runtime.h>
#include <mpi.h>
#include <immintrin.h>
#include "convolution.h"
#include "util.h"

#define CUDA_CALL(f)                                                           \
  {                                                                            \
    cudaError_t err = (f);                                                     \
    if (err != cudaSuccess) {                                                  \
      fprintf(stderr, "CUDA error at [%s:%d] %d %s\n", __FILE__, __LINE__,     \
              err, cudaGetErrorString(err));                                   \
      exit(1);                                                                 \
    }                                                                          \
  }

#define MAX_NUM_GPU (4)
#define SINGLE_NODE (0)
#define TIME_MEASURE (0)
#define ALIGN_UP(_A,_SIZE) ((((_A) + (_SIZE) - 1) / (_SIZE)) * (_SIZE))
#define MIN(_A,_B) ((_A) < (_B) ? (_A) : (_B))
#define OPTIMAL_FILTER_SIZE (16)
#define ENABLE_PREFETCH (1)
#if (ENABLE_PREFETCH)
#define MM_PREFETCH(__A, __B) _mm_prefetch(__A, __B)
#else
#define MM_PREFETCH(__A, __B)
#endif

#define OPT_LEVEL_2 (0)
#define OPTIMAL_NODE_CNT (2)
#define MPI_CH_CNT (4)
#define MPI_FILTER_CH_CNT (MPI_CH_CNT * OPTIMAL_NODE_CNT)
#define TS (8)

static float *__restrict input, *__restrict output, *__restrict 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_devices = 0;

// Array of device (GPU) pointers
static float *input_d[MAX_NUM_GPU];
static float *output_d[MAX_NUM_GPU];
static float *filter_d[MAX_NUM_GPU];
static int Nbegin[MAX_NUM_GPU], Nend[MAX_NUM_GPU], Nsize[MAX_NUM_GPU];


__global__ void k_convolution_Opt2(float *__restrict _input, float *__restrict _output, float *__restrict _filter, \
    int _N, int _C, int _H, int _W, \
    int _K, int _OH, int _OW)
{ 
  const int HW = _H * _W;
  const int CHW = _C * HW;
  const int RS = 16 * 16;
  const int CRS = _C * RS;
  const int OHOW = _OH * _OW;
  const int KOHOW = _K * OHOW;

  const int n = blockIdx.x;
  const int k = blockIdx.y;
  const int c = blockIdx.z;

  const int th_x = threadIdx.x;
  const int th_y = threadIdx.y;

// printf("Optimal calculation :) \n");
// N, C, K, H, W: 32 이상의 적당히 큰 2의 지수승
// R, S: 16

    __shared__ float Localfilter[OPTIMAL_FILTER_SIZE][OPTIMAL_FILTER_SIZE];
    for (int i = 1; i < OPTIMAL_FILTER_SIZE / TS; i++)
    {
        Localfilter[th_x * i][th_y *i] = _filter[k * CRS + (c * RS) + th_x * i * OPTIMAL_FILTER_SIZE + th_y * i];
    }

    __syncthreads();
    
    const int cHW = c * HW;
    const int nCHW_cHW = n * CHW + cHW;
    const int nKOHOW_kOHOW = n * KOHOW + k * OHOW;
    
    for (int oh = th_x; oh < _OH; oh += TS) {
        const int nKOHOW_kOHOW_ohOW = nKOHOW_kOHOW + oh * _OW;
        const int nCHW_ohW_cHW = nCHW_cHW + oh * _W; 
    for (int ow = th_y; ow < _OW; ow += TS) {
            const float* pnStartInput = &_input[nCHW_ohW_cHW + ow];

            if (oh < _OH && ow < _OW)
            {            
                float o = 0.f;
                for (int i = 0; i < OPTIMAL_FILTER_SIZE; i++)
                {
                    for (int j = 0; j < OPTIMAL_FILTER_SIZE; j++)
                    {
                        o += pnStartInput[i * OPTIMAL_FILTER_SIZE + j] * Localfilter[i][j];
                    }
                }

                _output[nKOHOW_kOHOW_ohOW + ow] = o;               

            }
            __syncthreads();
        }
    }
}

__global__ void k_convolution_Opt(float *__restrict _input, float *__restrict _output, float *__restrict _filter, \
    const int _N, const int _C, const int _H, const int _W, \
    const int _K, const int _OH, const int _OW)
{
    const int globalRow = blockDim.x * blockIdx.x + threadIdx.x;
    const int globalCol = blockDim.y * blockIdx.y + threadIdx.y;
    
    if (globalRow >= _OH || globalCol >= _N * _K * _OW)
    {
        // printf("Can't enter here!! :GR :%d, GC :%d\n",globalRow,globalCol);
        return;
    }
    // printf("Come in global row : [%x] Global Col : %d\n", globalRow, globalCol);
    

    int n, k, w;

    n = globalCol / (_K * _OW);
    w = globalCol - n * (_K * _OW);
    k = w / _OW;
    w = w - k * _OW;

    int col = w;
    float o = 0.f;
    const int nCWH_rowW_col = n *_C*_W*_H + globalRow*_W + col;
    const int kCRS = k*_C*OPTIMAL_FILTER_SIZE*OPTIMAL_FILTER_SIZE;
    const int WH = _W*_H;

    for (int c = 0; c < _C; c++){
        const int nCWH_cWH_rowW_col = nCWH_rowW_col + c*WH;
        const int kCRS_cRS = kCRS + c*OPTIMAL_FILTER_SIZE*OPTIMAL_FILTER_SIZE;
        for (int i = 0; i < OPTIMAL_FILTER_SIZE; i++){
#if 1
            const int in_offset = nCWH_cWH_rowW_col + i*_W;
            for (int j = 0; j < OPTIMAL_FILTER_SIZE / 4; j++)
            {
                float4 filter_v = *((float4*)&_filter[kCRS_cRS + i*OPTIMAL_FILTER_SIZE + j * 4]);
                o += (_input[in_offset + j * 4 + 0] * filter_v.x 
                     + _input[in_offset + j * 4 + 1]  * filter_v.y 
                     + _input[in_offset + j * 4 + 2]  * filter_v.z 
                     + _input[in_offset + j * 4 + 3]  * filter_v.w);
            }
#else
            for (int j = 0; j < OPTIMAL_FILTER_SIZE; j++)
            {
                float in = _input[nCWH_cWH_rowW_col + i*_W + j];
                float filter = _filter[kCRS_cRS + i*OPTIMAL_FILTER_SIZE + j];
                o += in * filter;
            }
#endif
        }
    }

    _output[n * _K*_OH*_OW + k*_OH*_OW + globalRow*_OW + col] = o;

#if 0 // print
    if ((n * _K*_OH*_OW + k*_OH*_OW + row*_OW + col) % 0x4000 == 0)
    {
        printf("_output[%x] set %3.f\n", n * _K*_OH*_OW + k*_OH*_OW + row*_OW + col, o);
    }
#endif
}


__global__ void k_convolution_base(float *__restrict _input, float *__restrict _output, float *__restrict _filter, \
    int _N, int _C, int _H, int _W, \
    int _K, int _R, int _S, \
    int _pad, int _dilation, int _stride, \
    int _OH, int _OW)
{
    const int globalRow = blockDim.x * blockIdx.x + threadIdx.x;
    const int globalCol = blockDim.y * blockIdx.y + threadIdx.y;
    
    if (globalRow >= _OH || globalCol >= _N * _K * _OW)
    {
        return;
    }
    // printf("Come in global row : [%x] Global Col : %d\n", globalRow, globalCol);
    

    int n, k, w;
    w = globalCol;
    n = w / (_K * _OW);
    w = w - n * (_K * _OW);
    k = w / _OW;
    w = w - k * _OW;

    int col = w;
    int row = globalRow;

    int start_row = row * _stride - _pad;
    int start_col = col * _stride - _pad;

    float o = 0.f;
    for (int c = 0; c < _C; c++){
        for (int i = 0; i < _R; i++){
            for (int j = 0; j < _S; j++)
            {
                int h = start_row + i * _dilation;
                int w = start_col + j * _dilation;
                if (h < 0 || w < 0 || h >= _H || w >= _W) continue;
                float in = _input[n *_C*_W*_H + c*_W*_H + h*_W + w];
                float filter = _filter[k*_C*_R*_S + c*_R*_S + i*_S + j];
                o += in * filter;
            }
        }
    }

    _output[n * _K*_OH*_OW + k*_OH*_OW + row*_OW + col] = o;

#if 0 // print
    if ((n * _K*_OH*_OW + k*_OH*_OW + row*_OW + col) % 0x4000 == 0)
    {
        printf("_output[%x] set %3.f\n", n * _K*_OH*_OW + k*_OH*_OW + row*_OW + col, o);
    }
#endif
}



#define OPTIMAL_TS_X (2)
#define OPTIMAL_TS_Y (32)
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 (!SINGLE_NODE)
  MPI_Status stMpiStatus;
  MPI_Request stMpiRequest[12];
#endif

  if (_pad == 0
      && _dilation == 1
      && _stride == 1
#if (!SINGLE_NODE)
      && mpi_world_size == OPTIMAL_NODE_CNT
  #endif
      && (((N | C | K | H | W) & (32 - 1)) == 0)
      && R == 16 && S == 16)
      {
#if (!SINGLE_NODE)
        const int SendNodeSize = N / OPTIMAL_NODE_CNT / MPI_CH_CNT;
        const int SendFilterSize = K / MPI_CH_CNT;
#endif
        // Optimal path
        if (mpi_rank == 0) {
          input = _input;
          output = _output;
          filter = _filter;
      #if (!SINGLE_NODE)
          MPI_Isend(input + (SendNodeSize * (MPI_CH_CNT + 0)) * C * H * W, (SendNodeSize) * C * H * W , MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &stMpiRequest[0]);
          MPI_Isend(input + (SendNodeSize * (MPI_CH_CNT + 1)) * C * H * W, (SendNodeSize) * C * H * W , MPI_FLOAT, 1, 1, MPI_COMM_WORLD, &stMpiRequest[1]);
          MPI_Isend(input + (SendNodeSize * (MPI_CH_CNT + 2)) * C * H * W, (SendNodeSize) * C * H * W , MPI_FLOAT, 1, 2, MPI_COMM_WORLD, &stMpiRequest[2]);
          MPI_Isend(input + (SendNodeSize * (MPI_CH_CNT + 3)) * C * H * W, (SendNodeSize) * C * H * W , MPI_FLOAT, 1, 3, MPI_COMM_WORLD, &stMpiRequest[3]);
          MPI_Isend(filter                                 , SendFilterSize * C * R * S, MPI_FLOAT, 1, 4, MPI_COMM_WORLD, &stMpiRequest[4]);
          MPI_Isend(filter + SendFilterSize * C * R * S    , SendFilterSize * C * R * S, MPI_FLOAT, 1, 5, MPI_COMM_WORLD, &stMpiRequest[5]);
          MPI_Isend(filter + SendFilterSize * 2 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 1, 6, MPI_COMM_WORLD, &stMpiRequest[6]);
          MPI_Isend(filter + SendFilterSize * 3 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 1, 7, MPI_COMM_WORLD, &stMpiRequest[7]);
          MPI_Irecv(output + (SendNodeSize * (MPI_CH_CNT + 0)) * K * OH * OW, (SendNodeSize) * K * OH * OW, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &stMpiRequest[8]);
          MPI_Irecv(output + (SendNodeSize * (MPI_CH_CNT + 1)) * K * OH * OW, (SendNodeSize) * K * OH * OW, MPI_FLOAT, 1, 1, MPI_COMM_WORLD, &stMpiRequest[9]);
          MPI_Irecv(output + (SendNodeSize * (MPI_CH_CNT + 2)) * K * OH * OW, (SendNodeSize) * K * OH * OW, MPI_FLOAT, 1, 2, MPI_COMM_WORLD, &stMpiRequest[10]);
          MPI_Irecv(output + (SendNodeSize * (MPI_CH_CNT + 3)) * K * OH * OW, (SendNodeSize) * K * OH * OW, MPI_FLOAT, 1, 3, MPI_COMM_WORLD, &stMpiRequest[11]);
          //printf("Master receive : %x, %x, %x, %x\n", (SendNodeSize * MPI_CH_CNT + 0), (SendNodeSize * MPI_CH_CNT + 1), (SendNodeSize * MPI_CH_CNT + 2), (SendNodeSize * MPI_CH_CNT + 3));
      #if (TIME_MEASURE)
          printf("Master send started : %f sec\n", timer_stop(0));
      #endif
      #endif
      // printf ("Optimized path! SendNodeOffset : %d, mpi_world_size : %d", N / OPTIMAL_NODE_CNT / MPI_CH_CNT, mpi_world_size);
        // Launch kernel on every GPU

        // Upload A and B matrix to every GPU
        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * C * H * W, (Nsize[i]) * C * H * W * sizeof(float),
                    cudaMemcpyHostToDevice) );
            CUDA_CALL( cudaMemcpy(filter_d[i], filter, K * C * R * S * sizeof(float), cudaMemcpyHostToDevice) );
        }
#if (OPT_LEVEL_2)
        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS, OPTIMAL_TS, 1);
            dim3 gridDim(Nsize[i], K, C);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_Opt2<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, OH, OW);
        }
#else
        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS_X, OPTIMAL_TS_Y, 1);
            dim3 gridDim((OH + OPTIMAL_TS_X - 1) /OPTIMAL_TS_X, (Nsize[i] * K * OW + OPTIMAL_TS_Y - 1) / OPTIMAL_TS_Y, 1);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_Opt<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, OH, OW);
        }
#endif

        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaDeviceSynchronize() );
        } 

        // Download C matrix from GPU
        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaMemcpy(output + Nbegin[i] * K * OH * OW, output_d[i], (Nsize[i]) * K * OH * OW * sizeof(float),
                    cudaMemcpyDeviceToHost) );
        }

      #if (TIME_MEASURE)
          printf("Master calculation complete : %f sec\n", timer_stop(0));
      #endif
      #if (!SINGLE_NODE)
          for(int i = 0; i < 12; i++)
          {
            MPI_Wait(&stMpiRequest[i], &stMpiStatus);
          }
      #if (TIME_MEASURE)
          printf("Master recieve complete : %f sec\n", timer_stop(0));
      #endif
    #endif
      }
#if (!SINGLE_NODE)
      else
      {
        // printf("Check memory pointer, input : %x Filter :% x output : %x\n", input, filter, output);
        MPI_Irecv(input + (SendNodeSize * 0) * C * H * W, SendNodeSize * C * H * W , MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &stMpiRequest[0]);
        MPI_Irecv(input + (SendNodeSize * 1) * C * H * W, SendNodeSize * C * H * W , MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &stMpiRequest[1]);
        MPI_Irecv(input + (SendNodeSize * 2) * C * H * W, SendNodeSize * C * H * W , MPI_FLOAT, 0, 2, MPI_COMM_WORLD, &stMpiRequest[2]);
        MPI_Irecv(input + (SendNodeSize * 3) * C * H * W, SendNodeSize * C * H * W , MPI_FLOAT, 0, 3, MPI_COMM_WORLD, &stMpiRequest[3]);
        MPI_Irecv(filter + SendFilterSize * 0 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 0, 4, MPI_COMM_WORLD, &stMpiRequest[4]);
        MPI_Irecv(filter + SendFilterSize * 1 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 0, 5, MPI_COMM_WORLD, &stMpiRequest[5]);
        MPI_Irecv(filter + SendFilterSize * 2 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 0, 6, MPI_COMM_WORLD, &stMpiRequest[6]);
        MPI_Irecv(filter + SendFilterSize * 3 * C * R * S, SendFilterSize * C * R * S, MPI_FLOAT, 0, 7, MPI_COMM_WORLD, &stMpiRequest[7]);

        for(int i = 0; i < 8; i++)
        {
            MPI_Wait(&stMpiRequest[i], &stMpiStatus);
        }

        // Upload A and B matrix to every GPU
        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * C * H * W, (Nsize[i]) * C * H * W * sizeof(float),
                    cudaMemcpyHostToDevice) );
            CUDA_CALL( cudaMemcpy(filter_d[i], filter, K * C * R * S * sizeof(float), cudaMemcpyHostToDevice) );
        }
#if (OPT_LEVEL_2)
        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS, OPTIMAL_TS, 1);
            dim3 gridDim(Nsize[i], K, C);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_Opt2<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, OH, OW);
        }
#else
        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS_X, OPTIMAL_TS_Y, 1);
            dim3 gridDim((OH + OPTIMAL_TS_X - 1) /OPTIMAL_TS_X, (Nsize[i] * K * OW + OPTIMAL_TS_Y - 1) / OPTIMAL_TS_Y, 1);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_Opt<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, OH, OW);
        }
#endif

        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaDeviceSynchronize() );
        } 

        // Download C matrix from GPU
        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaMemcpy(output + Nbegin[i] * K * OH * OW, output_d[i], (Nsize[i]) * K * OH * OW * sizeof(float),
                    cudaMemcpyDeviceToHost) );
        }
        
        // printf("Slave set from output[%d], stride all : [%d]\n", SendNodeOffset, SendNodeOffset * KOHOW);    
        // printf("Slave set end output[%d], stride all : [%d]\n", N, N * KOHOW);
    #if (TIME_MEASURE)
        printf("Slave calculation complete : %f sec\n", timer_stop(0));
    #endif
        // printf("Slave send from output[%d]\n", SendNodeOffset * C * OH * OW);
        MPI_Isend(output                                   , SendNodeSize * K * OH * OW, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &stMpiRequest[0]);
        MPI_Isend(output + (SendNodeSize * 1) * K * OH * OW, SendNodeSize * K * OH * OW, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &stMpiRequest[1]);
        MPI_Isend(output + (SendNodeSize * 2) * K * OH * OW, SendNodeSize * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD, &stMpiRequest[2]);
        MPI_Isend(output + (SendNodeSize * 3) * K * OH * OW, SendNodeSize * K * OH * OW, MPI_FLOAT, 0, 3, MPI_COMM_WORLD, &stMpiRequest[3]);
        //printf ("FirstData : %3.f, %3.f, %3.f, %3.f\n", *(output + (SendNodeSize *0) * K * OH * OW),*(output + (SendNodeSize *1) * K * OH * OW),*(output + (SendNodeSize *2) * K * OH * OW),*(output + (SendNodeSize *3) * K * OH * OW));
        //printf("Slave send : %x, %x, %x, %x\n",(SendNodeSize * 0) * K * OH * OW, (SendNodeSize * 1) * K * OH * OW, (SendNodeSize * 2) * K * OH * OW, (SendNodeSize * 3) * K * OH * OW);
        for(int i = 0; i < 4; i++)
        {
          MPI_Wait(&stMpiRequest[i], &stMpiStatus);
        }
    #if (TIME_MEASURE)
        printf("Slave send complete : %f sec\n", timer_stop(0));
    #endif
        }
#endif // SINGLE_NODE
      }
      else
      {
#if (!SINGLE_NODE)
        const int SendNodeSize = N / 2;
        const int SendNodeOffset = mpi_world_size > 1 ? (N - SendNodeSize) : N;
#endif
      if (mpi_rank == 0) {
        input = _input;
        output = _output;
        filter = _filter;

        // printf("Check memory pointer, input : %x Filter :% x output : %x, Rank :%d\n", input, filter, output, mpi_rank);
        // printf("Check memory size, input : %d Filter :%d output : %d, Rank :%d\n", sizeof(float) * N * C*H*W, sizeof(float) *K * C * R * S, sizeof(float) *N * K * OH * OW, mpi_rank);

    #if (!SINGLE_NODE)
        if (mpi_world_size > 1 && SendNodeSize > 0)
        {
          MPI_Isend(input + SendNodeOffset * C * H * W, SendNodeSize * C * H * W , MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &stMpiRequest[0]);
          MPI_Isend(filter, K * C * R * S, MPI_FLOAT, 1, 1, MPI_COMM_WORLD, &stMpiRequest[1]);
          MPI_Irecv(output + SendNodeOffset * K * OH * OW, SendNodeSize * K * OH * OW, MPI_FLOAT, 1, 2, MPI_COMM_WORLD, &stMpiRequest[2]);
    #if (TIME_MEASURE)
        printf("Master send started : %f sec\n", timer_stop(0));
    #endif
        }
    #endif

        // Upload A and B matrix to every GPU
        for (int i = 0; i < num_devices; i++) {
            // printf("Send to %d, size :%d amount : %d from %d MPI rank : %d\n", Nbegin[i],Nsize[i],(Nsize[i]) * C * H * W * sizeof(float), i, mpi_rank);
            // printf("Check memory pointer, input : %x Filter :% x output : %x, Rank :%d\n", input, filter, output, mpi_rank);
            CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * C * H * W, (Nsize[i]) * C * H * W * sizeof(float),
                    cudaMemcpyHostToDevice) );
            CUDA_CALL( cudaMemcpy(filter_d[i], filter, K * C * R * S * sizeof(float), cudaMemcpyHostToDevice) );
        }

        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS_X, OPTIMAL_TS_Y, 1);
            dim3 gridDim((OH + OPTIMAL_TS_X - 1) /OPTIMAL_TS_X, (Nsize[i] * K * OW + OPTIMAL_TS_Y - 1) / OPTIMAL_TS_Y, 1);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_base<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, _R, _S, _pad, _dilation, _stride, OH, OW);
        }

        
        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaDeviceSynchronize() );
        } 

        // Download C matrix from GPU
        for (int i = 0; i < num_devices; i++) {
            // printf("Recieved to %d, size %d amount : %d from %d MPI rank : %d\n", Nbegin[i], Nsize[i], (Nsize[i]) * K * OH * OW * sizeof(float), i, mpi_rank);
            const int nOffset =  Nbegin[i] * K * OH * OW;
            CUDA_CALL( cudaMemcpy(output + nOffset, output_d[i], (Nsize[i]) * K * OH * OW * sizeof(float),
                    cudaMemcpyDeviceToHost) );

            // printf("Output : %3.f, %3.f, %3.f, ... from %d MPI rank : %d\n", *(output + nOffset), *(output + nOffset + 1), *(output + nOffset + 2), i, mpi_rank);
        }

    #if (TIME_MEASURE)
        printf("Master calculation complete : %f sec\n", timer_stop(0));
    #endif
    #if (!SINGLE_NODE)
        if (mpi_world_size > 1 && SendNodeSize > 0)
        {
          MPI_Wait(&stMpiRequest[0], &stMpiStatus);
          MPI_Wait(&stMpiRequest[1], &stMpiStatus);
          MPI_Wait(&stMpiRequest[2], &stMpiStatus);
    #if (TIME_MEASURE)
        printf("Master recieve complete : %f sec\n", timer_stop(0));
    #endif
        }
    #endif
      }
    #if (!SINGLE_NODE)
      else if (SendNodeSize > 0)
      {
        // printf("Check memory pointer, input : %x Filter :% x output : %x, Rank :%d\n", input, filter, output, mpi_rank);
        MPI_Irecv(input, SendNodeSize * C * H * W , MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &stMpiRequest[0]);
        MPI_Irecv(filter, K * C * R * S, MPI_FLOAT, 0, 1, MPI_COMM_WORLD, &stMpiRequest[1]);

        MPI_Wait(&stMpiRequest[0], &stMpiStatus);
        MPI_Wait(&stMpiRequest[1], &stMpiStatus);
    #if (TIME_MEASURE)
        printf("Slave receive complete : %f sec\n", timer_stop(0));
    #endif
        // Upload A and B matrix to every GPU
        for (int i = 0; i < num_devices; i++) {
            // printf("Send to %d, amount : %x from %d MPI rank : %d\n", Nbegin[i],  (Nsize[i]) * C * H * W * sizeof(float), i, mpi_rank);
            CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * C * H * W, (Nsize[i]) * C * H * W * sizeof(float),
                    cudaMemcpyHostToDevice) );
            CUDA_CALL( cudaMemcpy(filter_d[i], filter, K * C * R * S * sizeof(float), cudaMemcpyHostToDevice) );
        }

        for (int i = 0; i < num_devices; i++) {
            dim3 blockDim(OPTIMAL_TS_X, OPTIMAL_TS_Y, 1);
            dim3 gridDim((OH + OPTIMAL_TS_X - 1) /OPTIMAL_TS_X, (Nsize[i] * K * OW + OPTIMAL_TS_Y - 1) / OPTIMAL_TS_Y, 1);
            CUDA_CALL( cudaSetDevice(i) );
            k_convolution_base<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], Nsize[i], _C, _H, _W, _K, _R, _S, _pad, _dilation, _stride, OH, OW);
        }

        for (int i = 0; i < num_devices; i++) {
            CUDA_CALL( cudaDeviceSynchronize() );
        } 

        // Download C matrix from GPU
        for (int i = 0; i < num_devices; i++) {
            // printf("Recieved to %d, amount : %x from %d MPI rank : %d\n", Nbegin[i],  (Nsize[i]) * K * OH * OW * sizeof(float), i, mpi_rank);
            CUDA_CALL( cudaMemcpy(output + Nbegin[i] * K * OH * OW, output_d[i], (Nsize[i]) * K * OH * OW * sizeof(float),
                    cudaMemcpyDeviceToHost) );
        }
        
        // printf("Slave set from output[%d], stride all : [%d]\n", SendNodeOffset, SendNodeOffset * KOHOW);    
        // printf("Slave set end output[%d], stride all : [%d]\n", N, N * KOHOW);
    #if (TIME_MEASURE)
        printf("Slave calculation complete : %f sec\n", timer_stop(0));
    #endif
        
        
        // printf("Slave send from output[%d]\n", SendNodeOffset * C * OH * OW);
        MPI_Isend(output, SendNodeSize * K * OH * OW, MPI_FLOAT, 0, 2, MPI_COMM_WORLD, &stMpiRequest[2]); 
        MPI_Wait(&stMpiRequest[2], &stMpiStatus);
    #if (TIME_MEASURE)
        printf("Slave send complete : %f sec\n", timer_stop(0));
    #endif
      }
    #endif
  }

    // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  } 
}

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;
  
  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);
  // printf("MPI rank : %d MPI world size :%d\n", mpi_rank, mpi_world_size);

  if (mpi_rank != 0)
  {
    const int SendNodeSize = N / 2;

    alloc_tensor((float**)&input, SendNodeSize, _C, _H, _W);
    alloc_tensor((float**)&output, SendNodeSize, _K, OH, OW);
    alloc_tensor((float**)&filter, _K, _C, _R, _S);   
    
    // printf("Set slave memory pointer, input : %x Filter :% x output : %x\n", input, filter, output);
  }

  CUDA_CALL(cudaGetDeviceCount(&num_devices));

  printf("Using %d devices\n", num_devices);
  for (int i = 0; i < num_devices; i++) {
    cudaDeviceProp prop;
    CUDA_CALL( cudaGetDeviceProperties(&prop, i) );

    // Try printing more detailed information here
    printf("[GPU %d] %s\n", i, prop.name);
  }

  if (num_devices <= 0) {
    printf("No CUDA device found. Aborting\n");
    exit(1);
  }

  // Setup problem size for each GPU
  for (int i = 0; i < num_devices; i++) {
    Nbegin[i] = (_N / num_devices / mpi_world_size) * i;
    Nend[i] = (_N / num_devices / mpi_world_size) * (i + 1);
    Nsize[i] = Nend[i] - Nbegin[i];
  }
  if (mpi_world_size == 2)
  {
    Nend[num_devices - 1] = mpi_rank == 0? (_N - (_N/mpi_world_size)) : (_N/mpi_world_size);
    Nsize[num_devices - 1] = Nend[num_devices - 1] - Nbegin[num_devices - 1];
  }
  else
  {
    Nend[num_devices - 1] = _N;
    Nsize[num_devices - 1] = Nend[num_devices - 1] - Nbegin[num_devices - 1];
  }
  
  // Allocate device memory for each GPU
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaSetDevice(i) );
    CUDA_CALL( cudaMalloc(&input_d[i], (Nsize[i]) * C * H * W * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&filter_d[i], K * C * R * S * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&output_d[i], (Nsize[i]) * K * OH * OW * sizeof(float)) );
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  }
}

void convolution_final(
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
  // Do any post-matmul cleanup work here.

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  }
}