/*
    -- MAGMA (version 1.2.1) --
       Univ. of Tennessee, Knoxville
       Univ. of California, Berkeley
       Univ. of Colorado, Denver
       June 2012

       @precisions normal z -> s d c

*/
#include "common_magma.h"

#define dlA(id, i, j)  (d_lA[id] + (j) * ldda + (i))
#define dlP(id, i, j)  (d_lP[id] + (j) * ldda + (i))

#define dlAT(id, i, j)  (d_lA[id] + (j) * lddat + (i))
#define dlPT(id, i, j)  (d_lP[id] + (j) * nb    + (i))

extern "C" magma_int_t
magma_zpotrf_mmic(int num_mics, magma_uplo_t uplo, magma_int_t n, 
                 magmaDoubleComplex **d_lA, magma_int_t ldda, magma_int_t *info, magma_queue_t queue )
{
/*  -- MAGMA (version 1.2.1) --
       Univ. of Tennessee, Knoxville
       Univ. of California, Berkeley
       Univ. of Colorado, Denver
       June 2012

    Purpose   
    =======   
    ZPOTRF computes the Cholesky factorization of a complex Hermitian   
    positive definite matrix dA.   

    The factorization has the form   
       dA = U**H * U,  if UPLO = 'U', or   
       dA = L  * L**H,  if UPLO = 'L',   
    where U is an upper triangular matrix and L is lower triangular.   

    This is the block version of the algorithm, calling Level 3 BLAS.   

    Arguments   
    =========   
    UPLO    (input) CHARACTER*1   
            = 'U':  Upper triangle of dA is stored;   
            = 'L':  Lower triangle of dA is stored.   

    N       (input) INTEGER   
            The order of the matrix dA.  N >= 0.   

    dA      (input/output) COMPLEX_16 array on the mic, dimension (LDDA,N)   
            On entry, the Hermitian matrix dA.  If UPLO = 'U', the leading   
            N-by-N upper triangular part of dA contains the upper   
            triangular part of the matrix dA, and the strictly lower   
            triangular part of dA is not referenced.  If UPLO = 'L', the   
            leading N-by-N lower triangular part of dA contains the lower   
            triangular part of the matrix dA, and the strictly upper   
            triangular part of dA is not referenced.   

            On exit, if INFO = 0, the factor U or L from the Cholesky   
            factorization dA = U**H * U or dA = L * L**H.   

    LDDA     (input) INTEGER   
            The leading dimension of the array dA.  LDDA >= max(1,N).
            To benefit from coalescent memory accesses LDDA must be
            dividable by 16.

    INFO    (output) INTEGER   
            = 0:  successful exit   
            < 0:  if INFO = -i, the i-th argument had an illegal value   
            > 0:  if INFO = i, the leading minor of order i is not   
                  positive definite, and the factorization could not be   
                  completed.   
    =====================================================================   */

	magma_int_t     j, jb, nb, nb0, nb2, d, id, j_local, j_local2;
	magmaDoubleComplex c_one     = MAGMA_Z_ONE;
	magmaDoubleComplex c_neg_one = MAGMA_Z_NEG_ONE;
	magmaDoubleComplex *work, *tmpwork;
	double          d_one     =  1.0;
	double          d_neg_one = -1.0;
	magmaDoubleComplex *d_lP[MagmaMaxDevices], *dlpanel;
	magma_int_t n_local[MagmaMaxDevices], lddat_local[MagmaMaxDevices], lddat, ldpanel;
	magma_event_t event[MagmaMaxDevices][4];

	*info = 0;
	if ( uplo != MagmaUpper && uplo != MagmaLower ) {
		*info = -1;
	} else if (n < 0) {
		*info = -2;
	} else if (ldda < max(1,n)) {
		*info = -4;
	}
	if (*info != 0) {
		magma_xerbla( __func__, -(*info) );
		return *info;
	}

	nb = magma_get_zpotrf_nb(n);

	if (MAGMA_SUCCESS != magma_zmalloc_pinned( &work, n * nb )) {
		*info = MAGMA_ERR_HOST_ALLOC;
		return *info;
	}

	if (MAGMA_SUCCESS != magma_zmalloc_pinned( &tmpwork, n * nb )) {
		*info = MAGMA_ERR_HOST_ALLOC;
		return *info;
	}


	if ((nb <= 1) || (nb >= n)) {
		/*  Use unblocked code. */
		magma_setdevice(0);
		magma_zgetmatrix( n, n, dlA(0, 0, 0), 0, ldda, work, 0, n, queue);
		if(uplo == MagmaLower)
			lapackf77_zpotrf( MagmaLowerStr, &n, work, &n, info);
		else
			lapackf77_zpotrf( MagmaUpperStr, &n, work, &n, info);
		magma_zsetmatrix( n, n, work, 0, n, dlA(0, 0, 0), 0, ldda, queue );
	} else {

		for( d = 0; d < num_mics; d++ ) {
			magma_setdevice(d);

			/* local-n and local-ld */
			n_local[d] = ((n / nb) / num_mics) * nb;
			if (d < (n / nb) % num_mics)
				n_local[d] += nb;
			else if (d == (n / nb) % num_mics)
				n_local[d] += n % nb;
			lddat_local[d] = ((n_local[d] + 31) / 32) * 32;
			lddat_local[d] = (lddat_local[d] % 256 == 0) ? lddat_local[d] + 32 : lddat_local[d];

			if (MAGMA_SUCCESS != magma_zmalloc( &d_lP[d], nb * ldda )) {
				for( j = 0; j < d; j++ ) {
					magma_setdevice(j);
					magma_free( d_lP[d] );
				}
				*info = MAGMA_ERR_DEVICE_ALLOC;
				return *info;
			}

			event[d][0] = NULL;
			event[d][1] = NULL;
			event[d][2] = NULL;
			event[d][3] = NULL;
		}

		/* Use blocked code. */
		if (uplo == MagmaUpper) 
		{     
			/* Compute the Cholesky factorization A = U'*U. */
			for (j = 0; j < n; j += nb) {

				/* Set the mic number that holds the current panel */
				id = (j / nb) % num_mics;
				magma_setdevice(id);

				/* Set the local index where the current panel is */
				j_local = j / (nb * num_mics);
				jb = min(nb, (n - j));

				if( j > 0 && (j + jb) < n) {
					/* wait for the off-diagonal column off the current diagonal *
					 * and send it to mics                                       */
					magma_event_sync( event[id][2] );
					for( d = 0; d < num_mics; d++ ) {
						magma_setdevice(d);
						magma_kernel_sync();
							if( d != id )
								magma_zsetmatrix_async( j, jb,
										&work[jb],   0, n,
										dlP(d, jb, 0), 0, ldda, queue, &event[d][3] );
					}
				}

				/* Update the current diagonal block */
				magma_setdevice(id);
				if(j > 0)
					magma_zherk(MagmaUpper, MagmaConjTrans, jb, j, 
							d_neg_one, dlA(id, 0, nb*j_local), 0, ldda, 
							d_one,     dlA(id, j, nb*j_local), 0, ldda, queue);

				magma_kernel_sync();

				/* send the diagonal to cpu */
				magma_zgetmatrix_async( jb, jb,
						dlA(id, j, nb * j_local), 0, ldda,
						work, 0, n, queue, &event[id][0] );

				if ( j > 0 && (j + jb) < n) {
					/* Compute the local block column of the panel. */
					for( d = 0; d < num_mics; d++ ) {
						magma_setdevice(d);
						j_local2 = j_local+1;
						if( d > id ) j_local2 --;

						/* wait for the off-diagonal */
						if( d != id ) {
							magma_event_sync( event[id][3] );
							dlpanel = dlP(d, jb, 0);
						} else {
							dlpanel = dlA(d, 0, nb*j_local);
						}

						/* update the panel */
						magma_zgemm(MagmaConjTrans, MagmaNoTrans, 
								jb, (n_local[d] - nb * (j_local2 - 1) - jb), j, 
								c_neg_one, dlpanel, 0, ldda, 
								dlA(d, 0, nb * j_local2), 0, ldda,
								c_one,     dlA(d, j, nb * j_local2), 0, ldda, queue);
					}
				}

				/* factor the diagonal */
				magma_setdevice(id);
				magma_event_sync( event[id][0] );
				lapackf77_zpotrf(MagmaUpperStr, &jb, work, &n, info);
				if (*info != 0) {
					*info = *info + j;
					break;
				}

				/* send the diagonal to mics */
				if ( (j + jb) < n) {
					for( d = 0; d < num_mics; d++ ) {
						magma_setdevice(d);
						if( d == id ) dlpanel = dlA(d, j, nb * j_local);
						else          dlpanel = dlP(d, 0, 0);

						magma_zsetmatrix_async( jb, jb,
								work,    0, n,
								dlpanel, 0, ldda, queue, &event[d][1] );
					}
				} else {
					magma_setdevice(id);
					magma_zsetmatrix_async( jb, jb,
							work,   0, n,
							dlA(id, j, nb * j_local), 0, ldda, queue, &event[id][1] );
				}

				/* panel-factorize the off-diagonal */
				if ( (j + jb) < n) {
						d = (j / nb + 1) % num_mics; 
						magma_setdevice(d);

						/* next column */
						j_local2 = j_local + 1;
						if( d > id ) j_local2--;
						if( d == id ) dlpanel = dlA(d, j, nb * j_local);
						else          dlpanel = dlP(d, 0, 0);
						nb0 = min(nb, n_local[d] - nb * j_local2 );

						magma_event_sync( event[d][1] );
						/* owns the next column, look-ahead the column */
						magma_ztrsm( MagmaLeft, MagmaUpper, MagmaConjTrans, MagmaNonUnit, 
								jb, nb0, c_one,
								dlpanel, 0, ldda, 
								dlA(d, j, nb * j_local2), 0, ldda, queue);
				}

				if ( (j + jb) < n) {
					for( d = 0; d < num_mics; d++ ) {
						magma_setdevice(d);

						/* next column */
						j_local2 = j_local + 1;
						if( d > id ) j_local2--;
						if( d == id ) dlpanel = dlA(d, j, nb * j_local);
						else          dlpanel = dlP(d, 0, 0);
						nb0 = min(nb, n_local[d] - nb * j_local2 );

						magma_event_sync( event[d][1] );
						if( d == (j / nb + 1) % num_mics ) {
							/* update the remaining blocks */
							   nb2 = n_local[d] - j_local2 * nb - nb0;
						} else {
							/* update the entire trailing matrix */
							nb2 = n_local[d] - j_local2 * nb;
							nb0 = 0;
						}

						magma_ztrsm( MagmaLeft, MagmaUpper, MagmaConjTrans, MagmaNonUnit, 
								jb, nb2, c_one,
								dlpanel, 0,                   ldda, 
								dlA(d, j, nb * j_local2 + nb0), 0, ldda, queue);
						magma_kernel_sync();
					}
				} /* end of ztrsm */
				if ( (j + jb) < n) {
					d = (j / nb + 1) % num_mics; 
					magma_setdevice(d);

					/* next column */
					j_local2 = j_local + 1;
					if( d > id ) j_local2--;
					if( d == id ) dlpanel = dlA(d, j, nb * j_local);
					else          dlpanel = dlP(d, 0, 0);
					nb0 = min(nb, n_local[d] - nb * j_local2 );

					/* send the column to cpu */
					if( j + jb + nb0 < n ) {
						magma_zgetmatrix_async( (j + jb), nb0,
								dlA(d, 0, nb * j_local2), 0, ldda,
								&work[nb0], 0,             n, queue, &event[d][2] );
					}
				}
			} /* end of for j=1, .., n */
		} else { 
			/* Compute the Cholesky factorization A = L*L'. */
			magma_int_t count = 0;

			for (j = 0; j < n; j += nb) {
				/* Set the mic number that holds the current panel */
				id = (j / nb) % num_mics;
				magma_setdevice(id);
				lddat = lddat_local[id];

				/* Set the local index where the current panel is */
				j_local = j / (nb * num_mics);
				jb = min(nb, (n - j));
				

				/* Update the current diagonal block */
				if( j > 0) {
					magma_zherk(MagmaLower, MagmaNoTrans, jb, j,
							d_neg_one, dlAT(id, nb * j_local, 0), 0, lddat,
							d_one,     dlAT(id, nb * j_local, j), 0, lddat, queue);
					magma_kernel_sync();
				}

				//if ( (j + jb) < n) {
				if ( j > 0 && j < n) {
					lddat = lddat_local[id];
					magma_setdevice(id);

					/* next column */
					j_local2 = j_local;
					nb0 = min(nb, n_local[id] - nb * j_local2 );

					/* send the column to cpu */
					if(j + nb0 < n && nb0) {
						magma_zgetmatrix_async( nb0, j,
								dlAT(id, nb * j_local2, 0), 0, lddat,
								&work[nb * nb], 0, nb, queue, &event[id][2] );
						magma_event_sync( event[id][2] );
					}
				}

				if( j > 0 && (j + jb) < n) {
					/* wait for the off-diagonal row off the current diagonal *
					 * and send it to mics                                    */
					count = 0;
					d = id;
					for( ; count < num_mics; count++, d = (d + 1) % num_mics ) {
						magma_setdevice(d);
						lddat = lddat_local[d];
						if( d != id ) { 
							magma_kernel_sync();
							magma_zsetmatrix_async( jb, j,
									&work[nb * nb], 0, nb,
									dlPT(d, 0, jb), 0, nb, queue, &event[d][3] );
						}
					}
				}

				if ( j > 0 && (j + jb) < n) {
					/* compute the block-rows of the panel */
					count = 0;
					d = (id + 1) % num_mics;
					for( ; count < num_mics; count++, d = (d + 1) % num_mics ) {
						magma_setdevice(d);
						lddat = lddat_local[d];
						j_local2 = j_local + 1;
						if( d > id ) j_local2 --;

						/* wait for the off-diagonal */
						if( d != id ) {
							magma_event_sync( event[id][3] );
							dlpanel = dlPT(d, 0, jb);
							ldpanel = nb;
						} else {
							dlpanel = dlAT(d, nb * j_local, 0);
							ldpanel = lddat;
						}
						
						if(d == id)
							magma_kernel_sync();

						/* update the panel */
						if((n_local[d] - nb * j_local2) > 0) {
							magma_zgemm( MagmaNoTrans, MagmaConjTrans,
									n_local[d] - nb * j_local2, jb, j,
									c_neg_one, dlAT(d, nb * j_local2, 0), 0, lddat,
									dlpanel, 0, ldpanel,
									c_one,     dlAT(d, nb * j_local2, j), 0, lddat, queue);
						} 

					}
				}

				/* factor the diagonal */
				magma_setdevice(id);
				lddat = lddat_local[id];

				/* send the diagonal to cpu */
				magma_zgetmatrix_async( jb, jb,
						dlAT(id, nb * j_local, j), 0, lddat,
						work, 0, nb, queue, &event[id][0] );
				
				magma_event_sync( event[id][0] );

				lapackf77_zpotrf(MagmaLowerStr, &jb, work, &nb, info);
				if (*info != 0) {
					*info = *info + j;
					break;
				}

				/* send the diagonal to mics */
				if ( (j + jb) < n) {
					magma_int_t count = 0;
					d = (j / nb + 1) % num_mics;
					for( ; count < num_mics; count++, d = (d + 1) % num_mics ) {
						magma_setdevice(d);
						lddat = lddat_local[d];
						if( d == id ) {
							dlpanel = dlAT(d, nb * j_local, j);
							ldpanel = lddat;
						} else {
							dlpanel = dlPT(d, 0, 0);
							ldpanel = nb;
						}

						magma_zsetmatrix_async( jb, jb,
								work,    0, nb,
								dlpanel, 0, ldpanel, queue, &event[d][1] );
					}
				} else {
					magma_zsetmatrix_async( jb, jb,
							work, 0,                   nb,
							dlAT(id, nb * j_local, j), 0, lddat, queue, &event[id][1] );
					magma_event_sync( event[d][1] );
				}

				if ( (j + jb) < n) {
					magma_int_t count = 0;
					d = (j / nb + 1) % num_mics;
					for( ; count < num_mics; count++, d = (d + 1) % num_mics ) {
						magma_setdevice(d);
						lddat = lddat_local[d];

						/* next column */
						j_local2 = j_local + 1;
						if( d > id ) j_local2--;
						if( d == id ) {
							dlpanel = dlAT(d, nb * j_local, j);
							ldpanel = lddat;
						} else {         
							dlpanel = dlPT(d, 0, 0);
							ldpanel = nb;
						}

						magma_event_sync( event[d][1] );

						/* update the entire trailing matrix */
						nb2 = n_local[d] - j_local2 * nb;

						if( nb2 > 0)
							magma_ztrsm( MagmaRight, MagmaLower, MagmaConjTrans, MagmaNonUnit, 
									nb2, jb, c_one,
									dlpanel, 0, ldpanel, 
									dlAT(d, nb * j_local2, j), 0, lddat, queue);
					}

				}
			}
		} /* end of else not upper */

		/* clean up */
		for( d=0; d<num_mics; d++ ) {
			magma_setdevice(d);
			magma_free( d_lP[d] );
		}

	} /* end of not lapack */

	/* free workspace */
	magma_free_pinned( work );
	magma_free_pinned( tmpwork );

	return *info;
} /* magma_zpotrf_mmic */
