#include "mex.h"
#include "cuda.h"

typedef struct{
int width;
int height;
float* elements;
}Matrix;

#define BLOCK_SIZE 16

__global__ void MatMulKernel(Matrix A, Matrix B, Matrix C)
{
float Cvalue=0;
int row=blockIdx.y*blockDim.y+threadIdx.y;
int col=blockIdx.x*blockDim.x+threadIdx.x;

for(int e=0; e<A.width; ++e)
Cvalue += A.elements[row*A.width + e]*B.elements[e*B.width+col];
C.elements[row*C.width + col]=Cvalue;
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
Matrix A,B,C;
Matrix dA,dB,dC;
int dim[2];

A.width=mxGetN(prhs[0]);
A.height=mxGetM(prhs[0]);
A.elements=(float*)mxGetData(prhs[0]);

B.width=mxGetN(prhs[1]);
B.height=mxGetM(prhs[1]);
B.elements=(float*)mxGetData(prhs[1]);

C.width=B.width;
C.height=A.height;

dim[0]=C.height;
dim[1]=C.width;

plhs[0]=mxCreateNumericArray(2,dim,mxSINGLE_CLASS,mxREAL);
C.elements=(float*)mxGetData(plhs[0]);

dA.width=A.width;
dA.height=A.height;
dB.width=B.width;
dB.height=B.height;
dC.width=C.width;
dC.height=C.height;

size_t size=A.width*A.height*sizeof(float);
cudaMalloc((void**)&dA.elements,size);
cudaMemcpy(dA.elements,A.elements,size,cudaMemcpyHostToDevice);

size=B.width*B.height*sizeof(float);
cudaMalloc((void**)&dB.elements,size);
cudaMemcpy(dB.elements,B.elements,size,cudaMemcpyHostToDevice);

size=C.width*C.height*sizeof(float);
cudaMalloc((void**)&dC.elements,size);

dim3 dimBlock(BLOCK_SIZE,BLOCK_SIZE);
dim3 dimGrid(B.width/dimBlock.x,A.height/dimBlock.y);

MatMulKernel<<<dimGrid,dimBlock>>>(dA,dB,dC);
cudaMemcpy(C.elements,dC.elements,size,cudaMemcpyDeviceToHost);

cudaFree(dA.elements);
cudaFree(dB.elements);
cudaFree(dC.elements);
}