/*
    -- 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 Multimics

extern "C" magma_int_t
magma_zgeqrf2_mmic( magma_int_t num_mics, magma_int_t m, magma_int_t n,
                    magmaDoubleComplex **dlA, magma_int_t ldda,
                    magmaDoubleComplex *tau, 
                    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
    =======
    ZGEQRF2_Mmic computes a QR factorization of a complex M-by-N matrix A:
    A = Q * R. This is a mic interface of the routine.

    Arguments
    =========
    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.

    dA      (input/output) COMPLEX_16 array on the mic, dimension (LDDA,N)
            On entry, the M-by-N matrix dA.
            On exit, the elements on and above the diagonal of the array
            contain the min(M,N)-by-N upper trapezoidal matrix R (R is
            upper triangular if m >= n); the elements below the diagonal,
            with the array TAU, represent the orthogonal matrix Q as a
            product of min(m,n) elementary reflectors (see Further
            Details).

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

    TAU     (output) COMPLEX_16 array, dimension (min(M,N))
            The scalar factors of the elementary reflectors (see Further
            Details).

    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.

    Further Details
    ===============

    The matrix Q is represented as a product of elementary reflectors

       Q = H(1) H(2) . . . H(k), where k = min(m,n).

    Each H(i) has the form

       H(i) = I - tau * v * v'

    where tau is a complex scalar, and v is a complex vector with
    v(1:i-1) = 0 and v(i) = 1; v(i+1:m) is stored on exit in A(i+1:m,i),
    and tau in TAU(i).
    =====================================================================    */

    #define dlA(mic, a_1, a_2) ( dlA[mic] + (a_2) * (ldda) + (a_1))
    #define work_ref(a_1)    ( work + (a_1))
    #define hwork            ( work + (nb) * (m))

    #define hwrk_ref(a_1)    ( local_work + (a_1))
    #define lhwrk            ( local_work + (nb) * (m32m))

	magmaDoubleComplex *dwork[4], *panel[4], *local_work;

	magma_int_t i, j, k, ldwork, lddwork, old_i, old_ib, rows;
	magma_int_t m32m, m32n, nbmin, nx, ib, nb;
	magma_int_t lhwork, lwork;

	//magma_device_t cdevice;
	//magma_getdevice(&cdevice);

	magma_int_t panel_micnum, i_local, n_local[4], la_mic, displacement; 

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

	k = min(m, n);
	if (k == 0)
		return *info;

	nb = magma_get_zgeqrf_nb(m);

	m32m = (m + 31) / 32 * 32;
	m32m = (m32m % 256 == 0) ? m32m + 32 : m32m;
	m32n = (n + 31) / 32 * 32;
	m32n = (m32n % 256 == 0) ? m32n + 32 : m32n;

	displacement = m32n * nb;
	lwork  = (m32m + m32n) * nb;
	lhwork = lwork - (m32m) * nb;

	for(i = 0; i < num_mics; i++){
		magma_setdevice(i);
		if (MAGMA_SUCCESS != magma_zmalloc( &(dwork[i]), (m32n + ldda) * nb )) {
			*info = MAGMA_ERR_DEVICE_ALLOC;
			return *info;
		}
	}

	/* Set the number of local n for each mic */
	for(i = 0; i < num_mics; i++){
		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;
	}

	if (MAGMA_SUCCESS != magma_zmalloc_pinned( &local_work, lwork )) {
		*info = -9;
		for(i = 0; i < num_mics; i++){
			magma_setdevice(i);
			magma_free( dwork[i] );
		}

		*info = MAGMA_ERR_HOST_ALLOC;
		return *info;
	}

	magma_event_t event[4][2];
	for(i = 0; i < num_mics; i++){
		event[i][0] = NULL;
		event[i][1] = NULL;
	}  

	nbmin = 2;
	nx    = nb;
	ldwork = m32m;
	lddwork = m32n;

	if (nb >= nbmin && nb < k && nx < k) {
		/* Use blocked code initially */
		old_i = 0; old_ib = nb;
		for (i = 0; i < (k - nx); i += nb) 
		{
			/* Set the mic number that holds the current panel */
			panel_micnum = (i / nb) % num_mics;

			/* Set the local index where the current panel is */
			i_local = i / (nb * num_mics) * nb;

			ib = min(k - i, nb);
			rows = m - i;

			magma_setdevice(panel_micnum);
			magma_kernel_sync( );

			if (i > 0){
				/* Apply H' to A(i:m, i + 2 * ib:n) from the left; this is the look-ahead
				   application to the trailing matrix                                     */

				/* only the mic that has next panel is done look-ahead */
				magma_setdevice(la_mic);
				magma_zlarfb( MagmaLeft, MagmaConjTrans, MagmaForward, MagmaColumnwise,
						m - old_i, n_local[la_mic] - i_local - old_ib, old_ib,
						panel[la_mic], 0, ldda, dwork[la_mic],  0,    lddwork,
						dlA(la_mic, old_i, i_local + old_ib), 0, ldda, 
						dwork[la_mic], old_ib, lddwork, queue);

				//la_mic = ((i - nb) / nb) % num_mics;
				//magma_setdevice(la_mic);

				//magma_zsetmatrix_async( old_ib, old_ib,
				//		hwrk_ref(old_i), 0, ldwork,
				//		panel[la_mic],   0, ldda, queue, &event[la_mic][0] );
			}

			//magma_setdevice(panel_micnum);
			/* Send current panel to the CPU */
			magma_zgetmatrix_async( rows, ib,
					dlA(panel_micnum, i, i_local), 0, ldda,
					hwrk_ref(i),    0, ldwork, queue, &event[panel_micnum][1]);

			magma_event_sync( event[panel_micnum][1] );

			lapackf77_zgeqrf(&rows, &ib, hwrk_ref(i), &ldwork, tau+i, lhwrk, &lhwork, info);

			// Form the triangular factor of the block reflector
			// H = H(i) H(i+1) . . . H(i+ib-1) 
			lapackf77_zlarft( MagmaForwardStr, MagmaColumnwiseStr,
					&rows, &ib,
					hwrk_ref(i), &ldwork, tau + i, lhwrk, &ib);

			//zpanel_to_q( MagmaUpper, ib, hwrk_ref(i), ldwork, lhwrk + ib * ib );
			// Send the current panel back to the mics 
			// Has to be done with asynchronous copies
			for(j = 0; j < num_mics; j++)
			{  
				if (j == panel_micnum)
					panel[j] = dlA(j, i, i_local);
				else
					panel[j] = dwork[j] + displacement;

				magma_setdevice(j);
				magma_kernel_sync( );

				magma_zsetmatrix_async( rows, ib,
						hwrk_ref(i), 0, ldwork,
						panel[j],    0, ldda, queue, &event[j][0] );
				magma_zsetmatrix_async( ib, ib,
						lhwrk,    0, ib,
						dwork[j], 0, lddwork, queue, &event[j][0] );
			}

			for(j = 0; j < num_mics; j++)
			{
				magma_setdevice(j);
				magma_event_sync( event[j][0] );
			}

			/* Restore the panel */
			//zq_to_panel( MagmaUpper, ib, hwrk_ref(i), ldwork, lhwrk + ib * ib);

			if (i + ib < n) 
			{
				/* Send the T matrix to the mic. 
				   Has to be done with asynchronous copies */
				//for(j = 0; j < num_mics; j++)
				//{
				//	magma_setdevice(j);
				//	magma_zsetmatrix_async( ib, ib,
				//			lhwrk,    0, ib,
				//			dwork[j], 0, lddwork, queue, &event[j][0] );
				//}

				if ((i + nb) < (k - nx))
				{
					/* Apply H' to A(i:m,i+ib:i+2*ib) from the left;
					   This is update for the next panel; part of the look-ahead    */
					la_mic = (panel_micnum + 1) % num_mics;
					int i_loc = (i + nb) / (nb * num_mics) * nb;
					int nmic = 0;

					//for(j = 0; j < num_mics; j++){
					for(j = la_mic; nmic < num_mics; nmic++, j = (j + 1) % num_mics){
						magma_setdevice(j);
						magma_event_sync( event[j][0] );
						if (j == la_mic) {
							magma_zlarfb( MagmaLeft, MagmaConjTrans, MagmaForward, MagmaColumnwise,
									rows, ib, ib,
									panel[j], 0, ldda, dwork[j],    0, lddwork,
									dlA(j, i, i_loc), 0, ldda, dwork[j], ib, lddwork, queue);
						}else if (j <= panel_micnum){
							magma_zlarfb( MagmaLeft, MagmaConjTrans, MagmaForward, MagmaColumnwise,
									rows, n_local[j] - i_local - ib, ib,
									panel[j], 0, ldda, dwork[j],    0, lddwork,
									dlA(j, i, i_local + ib), 0, ldda, dwork[j] + ib, 0, lddwork, queue);
						}else {
							magma_zlarfb( MagmaLeft, MagmaConjTrans, MagmaForward, MagmaColumnwise,
									rows, n_local[j] - i_local, ib,
									panel[j], 0, ldda, dwork[j],    0, lddwork,
									dlA(j, i, i_local), 0, ldda, dwork[j], ib, lddwork, queue);
						}
					}     
				}
				else {
					/* do the entire update as we exit and there would be no lookahead */
					la_mic = (panel_micnum + 1) % num_mics;
					int i_loc = (i + nb) / (nb * num_mics) * nb;

					magma_setdevice(la_mic);
					magma_zlarfb( MagmaLeft, MagmaConjTrans, MagmaForward, MagmaColumnwise,
							rows, n_local[la_mic] - i_loc, ib,
							panel[la_mic], 0, ldda, dwork[la_mic],    0, lddwork,
							dlA(la_mic, i, i_loc), 0, ldda, dwork[la_mic], ib, lddwork, queue);

					magma_setdevice(panel_micnum);
					magma_zsetmatrix( ib, ib,
							hwrk_ref(i), 0,                  ldwork,
							dlA(panel_micnum, i, i_local), 0, ldda, queue );
				}
				old_i  = i;
				old_ib = ib;
			}
		}
	} else {
		i = 0;
	}

	for(j = 0; j < num_mics; j++){
		magma_setdevice(j);
		magma_free( dwork[j] );
	}

	/* Use unblocked code to factor the last or only block. */
	if (i < k) {
		ib   = n - i;
		rows = m - i;
		lhwork = lwork - rows * ib;

		panel_micnum = (panel_micnum + 1) % num_mics;
		int i_loc = (i) / (nb * num_mics) * nb;

		magma_setdevice(panel_micnum);
		magma_zgetmatrix( rows, ib,
				dlA(panel_micnum, i, i_loc), 0, ldda,
				lhwrk, 0, rows, queue);

		lhwork = lwork - rows * ib;
		lapackf77_zgeqrf(&rows, &ib, lhwrk, &rows, tau + i, lhwrk + ib * rows, &lhwork, info);

		magma_zsetmatrix( rows, ib,
				lhwrk, 0,                      rows,
				dlA(panel_micnum, i, i_loc), 0, ldda, queue);
	}

	magma_setdevice(0);
	magma_free_pinned( local_work );

	return *info;
} /* magma_zgeqrf2_mmic */
