#include "mat_mul.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
#define TSIZE 32
#define WPT 8
#define RSIZE (TSIZE/WPT)

int num_devices = 4;   

__global__ void sgemm(float *A, float *B, float *C, int M, int N, int K) {
    int row = threadIdx.x; 
    int col = threadIdx.y;
    int grow = (blockDim.x*WPT) * blockIdx.x + threadIdx.x;
    int gcolumn = blockDim.y * blockIdx.y + threadIdx.y; 

    __shared__ float Sub_A[TSIZE][TSIZE];
    __shared__ float Sub_B[TSIZE][TSIZE];
  
    float result[WPT];
    
    for (int w=0; w<WPT; w++) {
        result[w] = 0.0f;
    }

    const int t_num = (K+TSIZE-1)/TSIZE;
	for (int t=0; t<t_num; t++) {
        for (int w=0; w<WPT; w++) {
            const int trow = TSIZE*t + row;
            const int tcolumn = TSIZE*t + col;
			if(grow + w*RSIZE >= M || tcolumn >= K) {
			    Sub_A[row + w*RSIZE][col] = 0.0f;   
			} 
            else {
                Sub_A[row + w*RSIZE][col] = A[(grow + w*RSIZE)*K + tcolumn];
			}
			if(trow + w*RSIZE >=K || gcolumn >=N) {
			    Sub_B[row + w*RSIZE][col] = 0.0f;
			}
			else {
                Sub_B[row + w*RSIZE][col] = B[(trow + w*RSIZE)*N + gcolumn];
			}
        }

		__syncthreads();

        for (int k=0; k<TSIZE; k++) {
            for (int w=0; w<WPT; w++) {
                result[w] += Sub_A[row + w*RSIZE][k] * Sub_B[k][col];
            }
        }

        __syncthreads();
    }
 
    for (int w=0; w<WPT; w++) {
	    if (grow + w*RSIZE >= M || gcolumn >=N) continue;
        C[(grow + w*RSIZE)*N + gcolumn] = result[w];
    }
}

// Array of device (GPU) pointers
static float *a_d[MAX_NUM_GPU];
static float *b_d[MAX_NUM_GPU];
static float *c_d[MAX_NUM_GPU];
static int M, N, K;
static int Mbegin[MAX_NUM_GPU], Mend[MAX_NUM_GPU];

void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K) {

    // Launch kernel on every GPU
    for (int i = 0; i < num_devices; i++) {  
        dim3 blockDim(TSIZE/WPT, TSIZE, 1);
        dim3 gridDim(((Mend[i] - Mbegin[i]+TSIZE-1)/TSIZE), (N+TSIZE-1)/TSIZE, 1);

        CUDA_CALL( cudaSetDevice(i) );
        sgemm<<<gridDim, blockDim>>>(a_d[i], b_d[i], c_d[i], Mend[i] - Mbegin[i], N, K);
    }

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

void mat_mul_init(float *A, float *B, float *C, int _M, int _N, int _K) {
    M = _M, N = _N, K = _K;

    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(&a_d[i], (Mend[i] - Mbegin[i]) * K * sizeof(float)) );
        CUDA_CALL( cudaMalloc(&b_d[i], K * N * sizeof(float)) );
        CUDA_CALL( cudaMalloc(&c_d[i], (Mend[i] - Mbegin[i]) * N * sizeof(float)) );
    }

    // Upload A and B matrix to every GPU
    for (int i = 0; i < num_devices; i++) {
        CUDA_CALL( cudaMemcpy(a_d[i], A + Mbegin[i] * K,
                   (Mend[i] - Mbegin[i]) * K * sizeof(float), 
                   cudaMemcpyHostToDevice));
        CUDA_CALL( cudaMemcpy(b_d[i], B, K * N * sizeof(float), cudaMemcpyHostToDevice) );   
    }

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

void mat_mul_final(float *A, float *B, float *C, int M, int N, int K) {
    // Do any post-matmul cleanup work here.

    // Download C matrix from GPUs    
    for (int i = 0; i < num_devices; i++) {
        CUDA_CALL( cudaMemcpy(C + Mbegin[i] * N, c_d[i],
                   (Mend[i] - Mbegin[i]) * N * sizeof(float), 
                   cudaMemcpyDeviceToHost));  
    }

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