// MP 3: Due Sunday, Dec 30, 2012 at 11:59 p.m. PST
#include    <wb.h>

#define TILE_SIZE 16
#define wbCheck(stmt) do {                                 \
        cudaError_t err = stmt;                            \
        if (err != cudaSuccess) {                          \
            wbLog(ERROR, "Failed to run stmt ", #stmt);    \
            return -1;                                     \
        }                                                  \
    } while(0)

// Compute C = A * B
__global__ void matrixMultiplyShared(float * A, float * B, float * C,
			             int numARows, int numAColumns,
			             int numBRows, int numBColumns,
			             int numCRows, int numCColumns) {

    //@@ Insert code to implement matrix multiplication here
    //@@ You have to use shared memory for this MP

    __shared__ float tileA[TILE_SIZE][TILE_SIZE];
    __shared__ float tileB[TILE_SIZE][TILE_SIZE];

	int col = blockIdx.x * TILE_SIZE + threadIdx.x;
	int row = blockIdx.y * TILE_SIZE + threadIdx.y;
	float pValue = 0;

	for(int m = 0; m < (numAColumns + 1) / TILE_SIZE -1; m++)
	{


		if(m * TILE_SIZE + threadIdx.x < numAColumns && row < numARows)
			tileA[threadIdx.y][threadIdx.x] = A[row * numAColumns + m * TILE_SIZE + threadIdx.x];
		else
			tileA[threadIdx.y][threadIdx.x] = 0;

		if(col && m * TILE_SIZE + threadIdx.y < numBRows)
			tileB[threadIdx.y][threadIdx.x] = B[(m * TILE_SIZE + threadIdx.y) * numBColumns + col];
		else
			tileB[threadIdx.y][threadIdx.x] = 0;

		__syncthreads();

		for(int k = 0; k < TILE_SIZE; k++)
			pValue += tileA[threadIdx.y][k] * tileB[k][threadIdx.x];

		__syncthreads();
	}

	if(row < numCRows && col < numCColumns)
		C[row * numCColumns + col] = pValue;

}

int main(int argc, char ** argv) {
    wbArg_t args;
    float * hostA; // The A matrix
    float * hostB; // The B matrix
    float * hostC; // The output C matrix
    float * deviceA;
    float * deviceB;
    float * deviceC;
    int numARows; // number of rows in the matrix A
    int numAColumns; // number of columns in the matrix A
    int numBRows; // number of rows in the matrix B
    int numBColumns; // number of columns in the matrix B
    int numCRows; // number of rows in the matrix C (you have to set this)
    int numCColumns; // number of columns in the matrix C (you have to set this)

    args = wbArg_read(argc, argv);

    wbTime_start(Generic, "Importing data and creating memory on host");
    hostA = (float *) wbImport(wbArg_getInputFile(args, 0), &numARows, &numAColumns);
    hostB = (float *) wbImport(wbArg_getInputFile(args, 1), &numBRows, &numBColumns);


    //@@ Set numCRows and numCColumns
    if(numAColumns != numBRows)
    {
    	wbLog(TRACE, "Incompatible matrices");
    	return 1;
    }

    numCRows = numARows;
    numCColumns = numBColumns;


    //@@ Allocate the hostC matrix
    int sizeOfA = numAColumns * numARows * sizeof(float);
    int sizeOfB = numBColumns * numBRows * sizeof(float);
  	int sizeOfC = numCColumns * numCRows * sizeof(float);

  	hostC = (float *)malloc(sizeOfC);

    wbTime_stop(Generic, "Importing data and creating memory on host");

    wbLog(TRACE, "The dimensions of A are ", numARows, " x ", numAColumns);
    wbLog(TRACE, "The dimensions of B are ", numBRows, " x ", numBColumns);

    wbTime_start(GPU, "Allocating GPU memory.");



    //@@ Allocate GPU memory here
    wbCheck(cudaMalloc((void**)&deviceA, sizeOfA));
  	wbCheck(cudaMalloc((void**)&deviceB, sizeOfB));
  	wbCheck(cudaMalloc((void**)&deviceC, sizeOfC));


    wbTime_stop(GPU, "Allocating GPU memory.");

    wbTime_start(GPU, "Copying input memory to the GPU.");


    //@@ Copy memory to the GPU here
    wbCheck(cudaMemcpy(deviceA, hostA, sizeOfA, cudaMemcpyHostToDevice));
    wbCheck(cudaMemcpy(deviceB, hostB, sizeOfB, cudaMemcpyHostToDevice));



    wbTime_stop(GPU, "Copying input memory to the GPU.");


  	//@@ Initialize the grid and block dimensions here
    dim3 dimBlock(TILE_SIZE, TILE_SIZE, 1);
  	dim3 dimGrid((numCColumns - 1)/dimBlock.x + 1, (numCRows - 1)/dimBlock.y + 1, 1);



    wbTime_start(Compute, "Performing CUDA computation");


  	//@@ Launch the GPU Kernel here
    matrixMultiplyShared<<<dimGrid, dimBlock>>>(deviceA, deviceB, deviceC,
			       numARows, numAColumns,
			       numBRows, numBColumns,
			       numCRows, numCColumns);

    cudaThreadSynchronize();
    wbTime_stop(Compute, "Performing CUDA computation");

    wbTime_start(Copy, "Copying output memory to the CPU");


    //@@ Copy the GPU memory back to the CPU here
    wbCheck(cudaMemcpy(hostC, deviceC, sizeOfC, cudaMemcpyDeviceToHost));


    wbTime_stop(Copy, "Copying output memory to the CPU");

    wbTime_start(GPU, "Freeing GPU Memory");


  	//@@ Free the GPU memory here
    wbCheck(cudaFree(deviceA));
    wbCheck(cudaFree(deviceB));
  	wbCheck(cudaFree(deviceC));



    wbTime_stop(GPU, "Freeing GPU Memory");

    wbSolution(args, hostC, numCRows, numCColumns);

    free(hostA);
    free(hostB);
    free(hostC);

    return 0;
}