//#include <mma.h>
//#include <cuda_fp16.h>
#include <iostream>
#include <stdlib.h>
#include <time.h>

using namespace std;
using namespace nvcuda::wmma;

// Matrix sizes.
// A is MxK, B is KxN, C=AB is MxN.
const int M = 32;
const int N = 8;
const int K = 16;

const int Asize = M*K;
const int Bsize = K*N;
const int Csize = M*N;

// Matrices are stored row major. The stride is the number of columns.
const int Astride = K;
const int Bstride = N;
const int Cstride = N;

// Tile sizes. Simplest case: only one tile for the entire matrix.
// Supported are M=16, N=16, K=16 or M=32, N=8, K=16 or M=8, N=32, K=16
const int Mtile = M;
const int Ntile = N;
const int Ktile = K;

// Kernel which multiplies only one tile.
__global__ void mykernel(half *a, half *b, float *c) 
{
  // Fragments (matrix tiles). Stored in registers in a tensor core.
  // matrix_a, matrix_b, accumulator are constants in WMMA namespace.
  // Templates are used instead of function arguments to have compile time constants.
  fragment<matrix_a, Mtile, Ntile, Ktile, half, row_major> a_frag;
  fragment<matrix_b, Mtile, Ntile, Ktile, half, col_major> b_frag;
  fragment<accumulator, Mtile, Ntile, Ktile, float> c_frag;

  // Initialize accumulator with zeros
  fill_fragment(c_frag, 0.0f);

  // Load A and B into fragments
  // All threads of the warp work together to load the tiles.
  // a,b point to the upper left corner of the tile.
  load_matrix_sync(a_frag, a, Astride, mem_row_major);
  load_matrix_sync(b_frag, b, Bstride, mem_row_major);

  // Do the matrix multiply-accumulate with a tensor core.
  mma_sync(c_frag, a_frag, b_frag, c_frag);

  // Store result back to memory
  store_matrix_sync(c, c_frag, Cstride, mem_row_major);
}

// Access to matrix elements.
template <class T>
T get(int i, int j, int n, T* a)
{
  return (float)a[i*n+j];
}
template <class T>
void set(int i, int j, int n, T* a, T value)
{
  a[i*n+j] = value;
}

// Matrix multiplication on host.
void multiply(half * a, half * b, float * c, int m, int n, int k)
{
  int i,j,l;
  float sum;

  for(i=0; i<m; i++)
  {
    for(j=0; j<n; j++)
    {
      sum = 0.0;
      for(l=0; l<k; l++)
      {
        sum += (float)get(i,l,Astride,a) * (float)get(l,j,Bstride,b);
      }
      set(i,j,Cstride,c,sum);
    }
  }
}

void compare(float * c1, float * c2)
{
  float max = 0.0;
  float diff;
  for(int i=0; i<Csize; i++)
    diff = fabs(c1[i]-c2[i]);
      if(max < diff) max = diff;
  cout << "max. absolute error: " << max << "\n";
}

int main() 
{
  half *a_dev, *b_dev;
  float *c_dev;
  int i;

  // Allocate memory on device
  cudaMalloc(&a_dev, Asize * sizeof(half));
  cudaMalloc(&b_dev, Bsize * sizeof(half));
  cudaMalloc(&c_dev, Csize * sizeof(float));

  // Host memory
  half a_host[Asize], b_host[Bsize];
  float c_host[Csize], check[Csize];

  // Initialize A and B with random numbers.
  srand48(clock());
  for(i=0; i<Asize; i++) a_host[i] = __float2half(drand48());
  for(i=0; i<Bsize; i++) b_host[i] = __float2half(drand48());

  cudaMemcpy(a_dev, a_host, Asize * sizeof(half), cudaMemcpyHostToDevice);
  cudaMemcpy(b_dev, b_host, Bsize * sizeof(half), cudaMemcpyHostToDevice);

  // Launch kernel with 1 warp (32 threads)
  mykernel<<<1, 32>>>(a_dev, b_dev, c_dev);
  cudaDeviceSynchronize();

  cudaMemcpy(c_host, c_dev, Csize * sizeof(float), cudaMemcpyDeviceToHost);

  // Check result.
  multiply(a_host,b_host,check,M,N,K);
  compare(c_host,check);

  cudaFree(a_dev);
  cudaFree(b_dev);
  cudaFree(c_dev);

  return 0;
}