/*
!===============================================================================
!
! cublas_wrap.c - cuda version for some lapack functions
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2011 Yoshifumi Nakamura
!
! This file is part of BQCD -- Berlin Quantum ChromoDynamics program
!
! BQCD is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! BQCD is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with BQCD.  If not, see <http://www.gnu.org/licenses/>.
!
!-------------------------------------------------------------------------------
*/
#ifdef NamesToLower
# define INIT_MAGMA init_magma
# define FIN_MAGMA fin_magma
# define WRAP_MAGMA_ZGEMM_GPU  wrap_magma_zgemm_gpu
# define WRAP_MAGMA_ZGETRF_MGPU wrap_magma_zgetrf_mgpu
# define WRAP_MAGMA_ZGETRF_GPU wrap_magma_zgetrf_gpu
# define WRAP_MAGMA_ZGETRI_GPU wrap_magma_zgetri_gpu
#endif
#ifdef NamesToLower_
# define INIT_MAGMA init_magma_
# define FIN_MAGMA fin_magma_
# define WRAP_MAGMA_ZGEMM_GPU  wrap_magma_zgemm_gpu_
# define WRAP_MAGMA_ZGETRF_MGPU wrap_magma_zgetrf_mgpu_
# define WRAP_MAGMA_ZGETRF_GPU wrap_magma_zgetrf_gpu_
# define WRAP_MAGMA_ZGETRI_GPU wrap_magma_zgetri_gpu_
#endif
#ifdef NamesToLower__
# define INIT_MAGMA init_magma__
# define FIN_MAGMA fin_magma__
# define WRAP_MAGMA_ZGEMM_GPU  wrap_magma_zgemm_gpu__
# define WRAP_MAGMA_ZGETRF_MGPU wrap_magma_zgetrf_mgpu__
# define WRAP_MAGMA_ZGETRF_GPU wrap_magma_zgetrf_gpu__
# define WRAP_MAGMA_ZGETRI_GPU wrap_magma_zgetri_gpu__
#endif

#include <cublas.h>
#include <stdio.h>
#define min(a, b) ((a) < (b) ? (a) : (b))


#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <cuda.h>
#include <cublas.h>
#include <stdio.h>
#define ADD_
#include <magma.h>
#include <magma_lapack.h>
#include <testings.h>
/* ========= definition of multiple GPU code ========= */
extern "C" magma_int_t
magma_zgetrf_mgpu(magma_int_t ngpus, 
		  magma_int_t m, magma_int_t n,
		  cuDoubleComplex **d_lA, magma_int_t ldda,
		  magma_int_t *ipiv, magma_int_t *info);
/* =================================================== */


//========== stored data =========================================================
int init_magma_alloc=1;
//cuDoubleComplex *d_B, *d_C;
cuDoubleComplex **d_lA, **d_lB, **d_lC, *dwork;
magma_int_t     *ipiv, ngpus;


//================================================================================
extern "C"  void  INIT_MAGMA(const int* m, const int* n, int* info){
  magma_int_t M = 0, N = 0, lda, ldda;
  magma_int_t min_mn, ldwork, nb, n_local, ldn_local, n_local_mgpu;
  int i;

  *info=0;
  if (init_magma_alloc == 0)return;
  cublasInit();

  M=*m;
  N=*n;
  ldda   = ((M+31)/32)*32;
  min_mn = min(M, N);
  ldwork = N * magma_get_zgetri_nb( N );
  nb     = magma_get_zgetrf_nb(M);
  cudaGetDeviceCount(&ngpus);

  //for test
  //ngpus=1;
  //for test


  n_local_mgpu = ((N/ngpus+1+31)/32)*32;

  /* allocate device memory, assuming fixed nb and ngpus */
  d_lA = new cuDoubleComplex*[ngpus];
  d_lB = new cuDoubleComplex*[ngpus];
  d_lC = new cuDoubleComplex*[ngpus];
  for(i=0; i<ngpus; i++){
    n_local = ((N/nb)/ngpus)*nb;
    if (i < (N/nb)%ngpus)
      n_local += nb;
    else if (i == (N/nb)%ngpus)
      n_local += N%nb;
    ldn_local = ((n_local+31)/32)*32;
    cudaSetDevice(i);
    //TESTING_DEVALLOC( d_lA[i], cuDoubleComplex, ldda*n_local );
    TESTING_DEVALLOC( d_lA[i], cuDoubleComplex, ldda*N ); // change ldn_local -> N to use as d_A
    TESTING_DEVALLOC( d_lB[i], cuDoubleComplex, ldda*n_local_mgpu );
    TESTING_DEVALLOC( d_lC[i], cuDoubleComplex, ldda*n_local_mgpu );
  }
  cudaSetDevice(0);

  TESTING_MALLOC(ipiv, magma_int_t, min_mn);
  TESTING_DEVALLOC(  dwork, cuDoubleComplex, ldwork );
  init_magma_alloc=0;
}

//================================================================================
extern "C" void FIN_MAGMA(){
  free(ipiv);
  cublasFree(dwork);
  for(int i=0; i<ngpus; i++) {
    TESTING_DEVFREE( d_lA[i] );
    TESTING_DEVFREE( d_lB[i] );
    TESTING_DEVFREE( d_lC[i] );
  }
  delete d_lA, d_lB, d_lC;
  cublasShutdown();
  int init_magma_alloc=1;
}

#define MaxGPUs 8
#define MAxStreams 20
//================================================================================
extern "C"     void  WRAP_MAGMA_ZGEMM_GPU(const char *transa, const char *transb, const int *m,
		  const int *n, const int *k, const cuDoubleComplex  *alpha,
		  const cuDoubleComplex  *A, const int *lda,
		  const cuDoubleComplex  *B, const int *ldb, 
		  const cuDoubleComplex  *beta, const cuDoubleComplex *C,
		  const int *ldc)
{
  magma_int_t M, N, K, lN, lK, ldda, lddb;
  M=*m;
  N=*n;
  K=*k;
  ldda   = ((M+31)/32)*32;
  
  if (N%ngpus == 0) {
    lN = N/ngpus;
  } else {
    printf( "ZGEMM_MGPU: N/ngpus is NOT integer\n");
    exit(1);
  }

  if (transb[0] == 'N'){
    lddb = ldda;
  } else {
    if (K%ngpus == 0) {
      lK = K/ngpus;
      lddb = ((lK+31)/32)*32;
    } else {
      printf( "ZGEMM_MGPU: K/ngpus is NOT integer\n");
      exit(1);
    }
  }

  // set matrices
  for (int d = 0; d < ngpus; ++d ) {
    cudaSetDevice(d);
    // set A
    if (transa[0] == 'N'){
      cublasSetMatrix(M, K, 16, A, *lda, d_lA[d], ldda);
    } else {
      cublasSetMatrix(K, M, 16, A, *lda, d_lA[d], ldda);
    }
    // set C
    if ((*beta).x !=0 || (*beta).y !=0) cublasSetMatrix(M, lN, 16, C+d*(*ldc)*lN, *ldc, d_lC[d], ldda);
    // set B
    if (transb[0] == 'N' || ngpus == 1){
      cublasSetMatrix(K, lN, 16, B+d*(*ldb)*lN, *ldb, d_lB[d], ldda);
      cublasZgemm(transa[0], transb[0], M, lN, K, *alpha, d_lA[d], ldda, d_lB[d], ldda, *beta, d_lC[d], ldda);
    } else {
      for (int j=0; j<N; j++) cublasSetMatrix(lK, 1, 16, B+j*(*ldb) +d*lK, *ldb, d_lB[d]+j*lddb, lddb);
      cublasZgemm(transa[0], transb[0], M, lN, K, *alpha, d_lA[d], ldda, d_lB[d], lddb, *beta, d_lC[d], ldda);
    }
  }

  // get C
  for (int d = 0; d < ngpus; ++d ) {
    cudaSetDevice(d);
    cublasGetMatrix(M, lN, 16, d_lC[d], ldda, (void*)(C+d*(*ldc)*lN), *ldc);
  }
  cudaSetDevice(0);

//  1 GPU
//  cudaSetDevice(0);
//  if (transa[0] == 'N'){
//    cublasSetMatrix(M, K, 16, A, *lda, d_lA[0], ldda);
//  } else {
//    cublasSetMatrix(K, M, 16, A, *lda, d_lA[0], ldda);
//  }
//  if (transb[0] == 'N'){
//    cublasSetMatrix(K, N, 16, B, *ldb, d_lB[0], ldda);
//  } else {
//    cublasSetMatrix(N, K, 16, B, *ldb, d_lB[0], ldda);
//  }
//  if ((*beta).x !=0 || (*beta).y !=0) cublasSetMatrix(M, N, 16, C, *ldc, d_lC[0], ldda);
//  cublasZgemm (transa[0], transb[0], M, N, K, *alpha, d_lA[0], ldda, d_lB[0], ldda, *beta, d_lC[0], ldda);
//  cublasGetMatrix(M, N, 16, d_lC[0], ldda, (void*)C, *ldc);


}

//================================================================================
extern "C" void WRAP_MAGMA_ZGETRF_GPU(const int* m, const int* n, cuDoubleComplex *A, const int* lda, int* myipiv, const int *info){
  magma_int_t ldda, nb;
  ldda = ((*m+31)/32)*32;
  nb   = magma_get_zgetrf_nb( min(*m, *n) );

  cudaSetDevice(0);
  magma_zsetmatrix( *m, *n, A, *lda, d_lA[0], ldda );
  magma_zgetrf_gpu( *m, *n, d_lA[0], ldda, ipiv, (magma_int_t*)info);
  if (*info < 0) printf("Argument %d of zgetrf had an illegal value.\n", (int) -(*info));
  magma_zgetmatrix( *m, *n, d_lA[0], ldda, A, *lda );
  for (int i=0; i<*m; i++) myipiv[i]=ipiv[i];
}

//================================================================================
extern "C" void WRAP_MAGMA_ZGETRF_MGPU(const int* m, const int* n, cuDoubleComplex *A, const int* lda, int* myipiv, const int* info){
  magma_int_t M = 0, N = 0, ldda, ngpus0;
  magma_int_t k, nb, nk;
  int j;

  M      = *m;
  N      = *n;
  ldda   = ((M+31)/32)*32;
  nb     = magma_get_zgetrf_nb(M);

  ngpus0 = ngpus;
  if( ngpus0 > N/nb ) {
    ngpus = N/nb;
    if( N%nb != 0 ) ngpus ++;
    printf( " * too many GPUs for the matrix size, using %d GPUs\n", (int) ngpus );
  } else {
    ngpus = ngpus0;
  }

  for(j=0; j<N; j+=nb){
    k = (j/nb)%ngpus;
    cudaSetDevice(k);
    nk = min(nb, N-j);
    magma_zsetmatrix( M, nk,
                      A+j* (*lda),                  *lda,
                      d_lA[k]+j/(nb*ngpus)*nb*ldda, ldda );
  }
  cudaSetDevice(0);

  /* == calling MAGMA with multiple GPUs == */
  magma_zgetrf_mgpu( ngpus, M, N, d_lA, ldda, ipiv, (magma_int_t*)info);
  if (*info < 0) {
    printf("Argument %d of magma_zgetrf_mgpu had an illegal value.\n", (int) -(*info));
  } else if (*info != 0 ) {
    printf("magma_zgetrf_mgpu returned info=%d.\n", (int) *info);
  }

  for(j=0; j<N; j+=nb){
    k = (j/nb)%ngpus;
    cudaSetDevice(k);
    nk = min(nb, N-j);
    magma_zgetmatrix( M, nk,
                      d_lA[k]+j/(nb*ngpus)*nb*ldda, ldda,
                      A+j*(*lda),                   *lda );
  }
  cudaSetDevice(0);
  for (int i=0; i<*m; i++) myipiv[i]=ipiv[i];
}

//================================================================================
extern "C" void WRAP_MAGMA_ZGETRI_GPU(const int* n, cuDoubleComplex *A, const int* lda, int* myipiv, const int *info) {
  magma_int_t ldda, ldwork;

  ldda = ((*n+31)/32)*32;
  ldwork = *n * magma_get_zgetri_nb( *n );

  cudaSetDevice(0);
  magma_zsetmatrix( *n, *n, A, *lda, d_lA[0], ldda );
  magma_zgetri_gpu( *n,     d_lA[0], ldda, ipiv, dwork, ldwork, (magma_int_t*)info );
  if (*info != 0) printf( "An error occured in magma_zgetri, info=%d\n", (int) *info );
  magma_zgetmatrix( *n, *n, d_lA[0], ldda, A, *lda );
}
