// Multiply two matrices A * B = C
 
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <time.h>

#include "MersenneTwister.h"
#include "book.h"



#define TILE 20 // tile size
#define NT 200 // number of tiles
#define MS TILE*NT //matrix size

// row-major storing
#define IDX2(i,j) (((i)*(MS))+(j))


// Matrices are stored in row-major order
typedef struct {
    int width;
    int height;
    double* elements;
} Matrix;





__global__ void matrixMul(Matrix A, Matrix B, Matrix C){

  int tidx=threadIdx.x;
  int tidy=threadIdx.y;

  int bidx=blockIdx.x;
  int bidy=blockIdx.y;

  int row = bidy*TILE+tidy;
  int col = bidx*TILE+tidx; 


  int i;
  double S=0.0;
  for(i=0;i<MS;i++){
    S += A.elements[IDX2(row,i)]*B.elements[IDX2(i,col)];
  }
  C.elements[IDX2(row,col)] = S;
}


int main(void){

  MTRand mt;
  Matrix A;
  Matrix B;
  Matrix C;

  A.width=MS;
  A.height=MS;
  A.elements = (double*) malloc(A.width*A.height*sizeof(double));

  B.width=MS;
  B.height=MS;
  B.elements = (double*) malloc(B.width*B.height*sizeof(double));

  C.width=MS;
  C.height=MS;
  C.elements = (double*) malloc(C.width*C.height*sizeof(double));

  int i,j;
  for(i=0;i<MS;i++)
    for(j=0;j<MS;j++){
      A.elements[IDX2(i,j)]=(double) mt();
      B.elements[IDX2(i,j)]=(double) mt();
    }

  int k;
  double S=0.0;
  clock_t start;


  /*
  start=clock();
  for(i=0;i<MS;i++){
    for(j=0;j<MS;j++){
      S=0;
      for(k=0;k<MS;k++){
	S += A.elements[IDX2(i,k)]*B.elements[IDX2(k,j)];
      }
    }
  }
  printf("Time to compute on CPU: %f seconds.\n ", ((double)clock() - start)/CLOCKS_PER_SEC );
  */



  Matrix dev_A;
  dev_A.width = A.width; 
  dev_A.height = A.height;
  HANDLE_ERROR(cudaMalloc((void **)&dev_A.elements, dev_A.width*dev_A.height*sizeof(double)));

  Matrix dev_B;
  dev_B.width = B.width; 
  dev_B.height = B.height;
  HANDLE_ERROR(cudaMalloc((void **)&dev_B.elements, dev_B.width*dev_B.height*sizeof(double)));

  Matrix dev_C;
  dev_C.width = C.width; 
  dev_C.height = C.height;
  HANDLE_ERROR(cudaMalloc((void **)&dev_C.elements, dev_C.width*dev_C.height*sizeof(double)));


  start=clock();
  
  HANDLE_ERROR(cudaMemcpy(dev_A.elements, A.elements, A.width*A.height*sizeof(double), cudaMemcpyHostToDevice));
  HANDLE_ERROR(cudaMemcpy(dev_B.elements, B.elements, B.width*B.height*sizeof(double), cudaMemcpyHostToDevice));

  dim3 threads(TILE, TILE);
  dim3 grid(NT,NT);

  matrixMul<<<grid,threads>>>(dev_A,dev_B,dev_C);

  HANDLE_ERROR(cudaMemcpy(C.elements, dev_C.elements, C.width*C.height*sizeof(double), cudaMemcpyDeviceToHost));

  printf("Time to compute on GPU: %f seconds.\n ", ((double)clock() - start)/CLOCKS_PER_SEC );

  /*
  for(i=0;i<MS;i++){
    for(j=0;j<MS;j++){
      printf("%f ", C.elements[IDX2(i,j)]);
    }
    printf("\n");
  }
  */


  free(A.elements);
  free(B.elements);
  free(C.elements);
  HANDLE_ERROR(cudaFree(dev_A.elements));
  HANDLE_ERROR(cudaFree(dev_B.elements));
  HANDLE_ERROR(cudaFree(dev_C.elements));
  return 1;
}