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

       @precisions normal z -> s d c

*/

#include <math.h>
#include <stdio.h>
#include "common_magma.h"


magma_int_t
magma_zgetrf_mmic(magma_int_t num_mics, 
                 magma_int_t m, magma_int_t n, 
                 magmaDoubleComplex_ptr *d_lA, magma_int_t ldda,
                 magma_int_t *ipiv, 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
    =======

    ZGETRF computes an LU factorization of a general M-by-N matrix A
    using partial pivoting with row interchanges.

    The factorization has the form
       A = P * L * U
    where P is a permutation matrix, L is lower triangular with unit
    diagonal elements (lower trapezoidal if m > n), and U is upper
    triangular (upper trapezoidal if m < n).

    This is the right-looking Level 3 BLAS version of the algorithm.

    Arguments
    =========

    NUM_MICS 
            (input) INTEGER
            The number of MICS to be used for the factorization.

    M       (input) INTEGER
            The number of rows of the matrix A.  M >= 0.

    N       (input) INTEGER
            The number of columns of the matrix A.  N >= 0.

    A       (input/output) COMPLEX_16 array on the MIC, dimension (LDDA,N).
            On entry, the M-by-N matrix to be factored.
            On exit, the factors L and U from the factorization
            A = P*L*U; the unit diagonal elements of L are not stored.

    LDDA     (input) INTEGER
            The leading dimension of the array A.  LDDA >= max(1,M).

    IPIV    (output) INTEGER array, dimension (min(M,N))
            The pivot indices; for 1 <= i <= min(M,N), row i of the
            matrix was interchanged with row IPIV(i).

    INFO    (output) INTEGER
            = 0:  successful exit
            < 0:  if INFO = -i, the i-th argument had an illegal value
                  or another error occured, such as memory allocation failed.
            > 0:  if INFO = i, U(i,i) is exactly zero. The factorization
                  has been completed, but the factor U is exactly
                  singular, and division by zero will occur if it is used
                  to solve a system of equations.
    =====================================================================    */

#define inA(id, i, j) (d_lA[(id)] + (j) * nb * ldda + (i) * nb)
	magmaDoubleComplex c_one     = MAGMA_Z_ONE;
	magmaDoubleComplex c_neg_one = MAGMA_Z_NEG_ONE;

	magma_int_t err, iinfo, nb, n_local[MagmaMaxDevices];
	magma_int_t mindim;
	magma_int_t i, j, d, rows, s, lddat, lddwork;
	magma_int_t id, i_local, i_local2, nb0, nb1;
	magmaDoubleComplex *d_lAT[MagmaMaxDevices], *d_lAP[MagmaMaxDevices];
	magmaDoubleComplex *d_panel[MagmaMaxDevices], *work, *next_panel, *panel_local[MagmaMaxDevices];
	magma_int_t *d_ipiv[MagmaMaxDevices];

	magma_event_t event[MagmaMaxDevices][2];

	/* Check arguments */
	*info = 0;
	if (m < 0)
		*info = -2;
	else if (n < 0)
		*info = -3;
	else if (ldda < max(1,m))
		*info = -5;

	if (*info != 0) {
		magma_xerbla( __func__, -(*info) );
		return *info;
	}

	/* Quick return if possible */
	if (m == 0 || n == 0)
		return *info;

	/* Function Body */
	mindim = min(m, n);
	nb     = magma_get_zgetrf_nb(m);

	if (nb <= 1 || nb >= n) {
		/* Use CPU code. */
		err = magma_zmalloc_pinned( &work, m * n );
		if ( work == NULL ) {
			*info = MAGMA_ERR_HOST_ALLOC;
			return *info;
		}
		magma_zgetmatrix( m, n, d_lA[0], 0, ldda, work, 0, m, queue);
		lapackf77_zgetrf(&m, &n, work, &m, ipiv, info);
		magma_zsetmatrix( m, n, work, 0, m, d_lA[0], 0, ldda, queue);
		magma_free_pinned(work);
	} else {
		/* Use hybrid blocked code. */
		if( num_mics > ceil((double)n / nb) ) {
			printf( " * too many MICs for the matrix size, using %d MICs\n", (int) num_mics );
			*info = -1;
			return *info;
		}

		/* allocate workspace for each MIC */
		for(i = 0; i < num_mics; i++){

			magma_setdevice(i);

			/* local-n and local-ld */
			n_local[i] = (( n / nb) / num_mics) * nb;
			if (i < (n / nb) % num_mics)
				n_local[i] += nb;
			else if (i == (n / nb) % num_mics)
				n_local[i] += n % nb;
			panel_local[i] = NULL;

			/* workspaces */
			if (MAGMA_SUCCESS != magma_zmalloc( &d_lAP[i], ldda * nb )) {
				for( j = 0; j < i; j++ ) {
					magma_setdevice(j);
					magma_free( d_lAP[j]   );
					magma_free( d_panel[j] );
				}

				panel_local[i] = inA(i, 0, 0);
				*info = MAGMA_ERR_DEVICE_ALLOC;
				return *info;
			}

			if (MAGMA_SUCCESS != magma_imalloc( &d_ipiv[i], mindim )) {
				for( j = 0; j <= i; j++ ) {
					magma_setdevice(j);
					magma_free( d_lAP[j] );
				}
				for( j = 0; j < i; j++ ) {
					magma_setdevice(j);
					magma_free( d_ipiv[j] );
				}

				*info = MAGMA_ERR_DEVICE_ALLOC;
				return *info;
			}
		}

		/* cpu workspace */
		lddwork = ldda;
		if (MAGMA_SUCCESS != magma_zmalloc_pinned( &work, lddwork * nb )) {
			for(i = 0; i < num_mics; i++ ) {
				magma_setdevice(i);
				magma_free( d_lAP[i]   );
			}
			*info = MAGMA_ERR_HOST_ALLOC;
			return *info;
		}

		if (MAGMA_SUCCESS != magma_zmalloc_pinned( &next_panel, lddwork * nb )) {
			magma_free_pinned(work);
			for(i = 0; i < num_mics; i++ ) {
				magma_setdevice(i);
				magma_free( d_lAP[i]   );
			}
			*info = MAGMA_ERR_HOST_ALLOC;
			return *info;
		}


		s = mindim / nb;
		for( i = 0; i < s; i++ )
		{
			/* Set the MIC number that holds the current panel */
			id = i % num_mics;
			magma_setdevice(id);
			
			/* Set the local index where the current panel is */
			i_local = i / num_mics;
			rows  = m - i * nb;

			/* make sure that mic queue is empty */
			magma_kernel_sync();

			/* the remaining updates */
			if ( i>0 ){
				/* id-th mic update the remaining matrix */
				if((n_local[id] - (i_local + 1) * nb) > 0) {
					magma_ztrsmgemm( MagmaLeft, MagmaLower, MagmaNoTrans, MagmaUnit, 
							nb, n_local[id] - (i_local + 1) * nb, c_one,
							panel_local[id], 0,       ldda,
							inA(id, i - 1, i_local + 1), 0, ldda, 
							MagmaNoTrans, MagmaNoTrans, 
							m - i * nb, n_local[id] - (i_local + 1) * nb, nb, 
							c_neg_one, panel_local[id] + nb,          0, ldda, 
							inA(id, i - 1, i_local + 1), 0, ldda, 
							c_one,     inA(id, i,  i_local + 1),           0, ldda, queue );
				}


				/* cpu update id-th mic panel that is used next by cpu*/
				/* start sending the panel to cpu */
				magma_zgetmatrix_async(m - (i - 1) * nb, nb,
						inA(id, i - 1, i_local), 0, ldda,
						next_panel,      0, lddwork, queue, &event[id][1] );

				/* stnchrnoize i-th panel from id-th mic into work */
				magma_event_sync(event[id][1] );

				//update of look ahead panel on cpu
				ztrsm_( "L", "L", "N", "U", 
						&nb, &nb, 
						&c_one, work, &lddwork, 
						next_panel, &lddwork);

				//rows = m - i * nb;
				zgemm_( "N", "N", 
						&rows, &nb, &nb, 
						&c_neg_one, work + nb, &lddwork,
						next_panel, &lddwork, 
						&c_one,     next_panel + nb, &lddwork );

				magma_zsetmatrix_async(m - (i - 1) * nb, nb,
						next_panel,    0, lddwork,
						inA(id, i - 1, i_local), 0, ldda, queue, &event[id][0] );

				magma_event_sync(event[id][0] );

			}

			/* start sending the panel to cpu */
			magma_zgetmatrix_async( rows, nb,
					inA(id, i, i_local), 0, ldda,
					work,      0, lddwork, queue, &event[id][1] );

			/* stnchrnoize i-th panel from id-th mic into work */
			magma_event_sync(event[id][1] );

			/* i-th panel factorization */
			lapackf77_zgetrf( &rows, &nb, work, &lddwork, ipiv + i * nb, &iinfo);

			if ( (*info == 0) && (iinfo > 0) ) {
				*info = iinfo + i * nb;
				break;
			}
			
			//transfer pivot array
			for(d = 0; d < num_mics; d++) {
				magma_setdevice(d);
				magma_zsetintarray_async(mindim, ipiv, 0, d_ipiv[d], 0, queue, &event[id][1]);
			}

			for(d = 0; d < num_mics; d++) {
				magma_setdevice(d);
				magma_event_sync(event[d][1]);
				magma_zlaswp(n_local[d], inA(d, i, 0), ldda, 1, nb, d_ipiv[d] + i * nb, 1, queue);
			}

			for (int jj = 0; jj < nb; jj++) {
				*(ipiv + i * nb + jj) += i * nb;
			}

			for( d = 0; d < num_mics; d++ ) {
				/* storage for panel */
				if( d == id ) {
					/* the panel belong to this mic */
					panel_local[d] = inA(d, i, i_local);
				} else {
					/* the panel belong to another mic */
					panel_local[d] = d_lAP[d];
				}
			}

			/* start sending the panel to all the mics */
			for( d = 0; d < num_mics; d++ ) {
				magma_setdevice(d);
				magma_kernel_sync();
				magma_zsetmatrix_async( rows, nb,
						work,    0, lddwork,
						panel_local[d], 0, ldda, queue, &event[d][0] );
			}


			for( d = 0; d < num_mics; d++ ) {
				magma_setdevice(d);
				/* storage for panel */
				if( d == id ) {
					/* next column */
					i_local2 = i_local + 1;
				} else {
					/* next column */
					i_local2 = i_local;
					if( d < id ) i_local2++;
				}
				/* the size of the next column */
				if ( s > (i + 1) ) {
					nb0 = nb;
				} else {
					nb0 = n_local[d] - nb * (s / num_mics);
					if( d < s % num_mics ) nb0 -= nb;
				}

				/* synchronization */
				magma_event_sync( event[d][0] );

				if( d == (i + 1) % num_mics) {
					/* owns the next column, look-ahead the column */
					nb1 = nb0;
				} else {
					/* update the entire trailing matrix */
					nb1 = n_local[d] - i_local2 * nb;
				}

				if( (d == (i + 1) % num_mics) && s > (i + 1)) {
				} else {
					/* mic updating the trailing matrix */
					if(nb1 > 0) {
						magma_ztrsmgemm( MagmaLeft, MagmaLower, MagmaNoTrans, MagmaUnit, 
								nb, nb1, c_one,
								panel_local[d], 0,       ldda,
								inA(d, i, i_local2), 0, ldda,
								MagmaNoTrans, MagmaNoTrans, 
								m - (i + 1) * nb, nb1, nb, 
								c_neg_one, panel_local[d] + nb, 0,         ldda,
								inA(d, i, i_local2), 0, ldda, 
								c_one,     inA(d, i + 1, i_local2), 0,        ldda, queue );
					}
				}

			} /* end of mic updates */
		} /* end of for i=1..s */

		/* Set the MIC number that holds the last panel */
		id = s % num_mics;

		/* Set the local index where the last panel is */
		i_local = s / num_mics;

		/* size of the last diagonal-block */
		nb0 = min(m - s * nb, n - s * nb);
		rows = m    - s * nb;

		if( nb0 > 0 ) {
			/* send the last panel to cpu */
			magma_setdevice(id);
			//magma_kernel_sync();
			magma_zgetmatrix( rows, nb0, inA(id, s, i_local), 0, ldda, work, 0, lddwork, queue );

			/* factor on cpu */
			lapackf77_zgetrf( &rows, &nb0, work, &lddwork, ipiv + s * nb, &iinfo);
			if ( (*info == 0) && (iinfo > 0) )
				*info = iinfo + s * nb;

			//transfer pivot array
			for(d = 0; d < num_mics; d++) {
				magma_setdevice(d);
				magma_zsetintarray_async(mindim, ipiv, 0, d_ipiv[d], 0, queue, &event[id][1]);
			}

			for(d = 0; d < num_mics; d++) {
				magma_setdevice(d);
				magma_event_sync(event[d][1]);
				magma_zlaswp(n_local[d], inA(d, s, 0), ldda, 1, nb0, d_ipiv[d] + s * nb, 1, queue);
			}

			for (int jj = 0; jj < nb0; jj++) {
				*(ipiv + s * nb + jj) += s * nb;
			}


			magma_setdevice(id);
			magma_kernel_sync();

			magma_zsetmatrix_async( rows, nb0,
					work,    0, lddwork,
					inA(id, s, i_local), 0, ldda, queue, &event[id][0] );

			/* send the factor to mics */
			for( d = 0; d < num_mics; d++ ) {
				magma_setdevice(d);
				magma_kernel_sync();
				i_local2 = i_local;
				if( d < id ) i_local2 ++;

				if( d == id || n_local[d] > i_local2 * nb ) 
				{
					magma_zsetmatrix_async( rows, nb0,
							work,     0, lddwork,
							d_lAP[d], 0, ldda, queue, &event[d][0] );
				}
			}
		}

		/* clean up */
		for( d = 0; d < num_mics; d++ ) {
			magma_setdevice(d);
			magma_kernel_sync();

			if( nb0 > 0 ) {

				i_local2 = i_local;
				if( d < id ) 
					i_local2++;

				if( d == id ) {
					/* next column */
					nb1 = n_local[d] - i_local * nb - nb0;

					magma_event_sync( event[d][0] );

					if( nb1 > 0 ) {
						magma_ztrsm( MagmaLeft, MagmaLower, MagmaNoTrans, MagmaUnit, 
							nb0, nb1, c_one,
								d_lAP[d],        0, ldda, 
								inA(d, s, i_local) + nb0 * ldda, 0, ldda, queue);
					}

				} else if( n_local[d] > i_local2 * nb ) {
					/* next column */
					nb1 = n_local[d] - i_local2 * nb;

					magma_event_sync( event[d][0] );
					if(nb1 > 0) 
						magma_ztrsm( MagmaLeft, MagmaLower, MagmaNoTrans, MagmaUnit, 
								nb0, nb1, c_one,
								d_lAP[d],  0, ldda, 
								inA(d, s, i_local2), 0, ldda, queue);
				}
			}

			/* save on output */
			magma_setdevice(d);
			magma_free( d_lAP[d]   );
			magma_free( d_ipiv[d] );
		} /* end of for d=1,..,num_mics */

		magma_free_pinned( work );
		magma_free_pinned( next_panel );
	}

	magma_setdevice(0);

	return *info;       
	/* End of MAGMA_ZGETRF_MMIC */
}

#undef inAT
