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

#include <cstdio>
#include <cuda_runtime.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

static float *input, *output, *filter;
static float *in_d[MAX_NUM_GPU], *out_d[MAX_NUM_GPU], *fil_d[MAX_NUM_GPU];
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;  // ORG
static int num_devices = 1;    // DEBUG
static int Nsizei[MAX_NUM_GPU];		// Number of Nsize for each GPU[i]
static int Nodesize[2];		// Number of size for each node

/*
// Array of device (GPU) pointers
static float *input[MAX_NUM_GPU];
static float *filter[MAX_NUM_GPU];
static float *output[MAX_NUM_GPU];
static int Mbegin[MAX_NUM_GPU], Mend[MAX_NUM_GPU];
*/

// 4279 GFLOPS (TS = 4)
// 4887 GFLOPS (TS = 8)
// 3303 GFLOPS (TS = 16)
//#define TS 4
#define TS 8
//#define TS 16

__global__ void conv_kernel(
    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;

	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;

	if( (globalRow >= OH) || (globalCol >= _N* _K * OW)) return;

	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 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;
                  int h = start_row + r * _dilation;
                  int w = start_col + 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;
                }		// for s
              }		// for r
            }	// for c
            //_output[n * _K * OH * OW + k * OH * OW + oh * OW + ow] = o;
            _output[n * _K * OH * OW + k * OH * OW + row * OW + col] = 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 offset = 0;

	MPI_Request request;
	MPI_Status status;

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

	// printf("DEBUG: mpi_world_size = %d\n", mpi_world_size);

	/*
	// for two processors
	if(mpi_world_size == 2) Nodesize[1] = _N / 2;
	else Nodesize[1] = 0;
	Nodesize[0] = _N - Nodesize[1];

  OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
  OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;
	*/
  
	// For two nodes, separate the jobs
	if(mpi_rank == 0 && mpi_world_size == 2 && Nodesize[1] != 0) {		// for the first node
			MPI_Isend(&input[Nodesize[0]*C*H*W], Nodesize[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);
			if(Nodesize[mpi_rank] < MAX_NUM_GPU) {	// smaller number of size than MAX_NUM_GPU
					num_devices = Nodesize[mpi_rank];
			}
	}
	else if(mpi_rank == 1 && Nodesize[mpi_rank] != 0) {			// for the second node
			alloc_tensor(&input, Nodesize[1], C, H, W);
			alloc_tensor(&output, Nodesize[1], K, OH, OW);
			alloc_tensor(&filter, _K, _C, _R, _S);
			MPI_Recv(input, Nodesize[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(Nodesize[mpi_rank] < MAX_NUM_GPU) {
					num_devices = Nodesize[mpi_rank];
			}
	}

	offset = 0;
	for(int i=0; i < num_devices; i++) {
			CUDA_CALL( cudaMemcpy(in_d[i], input + offset, Nsizei[i]*C*H*W*sizeof(float), cudaMemcpyHostToDevice));
			CUDA_CALL( cudaMemcpy(fil_d[i], filter, K*C*R*S*sizeof(float), cudaMemcpyHostToDevice));
			offset += Nsizei[i] * C * H * W;
	}

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

			
	for(int i=0; i < num_devices; i++) {
			dim3 gridDim ( (OH+TS-1)/TS, (Nsizei[i]*K*OW + TS -1)/TS, 1);
			dim3 blockDim(TS, TS, 1);

			CUDA_CALL(cudaSetDevice(i));
			conv_kernel<<<gridDim, blockDim>>>(in_d[i], out_d[i], fil_d[i], Nsizei[i], _C, _H, _W, _K, _R, _S, _pad, _dilation, _stride);
	}


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

	// gather for the same node
	offset = 0;
	for(int i=0; i < num_devices; i++) {
			CUDA_CALL( cudaSetDevice(i) );
			CUDA_CALL( cudaMemcpy(output + offset, out_d[i], Nsizei[i]*K*OH*OW *sizeof(float), cudaMemcpyDeviceToHost) );
			offset += Nsizei[i]*K*OH*OW;
	}


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

	// gather for the different node
	if(mpi_rank==0 && mpi_world_size == 2 && Nodesize[1] != 0) {
			MPI_Recv(&output[Nodesize[0]*K*OH*OW], Nodesize[1]*K*OH*OW, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &status);
	}
	else if(mpi_rank==1 && Nodesize[1] != 0) {
			MPI_Isend(output, Nodesize[1]*K*OH*OW, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &request);
	}

}


/*

//  if (mpi_rank == 0) {

		//#pragma omp parallel for schedule(dynamic) num_threads(num_threads)
		//#pragma omp parallel for schedule(dynamic) 
		//#pragma omp parallel for schedule(dynamic) num_threads(80)
		#pragma omp parallel for schedule(dynamic) num_threads(80) collapse(3)
		//#pragma omp parallel for schedule(dynamic) num_threads(80) collapse(2)
    //for (int n = 0; n < N; ++n) {
    for (int n = 0; n < Nodesize[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) {
                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;
                }		// for s
              }		// for r
            }	// for c
            output[n * K * OH * OW + k * OH * OW + oh * OW + ow] = o;
          }	// for ow
        }	// for oh
      }	// for k
    }	// for n

//  }


	if(mpi_rank==0 && mpi_world_size == 2) {
			MPI_Recv(&output[Nodesize[0]*K*OH*OW], Nodesize[1]*K*OH*OW, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &status);
	}
	else if(mpi_world_size == 2) {
			MPI_Isend(output, Nodesize[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);

	// set initial value for the cuda and its kernels
  OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
  OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;

	// seperate the jobs for two nodes
	if(mpi_world_size == 2 && _N > 4) Nodesize[1] = _N/2;		// two mpi_rank
	else Nodesize[1] = 0;
	Nodesize[0] = N - Nodesize[1];		// allocate more jobs for the node0 (mpi_rank=0)

	// set more Nodesize for mpi_rank=0
	if(Nodesize[mpi_rank] < MAX_NUM_GPU) {
			num_devices = Nodesize[mpi_rank];
			for(int i=0; i < Nodesize[mpi_rank]; i++) Nsizei[i] = 1;
	} 
	else {
			num_devices = MAX_NUM_GPU;
			int Nodesize_per_GPU = Nodesize[mpi_rank] / MAX_NUM_GPU;
			int remainder= Nodesize[mpi_rank] % MAX_NUM_GPU;

			for(int i=0; i < MAX_NUM_GPU; i++) {
					Nsizei[i] = Nodesize_per_GPU;

					if(i < remainder) Nsizei[i]++;			// increase the more Nsizei for the last term
			}
	}


	// Allocate device memory for each GPU
	for (int i = 0; i < num_devices; i++) {
			CUDA_CALL( cudaSetDevice(i) );
			CUDA_CALL( cudaMalloc(&in_d[i], Nsizei[i]*C*H*W* sizeof(float)) );
			CUDA_CALL( cudaMalloc(&out_d[i], Nsizei[i]*K*OH*OW * sizeof(float)) );
			CUDA_CALL( cudaMalloc(&fil_d[i], K*C*R*S* sizeof(float)) );
	}

}







/*
	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++) {
			Mbegin[i] = (M / num_devices) * i;
			Mend[i] = (M / num_devices) * (i + 1);
	}
	Mend[num_devices - 1] = M;

	// Allocate device memory for each GPU
	for (int i = 0; i < num_devices; i++) {
			CUDA_CALL( cudaSetDevice(i) );
			CUDA_CALL( cudaMalloc(&in_d[i], (Mend[i] - Mbegin[i]) * K * sizeof(float)) );
			CUDA_CALL( cudaMalloc(&out_d[i], K * N * sizeof(float)) );
			CUDA_CALL( cudaMalloc(&fil_d[i], (Mend[i] - Mbegin[i]) * N * sizeof(float)) );
	}

}
*/


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