#include <unistd.h>
#include <stdio.h>
#include <cublasLt.h>

#define N 17
#define M 26
#define K 19

void setElement(float *matrix, int nrows, int row, int col, float value)
{ 
  matrix[nrows*col+row] = value;
} 
float getElement(float *matrix, int nrows, int row, int col)
{ 
  return matrix[nrows*col+row];
} 

void check(float alpha, float beta, float* A, int nrowsA, float* B, int nrowsB, float* C, int nrowsC, float *Cnew)
{ 
  int i,j,k;
  float sum;
  float maxerr = 0.0;
  float err;
    
  for(i=0; i<M; i++)
  { 
    for(j=0; j<N; j++) 
    {
      sum = 0.0;
      for(k=0; k<K; k++)
      {
        float aik = getElement(A,nrowsA,i,k);
        float bkj = getElement(B,nrowsB,k,j);
        sum += aik*bkj;
      }
      float cij = getElement(C,nrowsC,i,j);
      sum = sum*alpha + beta*cij;
  
      float cijnew = getElement(Cnew,nrowsC,i,j);
  
      err = abs(double(sum-cijnew));
      if(maxerr < err) maxerr = err; 
    }
  }
  printf("max. error %.15f\n",maxerr);
} 


int main() 
{
  // Host memory
  float Ahost[M*K], Bhost[K*N], Chost[M*N], Dhost[M*N];
  int nrowsA=M, nrowsB=K, nrowsC=M, nrowsD=M;
  int i,j;
  float alpha = 1.0f, beta = 1.0f;

  // Initialize A,B,C with random numbers.
  srand48(clock());
  for(i=0; i<M; i++)
    for(j=0; j<K; j++)
      setElement(Ahost,nrowsA,i,j,drand48());

  for(i=0; i<K; i++)
    for(j=0; j<N; j++)
      setElement(Bhost,nrowsB,i,j,drand48());

  for(i=0; i<M; i++)
    for(j=0; j<N; j++)
      setElement(Chost,nrowsC,i,j,drand48());

  // Device memory
  float *Adevice, *Bdevice, *Cdevice, *Ddevice;
  cudaMalloc(&Adevice, M*K*sizeof(float)); // M rows, K cols
  cudaMalloc(&Bdevice, K*N*sizeof(float)); // K rows, N cols
  cudaMalloc(&Cdevice, M*N*sizeof(float)); // M rows, N cols
  cudaMalloc(&Ddevice, M*N*sizeof(float)); // M rows, N cols

  // Copy matrices to device.
  cudaMemcpy(Adevice, Ahost, M*K*sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(Bdevice, Bhost, K*N*sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(Cdevice, Chost, M*N*sizeof(float), cudaMemcpyHostToDevice);

  // Create cuBLASLt handle.
  cublasLtHandle_t ltHandle;
  cublasLtCreate(&ltHandle);

  // Operation descriptor: FP32 compute, FP32 scale
  cublasLtMatmulDesc_t operationDesc;
  cublasLtMatmulDescCreate
  ( 
    &operationDesc,
    CUBLAS_COMPUTE_32F_FAST_TF32, // Precision for intermediate results.
    CUDA_R_32F           // Precision for alpha and beta.
  );

  // Matrix layouts, column-major.
  cublasLtMatrixLayout_t layoutA, layoutB, layoutC, layoutD;
  cublasLtMatrixLayoutCreate(&layoutA, CUDA_R_32F, M, K, nrowsA);
  cublasLtMatrixLayoutCreate(&layoutB, CUDA_R_32F, K, N, nrowsB);
  cublasLtMatrixLayoutCreate(&layoutC, CUDA_R_32F, M, N, nrowsC);
  cublasLtMatrixLayoutCreate(&layoutD, CUDA_R_32F, M, N, nrowsD);

  // Compute D = alpha*A * B + beta*C
  cublasLtMatmul
  (
    ltHandle,
    operationDesc,
    &alpha,
    Adevice, layoutA,
    Bdevice, layoutB,
    &beta,
    Cdevice, layoutC,
    Ddevice, layoutD,
    NULL, nullptr, 0, 0
  );

  cudaDeviceSynchronize();

  // Copy result.
  cudaMemcpy(Dhost, Ddevice, M*N*sizeof(float), cudaMemcpyDeviceToHost);

  // Check result.
  printf("cuBLASLt result D[0,0] = %f\n",Dhost[0]);
  check(alpha,beta,Ahost, nrowsA, Bhost, nrowsB, Chost, nrowsC, Dhost);


  // Cleanup
  cublasLtMatmulDescDestroy(operationDesc);
  cublasLtMatrixLayoutDestroy(layoutA);
  cublasLtMatrixLayoutDestroy(layoutB);
  cublasLtMatrixLayoutDestroy(layoutC);
  cublasLtMatrixLayoutDestroy(layoutD);
  cublasLtDestroy(ltHandle);

  cudaFree(Adevice); cudaFree(Bdevice); cudaFree(Cdevice); cudaFree(Ddevice);
  return 0;
}