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

#define TS 8
#define MAX_NODE 2
#define MAX_NUM_GPU 8
#define TILE_WIDTH 32 

#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);                                                            \
   }                                                                      \
  }                                                                        



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 float *in_d[MAX_NUM_GPU];
static float *out_d[MAX_NUM_GPU];
static float *fil_d[MAX_NUM_GPU];
static int Mbegin[MAX_NUM_GPU], Mend[MAX_NUM_GPU];
static int MM[MAX_NUM_GPU];
int num_devices;


__global__ void conv(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){

 //  const int globalRow = blockDim.x * blockIdx.x + threadIdx.x;
 //  const int globalCol = blockDim.y * blockIdx.y + threadIdx.y;
   int OH, OW;

   OH = (_H + 2* _pad - _dilation*(_R-1) -1) /_stride + 1;
   OW = (_W + 2* _pad - _dilation*(_S-1) -1) /_stride + 1;

//   __shared__ float ds_i[TILE_WIDTH][TILE_WIDTH];
//   __shared__ float ds_f[TILE_WIDTH][TILE_WIDTH];

   int n = blockIdx.x;
   int k = blockIdx.y;
   int oh = blockIdx.z;
   int ow = threadIdx.x;
   
  
   float o = 0;
   for(int c = 0; c < _C; c++){
     for(int r=0; r< _R; r++){
       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;
       }
     }
   }
   _output[n*_K*OH*OW + k*OH*OW + oh*OW + ow] = o;
}

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) {

  int size[MAX_NODE];

  input = _input;
  output = _output;
  filter = _filter;

  MPI_Request request;
  MPI_Status status;

//  if(mpi_world_size == 2)
//      size[1] = _N/2;
//  else 
      size[1] = 0;
  
  size[0] = _N - size[1];

/* if(mpi_world_size == 2){
      
      for(int i=0; i< num_devices; i++){
          Mbegin[i] = (N/2) /num_devices * i;
          Mend[i] = (N/2) /num_devices*(i+1);
      }
      for(int i=0; i< num_devices; i++){
          Mbegin[i+4] = (N/2)/num_devices * i;
          Mend[i+4] = (N/2)/num_devices*(i+1);
      }
      Mend[num_devices*2-1] = N;

      for(int i=0; i<num_devices*2; i++){
          MM[i] = Mend[i] - Mbegin[i];
      }
  }else{
*/  
      for(int i=0; i< num_devices; i++){
          Mbegin[i] = (N/num_devices) * i;
          Mend[i] = (N/num_devices)*(i+1);
      }
      Mend[num_devices-1] = N;

      for(int i=0; i<num_devices; i++){
          MM[i] = Mend[i] - Mbegin[i];
      }
//  }

  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);
  }
 */ 
  if(mpi_rank == 0){
      for(int i=0; i< num_devices ; i++){
          CUDA_CALL( cudaSetDevice(i) );
          CUDA_CALL( cudaMalloc(&in_d[i], MM[i]*_C*_H*_W*sizeof(float)));
          CUDA_CALL( cudaMalloc(&out_d[i], MM[i]*_K*OH*OW*sizeof(float)));
          CUDA_CALL( cudaMalloc(&fil_d[i], _K*_C*_R*_S*sizeof(float)));
      }
  
      for(int i=0; i< num_devices; i++){
          CUDA_CALL( cudaMemcpy(in_d[i], _input + Mbegin[i]*_C*_H*_W, MM[i]*_C*_H*_W*sizeof(float), cudaMemcpyHostToDevice));
          CUDA_CALL( cudaMemcpy(fil_d[i], _filter, _K*_C*_R*_S*sizeof(float), cudaMemcpyHostToDevice));
      }

      for(int i=0; i <num_devices; i++){
          CUDA_CALL( cudaDeviceSynchronize() );
      }
  //printf("check1\n");
      for(int i=0; i<num_devices; i++){
	  dim3 blockDim(OW, 1);
          dim3 gridDim(MM[i],_K,OH);

          CUDA_CALL( cudaSetDevice(i) );
          conv<<<gridDim, blockDim>>>(in_d[i],out_d[i],fil_d[i],MM[i],_C,_H,_W,_K,_R,_S,_pad,_dilation,_stride);     
      }
 // printf("check2\n");
      for(int i=0; i< num_devices; i++){
          CUDA_CALL( cudaDeviceSynchronize());
      }
  }
  else{
/*      printf("no 1\n");	  
      for(int i=0; i< num_devices; i++){
          CUDA_CALL( cudaSetDevice(i) );
          CUDA_CALL( cudaMalloc(&in_d[i+4], MM[i+4]*_C*_H*_W*sizeof(float)));
          CUDA_CALL( cudaMalloc(&out_d[i+4], MM[i+4]*_K*OH*OW*sizeof(float)));
          CUDA_CALL( cudaMalloc(&fil_d[i+4], _K*_C*_R*_S*sizeof(float)));
      }
      printf("no 2\n");
      for(int i=0; i< num_devices; i++){
          CUDA_CALL( cudaMemcpy(in_d[i+4], _input + Mbegin[i+4]*_C*_H*_W, MM[i+4]*_C*_H*_W*sizeof(float), cudaMemcpyHostToDevice));
          CUDA_CALL( cudaMemcpy(fil_d[i+4], _filter, _K*_C*_R*_S*sizeof(float), cudaMemcpyHostToDevice));
      }
      printf("no 3\n");
      for(int i=0; i <num_devices; i++){
          CUDA_CALL( cudaDeviceSynchronize() );
      }
      printf("check1\n");
      for(int i=0; i<num_devices; i++){
          dim3 blockDim(OW, 1);
          dim3 gridDim(MM[i+4],_K,OH);

          CUDA_CALL( cudaSetDevice(i) );
          conv<<<gridDim, blockDim>>>(in_d[i+4],out_d[i+4],fil_d[i+4],MM[i+4],_C,_H,_W,_K,_R,_S,_pad,_dilation,_stride);
      }
      printf("check2\n");
      for(int i=0; i< num_devices; i++){
          CUDA_CALL( cudaDeviceSynchronize());
*///      }
  }
 // printf("check3\n");
}


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);

  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) );

      printf("[GPU %d] %s\n", i, prop.name);
  }

  if(num_devices <= 0){
     printf("No CUDA device found. Aborting\n");
     exit(1);
  }     
/*
  if(mpi_world_size == 2){
      size[1] = _N/2;
      node_num = 2;
  }
  else {
      size[1] = 0;
      node_num = 1;
  }

  size[0] = N - size[1];  

  for(int i=0; i< num_devices; i++){
      Mbegin[i] = (N/num_devices) * i;
      Mend[i] = (N/num_devices)*(i+1);
  }
  Mend[num_devices-1] = N;

  for(int i=0; i<num_devices; i++){
      MM[i] = Mend[i] - Mbegin[i];
  }
*/  
/*
  OH = (H + 2*pad - dilation*(R -1) -1)/stride +1;
  OW = (W + 2*pad - dilation*(S -1) -1)/stride +1;

  for(int i=0; i< num_devices ; i++){
      CUDA_CALL( cudaSetDevice(i) );
      CUDA_CALL( cudaMalloc(&in_d[i], MM[i]*_C*_H*_W*sizeof(float)));
      CUDA_CALL( cudaMalloc(&out_d[i], MM[i]*_K*OH*OW*sizeof(float)));
      CUDA_CALL( cudaMalloc(&fil_d[i], _K*_C*_R*_S*sizeof(float)));
  }

  for(int i=0; i< num_devices; i++){
      CUDA_CALL( cudaMemcpy(in_d[i], _input + Mbegin[i]*_C*_H*_W, MM[i]*_C*_H*_W*sizeof(float), cudaMemcpyHostToDevice));
      CUDA_CALL( cudaMemcpy(fil_d[i], _filter, _K*_C*_R*_S*sizeof(float), cudaMemcpyHostToDevice));
  }

  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) {
    
/*    if(mpi_rank == 0 && mpi_world_size == 2){	
        printf("final 1\n");	
        for(int i=0; i<num_devices*2; i++){
            CUDA_CALL( cudaMemcpy(output+Mbegin[i]*K*OH*OW, out_d[i], MM[i]*K*OH*OW*sizeof(float), cudaMemcpyDeviceToHost));
        }
        printf("final 2\n");
        for(int i = 0; i < num_devices*2;i++){
	    CUDA_CALL( cudaDeviceSynchronize() );
        }
    }else if(mpi_rank == 0 && mpi_world_size == 1){
     
 */ 
    if(mpi_rank == 0){	
    	for(int i=0; i<num_devices; i++){
            CUDA_CALL( cudaMemcpy(output+Mbegin[i]*K*OH*OW, out_d[i], MM[i]*K*OH*OW*sizeof(float), cudaMemcpyDeviceToHost));
        }

        for(int i = 0; i < num_devices;i++){
            CUDA_CALL( cudaDeviceSynchronize() );
        }}
 //    }
    //printf("Done\n");    
}