////////////////////////////////////////////////////////////////////////////////////////////////
// - Naive matrix multiplication algorithm
//   for (size_t i=0 ; i<N ; i++)
//      for (size_t j=0 ; j<N ; j++)
//         for (size_t k=0 ; k<_N ; k++)
//            C[(i * N) + j] += A[(i * N) + k] * B[(k * N) + j];
//
////////////////////////////////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////////////////////////////////
// Author: David Goz
// mail  : david.goz@inaf.it
// date  : 20.08.2024
// code tested using nvhpc
//
// - Compile the code:
//   $ nvc++ mat_mult.cu -o mat_mult
// - Run the code:
//   $ ./mat_mult
//////////////////////////////////////////////////////////////////////////////////////////////////

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <assert.h>
#include <cuda.h>
#include <string.h>
#include <math.h>
#include <float.h>

#define N                     1024
#define SIZE                  (N * N) // matrix size
typedef double MyData;                // do not change
#define BLOCKSIZE             1024

// sanity check
#if BLOCKSIZE > 1024
#error BLOCKSIZE must be <= 1024
#endif

#define LOOP 100
#define NDEBUG

double wall_time()
{
  struct timespec ts;
  clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &ts);
  const double ret = (double) (ts.tv_sec) + (double) ts.tv_nsec * 1.0e-9;

  return ret;
}

void CPU_mat_mult(const MyData *const restrict A,
		  const MyData *const restrict B,
	                MyData *const restrict C,
		  const size_t                 size)
{
  for (size_t i=0 ; i<size ; i++)
    for (size_t j=0 ; j<size ; j++)
      for (size_t k=0 ; k<size ; k++)
	C[(i * size) + j] += (A[(i * size) + k] * B[(k * size) + j]);

  return;
}

__global__ void GPU_mat_mult(const MyData *const restrict A,
			     const MyData *const restrict B,
			           MyData *const restrict C,
			     const size_t                 size)
{
  const size_t globalID = threadIdx.x + (blockIdx.x * blockDim.x);
  if (globalID >= SIZE)
    return;
  
  const size_t i = globalID / size;
  const size_t j = globalID % size;

  MyData value = (MyData)0;
  for (size_t k=0 ; k<size ; k++)
    value += A[(i * N) + k] * B[(k * N) + j];

  C[(i * N) + j] = value;
  
  return;
}

void check(const MyData *const __restrict__ cpu_matrix,
	   const MyData *const __restrict__ gpu_matrix)
{
  int flag = 0;
  for (size_t i=0 ; i<SIZE ; i++)
    flag = ((fabs(cpu_matrix[i] - gpu_matrix[i]) > FLT_EPSILON) ? 1 : flag);

  if (!flag)
    printf("\n\t Result OK");
  else
    printf("\n\t Result wrong");

#if !defined(NDEBUG)

  for (size_t i=0 ; i<N ; i++)
    {
      printf("\n");
      for (size_t j=0 ; j<N ; j++)
	{
	  const size_t index = ((i * N) + j);
	  printf("\n\t cpu_matrix[%d] = %lg - gpu_matrix[%d] = %lg - diff = %lg",
		 index, cpu_matrix[index], index, gpu_matrix[index], fabs(cpu_matrix[index] - gpu_matrix[index]));
	}
    }
  printf("\n");

#endif // NDEBUG  
  
  return;
}

int main()
{
  double time;
  MyData *buffer_cpu = (MyData *)calloc(4 * SIZE, sizeof(MyData));
  assert(buffer_cpu != NULL);

  // host reference matrix A
  MyData *const restrict A_CPU = buffer_cpu;
  MyData *const restrict B_CPU = A_CPU + SIZE;
  MyData *const restrict C_CPU = B_CPU + SIZE;
  MyData *const restrict C_GPU_host = C_CPU + SIZE;
  for (size_t i=0 ; i<SIZE ; i++)
    {
      A_CPU[i] = drand48();
      B_CPU[i] = drand48();
    }

  ////////////////////////// CPU naive algorithm //////////////////////////////////////////
  CPU_mat_mult(A_CPU, B_CPU, C_CPU, N);
  /////////////////////////////////////////////////////////////////////////////////////////

  // copy/alloc data to the GPU
  MyData *buffer_gpu = NULL;
  cudaMalloc((void **)&buffer_gpu, (3 * SIZE * sizeof(MyData)));
  MyData *const restrict A_GPU = buffer_gpu;
  MyData *const restrict B_GPU = A_GPU + SIZE;
  MyData *const restrict C_GPU = B_GPU + SIZE;
  cudaMemcpy(A_GPU, A_CPU, (2 * SIZE * sizeof(MyData)), cudaMemcpyHostToDevice);

  const dim3 block   = {BLOCKSIZE, 1, 1};
  const dim3 nblocks = {(SIZE + BLOCKSIZE  -1)/BLOCKSIZE, 1, 1};
  
  /////////////////////////// GPU naive algorithm ////////////////////////////////////////
  time = 0.0;
  GPU_mat_mult<<< nblocks, block >>>(A_GPU, B_GPU, C_GPU, N); // warm-up
  cudaDeviceSynchronize();
  for (unsigned short int loop=0 ; loop<LOOP ; loop++)
    {
      const double start = wall_time();
      GPU_mat_mult<<< nblocks, block >>>(A_GPU, B_GPU, C_GPU, N);
      cudaDeviceSynchronize();
      time += (wall_time() - start);
    }
  cudaMemcpy(C_GPU_host, C_GPU, (SIZE * sizeof(MyData)), cudaMemcpyDeviceToHost);
  
  check(C_CPU, C_GPU_host);
  printf("\n\t GPU naive time %lg [s]\n", (time / LOOP));
  /////////////////////////////////////////////////////////////////////////////////////

  // free CPU memory
  free(buffer_cpu);
  // free GPU memory
  cudaFree(buffer_gpu);
  
  printf("\n");

  return EXIT_SUCCESS;
}
