#include "mat_mul.h"

#include <cstdio>
#include <cuda_runtime.h>

#define CUDA_CALL(d, f)                                                                       \
  {                                                                                           \
    cudaError_t err = (f);                                                                    \
    if (err != cudaSuccess) {                                                                 \
      fprintf(stderr, "[GPU %d] CUDA error at [%s:%d] %d %s\n",                               \
              (d), __FILE__, __LINE__, err, cudaGetErrorString(err));                         \
      exit(1);                                                                                \
    }                                                                                         \
  }

#define MAX_NUM_GPU 4
int num_devices = 0;

// A: M rows, K columns
// B: K rows, N columns
// C: M rows, N columns
//                         
//                   N     
//                o-----o  
//                |     |  
//              K | [B] |  
//                |     |  
//                o-----o  
//        K          N     
//    o-------o   o-----o  
//  M |  [A]  | M | [C] |  
//    |       |   |     |  
//    o-------o   o-----o  
//                         

#define TS_M    64                                  // The tile-size in dimension M
#define TS_N    64                                  // The tile-size in dimension N
#define TS_K    64                                  // The tile-size in dimension K
#define WPT_M   16                                  // The amount of work-per-thread in dimension M
#define WPT_N   8                                   // The amount of work-per-thread in dimension N

#define CEIL_DIV(x,y) ( ((x) + (y) - 1) / (y) )
#define CEIL(x,y)     ( CEIL_DIV((x),(y)) * (y) )

__global__ void sgemm(float *A, float *B, float *C, int M, int N, int K)
{
  // Thread identifiers
  const int row = threadIdx.x;                                  // Local row ID (max: TS_M/WPT_M)
  const int col = threadIdx.y;                                  // Local col ID (max: TS_N/WPT_N)
  const int globalRow = TS_M * blockIdx.x + row * WPT_M;  // Row ID of C (0..M)
  const int globalCol = TS_N * blockIdx.y + col * WPT_N;  // Col ID of C (0..N)

  // printf("[R%03d, C%03d] GR=%04d GC=%04d\n", row, col, globalRow, globalCol);

  // Local memory to fit a tile of TS*TS elements of A and B
  __shared__ float Asub[TS_M][TS_K];
  __shared__ float Bsub[TS_K][TS_N];

  // Initialize the accumulation registers
  float acc[WPT_M][WPT_N];
  for (int wm = 0; wm < WPT_M; wm++)
  {
      for (int wn = 0; wn < WPT_N; wn++)
      {
          acc[wm][wn] = 0.0f;
      }
  }

  // Loop over all tiles
  const int numTiles = CEIL_DIV(K, TS_K);
  for (int t = 0; t < numTiles; t++)
  {
      const int rowInB = TS_M * t + row * WPT_M;
      const int colInA = TS_N * t + col * WPT_N;

      // Load one tile of A and B into local memory
      #pragma unroll
      for (int wm = 0; wm < WPT_M; wm++)
      {
          #pragma unroll
          for (int wn = 0; wn < WPT_N; wn++)
          {
              int r, c;

              r = globalRow + wm;
              c = colInA + wn;
              Asub[row * WPT_M + wm][col * WPT_N + wn] = (r >= M || c >= K) ? 0.0f : A[r * K + c];

              r = rowInB + wm;
              c = globalCol + wn;
              Bsub[row * WPT_M + wm][col * WPT_N + wn] = (r >= K || c >= N) ? 0.0f : B[r * N + c];
          }
      }

      // Synchronize to make sure the tile is loaded
      __syncthreads();;

      // Loop over the values of a single tile
      for (int k = 0; k < TS_K; k++)
      {
          // Cache the values of Bsub in registers
          float bs[WPT_N];

          #pragma unroll
          for (int wn = 0; wn < WPT_N; wn++) {
              bs[wn] = Bsub[k][col * WPT_N + wn];
          }

          // Perform the computation
          #pragma unroll
          for (int wm = 0; wm < WPT_M; wm++)
          {
              float a = Asub[row * WPT_M + wm][k];

              #pragma unroll
              for (int wn = 0; wn < WPT_N; wn++)
              {
                  acc[wm][wn] += a * bs[wn];
              }
          }
      }

      // Synchronize before loading the next tile
      __syncthreads();;
  }

  // Store the final results in C
  #pragma unroll
  for (int wm = 0; wm < WPT_M; wm++)
  {
      #pragma unroll
      for (int wn = 0; wn < WPT_N; wn++)
      {
          int r = globalRow + wm;
          int c = globalCol + wn;

          if (r < M && c < N)
          {
              C[r * N + c] = acc[wm][wn];
          }
      }
  }
}

// Array of device (GPU) pointers
static float *a_d[MAX_NUM_GPU];
static float *b_d[MAX_NUM_GPU];
static float *c_d[MAX_NUM_GPU];
static int M, N, K;
static int Mbegin[MAX_NUM_GPU], Mend[MAX_NUM_GPU];

void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K) {

  // Launch kernel on every GPU
  for (int i = 0; i < num_devices; i++) {
    // dim3 blockDim(1, 1, 1);
    // dim3 gridDim(Mend[i] - Mbegin[i], N, 1);
    dim3 blockDim(CEIL_DIV(TS_M, WPT_M), CEIL_DIV(TS_N, WPT_N), 1);
    dim3 gridDim(CEIL_DIV(Mend[i] - Mbegin[i], TS_M), CEIL_DIV(N, TS_N), 1);

    CUDA_CALL( i, cudaSetDevice(i) );
    sgemm<<<gridDim, blockDim>>>(a_d[i], b_d[i], c_d[i], Mend[i] - Mbegin[i], N, K);
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaSetDevice(i) );
    CUDA_CALL( i, cudaDeviceSynchronize() );
  }
}

void mat_mul_init(float *A, float *B, float *C, int _M, int _N, int _K) {
  M = _M, N = _N, K = _K;

  CUDA_CALL( 0, cudaGetDeviceCount(&num_devices) );
  if (num_devices > MAX_NUM_GPU)
  {
    num_devices = MAX_NUM_GPU;
  }

  printf("Using %d devices\n", num_devices);
  for (int i = 0; i < num_devices; i++) {
    cudaDeviceProp prop;
    CUDA_CALL( i, cudaGetDeviceProperties(&prop, i) );

    // Try printing more detailed information here
    printf("[GPU %d] %s\n", i, prop.name);
  }

  if (num_devices <= 0) {
    printf("No CUDA device found. Aborting\n");
    exit(1);
  }

  // Setup problem size for each GPU
  for (int i = 0; i < num_devices; i++) {
    Mbegin[i] = (M / num_devices) * i;
    Mend[i] = (M / num_devices) * (i + 1);
  }
  Mend[num_devices - 1] = M;

  // Allocate device memory for each GPU
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaSetDevice(i) );
    CUDA_CALL( i, cudaMalloc(&a_d[i], (Mend[i] - Mbegin[i]) * K * sizeof(float)) );
    CUDA_CALL( i, cudaMalloc(&b_d[i], K * N * sizeof(float)) );
    CUDA_CALL( i, cudaMalloc(&c_d[i], (Mend[i] - Mbegin[i]) * N * sizeof(float)) );
  }

  // Upload A and B matrix to every GPU
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaSetDevice(i) );
    CUDA_CALL( i, cudaMemcpy(a_d[i], A + Mbegin[i] * K,
               (Mend[i] - Mbegin[i]) * K * sizeof(float),
               cudaMemcpyHostToDevice) );
    CUDA_CALL( i, cudaMemcpy(b_d[i], B, K * N * sizeof(float), cudaMemcpyHostToDevice) );
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaSetDevice(i) );
    CUDA_CALL( i, cudaDeviceSynchronize() );
  }
}

void mat_mul_final(float *A, float *B, float *C, int M, int N, int K) {
  // Do any post-matmul cleanup work here.

  // Download C matrix from GPUs
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaMemcpy(C + Mbegin[i] * N, c_d[i],
               (Mend[i] - Mbegin[i]) * N * sizeof(float),
               cudaMemcpyDeviceToHost) );
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( i, cudaDeviceSynchronize() );
  }
}