// 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 30 // 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; 


  double S=0.0;
  int i;
  for(i=0;i<NT;i++){
     __shared__ double shared_A1[TILE][TILE];
     __shared__ double shared_B1[TILE][TILE];


     shared_A1[tidy][tidx]=A.elements[IDX2(row,i*TILE+tidx)];
     shared_B1[tidy][tidx]=B.elements[IDX2(i*TILE+tidy,col)];
     
     __syncthreads();

     for (int k = 0; k < TILE; k++)
            S += shared_A1[tidy][k] * shared_B1[k][tidx];

     __syncthreads();

  }


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

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


  clock_t 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: %f seconds.\n ", ((double)clock() - start)/CLOCKS_PER_SEC );



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