////////////////////////////////////////////////////////////////////////////////////////////////
// - Block matrix multiplication algorithm
//
//   const size_t Nblocks = (N / Bsize);
//
//   // loop over blocks of matrix C
//   for (size_t ib=0 ; ib<Nblocks ; ib++)
//   {
//      for (size_t jb=0 ; jb<Nblocks ; jb++)
//      {
//        
//         // loop over blocks of rows of A       
//         for (size_t kb=0 ; kb<Nblocks ; kb++)
//         {
//
//           // Matrix multiplication within a block
//           for (size_t i=(ib * Bsize) ; i<((ib + 1) * Bsize) ; i++)
//             for (size_t j=(jb * Bsize) ; j<((jb + 1) * Bsize) ; j++)
//               for (size_t k=(kb * Bsize) ; k<((kb + 1) * Bsize) ; k++)
//                  C[(i * N) + j] += A[(i * N) + k] * B[(k * N) + j];
//
//         } // kb
//      } // jb
//   } // ib
//
// - Exploit GPU shared-memory.
////////////////////////////////////////////////////////////////////////////////////////////////

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

#include <stdio.h>
#include <string.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>

#include "energy/energy_pmt.h"

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

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

#if BLOCKSIZE > SIZE
#error BLOCKSIZE must be <= SIZE
#endif

#define LOOP 3
#define NDEBUG

double wall_time()
{
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &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;
}

void CPU_mat_mult_block(const MyData *const __restrict__ A,
			const MyData *const __restrict__ B,
	                      MyData *const __restrict__ C,
			const size_t                 size)
{
  const size_t Nblocks = (size / BLOCK_);

  // loop over blocks of matrix C
  for (size_t ib=0 ; ib<Nblocks ; ib++)
    {
      for (size_t jb=0 ; jb<Nblocks ; jb++)
	{       
	  // loop over blocks of rows of A
	  for (size_t kb=0 ; kb<Nblocks ; kb++)
	    {
	      // Matrix multiplication within a block
	      for (size_t i=(ib * BLOCK_) ; i<((ib + 1) * BLOCK_) ; i++)
		for (size_t j=(jb * BLOCK_) ; j<((jb + 1) * BLOCK_) ; j++)
		  for (size_t k=(kb * BLOCK_) ; k<((kb + 1) * BLOCK_) ; k++)
		    C[(i * N) + j] += A[(i * N) + k] * B[(k * N) + j];
	    } // kb
	} // jb
    } // ib
  
  return;
}

__global__ void GPU_mat_mult_block(const MyData *const __restrict__ A,
				   const MyData *const __restrict__ B,
			                 MyData *const __restrict__ C,
				   const size_t                 size)
{
  const size_t size2 = (size * size);
  
  const size_t globalID = threadIdx.x + (blockIdx.x * blockDim.x);
  if (globalID >= size2)
    return;
  
  const size_t Nblocks  = size / BLOCK_;           // number of blocks to loop over A and B
  const size_t ib       = blockIdx.x / Nblocks;    // indexes of starting matrix's block
  const size_t jb       = blockIdx.x % Nblocks;
  const size_t i_local  = threadIdx.x / BLOCK_;    // local matrix's indexes mapped to each CUDA thread
  const size_t j_local  = threadIdx.x % BLOCK_;    // within its own block
  const size_t i_global = i_local + (ib * BLOCK_); // global matrix's indexes mapped to each CUDA thread
  const size_t j_global = j_local + (jb * BLOCK_); // N.B. uncoalescent memory accesses to A and B matrices
  
  C[(i_global * size) + j_global] = (MyData)0;

  // loop over blocks
  for (size_t block=0 ; block<Nblocks ; block++)
    {
      const size_t j_A = (block * BLOCK_);
      const size_t i_B = (block * BLOCK_);

      // perform the matrix multiplication within the block
      MyData value = (MyData)0;
      for (size_t k=0 ; k<BLOCK_ ; k++)
	value += A[(i_global * size) + k + j_A] * B[((k + i_B) * size) + j_global];

      C[(i_global * size) + j_global] += value;
    }

  return;
}

__global__ void GPU_mat_mult_block_shared(const MyData *const __restrict__ A,
					  const MyData *const __restrict__ B,
					        MyData *const __restrict__ C,
					  const size_t                 size)
{
  __shared__ MyData Ablock[BLOCKSIZE];
  __shared__ MyData Bblock[BLOCKSIZE];
  __shared__ MyData Cblock[BLOCKSIZE];
  
  const size_t globalID = threadIdx.x + (blockIdx.x * blockDim.x);
  const size_t size2 = (size * size);
  if (globalID >= size2)
    return;
  
  const size_t Nblocks  = size / BLOCK_;           // number of blocks to loop over
  const size_t ib       = blockIdx.x / Nblocks;    // indexes of starting matrix's block
  const size_t jb       = blockIdx.x % Nblocks;
  const size_t i_local  = threadIdx.x / BLOCK_;    // local matrix's indexes mapped to each CUDA thread
  const size_t j_local  = threadIdx.x % BLOCK_;    // within its own block
  const size_t i_global = i_local + (ib * BLOCK_); // global matrix's indexes mapped to each CUDA thread
  const size_t j_global = j_local + (jb * BLOCK_); // N.B. uncoalescent memory accesses to A and B matrices

  // Init Cblock
  Cblock[(i_local * BLOCK_) + j_local] = (MyData)0;

  // loop over blocks
  for (size_t block=0 ; block<Nblocks ; block++)
    {
      const size_t j_A = (block * BLOCK_);
      const size_t i_B = (block * BLOCK_);

      // waits until all threads in the thread block have reached this point and shared memory accesses
      // made by these threads prior to __syncthreads() are visible to all threads in the block.
      __syncthreads();
      
      // copy block of A into shared memory
      Ablock[(i_local * BLOCK_) + j_local] = A[(i_global * size) + j_local + j_A];
      // copy block of B into shared memory
      Bblock[(i_local * BLOCK_) + j_local] = B[((i_local + i_B) * size) + j_global];
      
      // waits until all threads in the thread block have reached this point and shared memory accesses      // made by these threads prior to __syncthreads() are visible to all threads in the block.
      __syncthreads();

      // perform the matrix multiplication within the block
      MyData value = (MyData)0;
      for (size_t k=0 ; k<BLOCK_ ; k++)
	value += Ablock[(i_local * BLOCK_) + k] * Bblock[(k * BLOCK_) + j_local];

      // store the partial result in shared memory
      Cblock[(i_local * BLOCK_) + j_local] += value;
    }

  C[(i_global * size) + j_global] = Cblock[(i_local * BLOCK_) + j_local];
  
  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");
  
  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();
    }

  // Get the number of compute-capable devices
  int NumGPUs{-1};
  cudaGetDeviceCount(&NumGPUs);
  if (NumGPUs <= 0)
    {
      printf("\n\t Nvidia-GPU is not available \n\n");
      return EXIT_FAILURE;
    }

  int devID = 0;
  // Init PMT
  PMT_CREATE(&devID, 1);
  
  ////////////////////////// CPU naive algorithm //////////////////////////////////////////
  for (unsigned short int loop=0 ; loop<LOOP ; loop++)
    {
      PMT_CPU_START("CPU_mat_mult_block");
      memset(C_CPU, 0, SIZE * sizeof(MyData));
      CPU_mat_mult_block(A_CPU, B_CPU, C_CPU, N);
      PMT_CPU_STOP("CPU_mat_mult_block");
    }
  /////////////////////////////////////////////////////////////////////////////////////////
  
  // copy/alloc data to the GPU
  MyData *buffer_gpu = NULL;
  cudaMalloc((void **)&buffer_gpu, (3 * SIZE * sizeof(MyData)));
  MyData *const A_GPU = buffer_gpu;
  MyData *const B_GPU = A_GPU + SIZE;
  MyData *const C_GPU = B_GPU + SIZE;
  cudaMemcpy(A_GPU, A_CPU, (2 * SIZE * sizeof(MyData)), cudaMemcpyHostToDevice);

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

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

  // free CPU memory
  free(buffer_cpu);
  // free GPU memory
  cudaFree(buffer_gpu);

  printf("\n");

  PMT_CPU_SHOW("CPU_mat_mult_block");
  PMT_GPU_SHOW("GPU_mat_mult_block");
  PMT_GPU_SHOW("GPU_mat_mult_block_shared");
  
  return EXIT_SUCCESS;
}
