/* Illustration of dense matrix inversion using CUDA cuBLAS library 
   Sergey Voronin, 01.22 
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <cuda_runtime.h>
#include "cublas_v2.h"

int main (void){
    cublasStatus_t stat;
    cublasHandle_t handle;
    int i, j, n, lda, batchsize = 1;
    if (cudaStat != cudaSuccess) {
        printf ("device memory allocation failed");
        return EXIT_FAILURE;
    }
    stat = cublasCreate(&handle);
    if (stat != CUBLAS_STATUS_SUCCESS) {
        printf ("CUBLAS initialization failed\n");
        return EXIT_FAILURE;
    }

    /* define matrix */
    n = 3; lda = 3;
    float a_mat[] = {
    1,   0,   3,
    4,   6,   5,
    1,   5,   9 };
    
    /* single array to hold matrix, device version, and device version of double array */
    float *am = (float *)malloc(n*n*sizeof(float));
    float *am_d; cudaMalloc(&am_d, n*n*sizeof(float));
    float **am_dd; cudaMalloc(&am_dd, sizeof(float *));

    /* initialize mat */
    for(i=0; i<n; i++){
     for(j=0; j<n; j++){
        am[i*n+j] = a_mat[i*n + j];
    }
    }

    printf("copying matrix to gpu..\n");
    cudaMemcpy(am_d, am, n*n*sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(am_dd, &am_d, sizeof(float *), cudaMemcpyHostToDevice);

  printf("print mat for inverse:\n");
  for (i = 0; i < n; i++) {
        for (j = 0; j < n; j++) {
            printf("%4.3f ", (a_mat[i*n+j]));
        }
	printf("\n");
  } 

  /* init option structures and perform inverse using two step solver, transferring data to and from device as appropriate */
  int *p = (int *)malloc(sizeof(int));
  p[0] = 0;
  int *p_d; cudaMalloc(&p_d, sizeof(int));
  cudaMemcpy(p_d, p, sizeof(int), cudaMemcpyHostToDevice);
  int *info = (int *)malloc(sizeof(int));
  info[0] = 0;
  int *info_d; cudaMalloc(&info_d, sizeof(int));
  cudaMemcpy(info_d, info, sizeof(int), cudaMemcpyHostToDevice);
  stat = cublasSgetrfBatched(handle, n, am_dd, lda, p_d, info_d, batchsize);
  cudaDeviceSynchronize();

  cudaMemcpy(p, p_d, sizeof(int), cudaMemcpyDeviceToHost);
  cudaMemcpy(info, info_d, sizeof(int), cudaMemcpyDeviceToHost);

  /* init output structures */
  float *c = (float *)malloc(n*n*sizeof(float));
  float *c_d; cudaMalloc(&c_d, n*n*sizeof(float));
  float **c_dd; cudaMalloc(&c_dd, sizeof(float *));

  cudaMemcpy(c_d, c, n*n*sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(c_dd, &c_d, sizeof(float*), cudaMemcpyHostToDevice);

  stat = cublasSgetriBatched(handle, n, am_dd, lda, p_d, c_dd, lda, info_d, batchsize);
  cudaDeviceSynchronize();
  //printf("stat = %d\n", stat);
  cudaMemcpy(&c_d, c_dd, sizeof(float*), cudaMemcpyDeviceToHost);
  cudaMemcpy(c, c_d, 9*sizeof(float), cudaMemcpyDeviceToHost);

   printf("printing inverse stored in c after cublasSgetriBatched..\n");
   for (i = 0; i < n; i++) {
        for (j = 0; j < n; j++) {
            printf("%4.3f  ", c[i*n+j]);
        }
	printf("\n");
    }

    /* free and exit */
    cublasDestroy(handle);
    return EXIT_SUCCESS;
}