#include "../include/lass.h"

/**
 *
 * 	@file kdppgetrf.c
 *
 * 	@brief LASs-DDSs kdppgetrf routine.
 *
 * 	LASs-DDSs is a software package provided by:
 * 	Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
 *
 * 	@author Pedro Valero-Lara pedro.valero@bsc.es
 *      @author Davor Davidovic davor.davidovic@irb.hr
 * 	@date 2019-04-26
 * 	@reviewer 
 * 	@modified 2019-05-07
 *
 **/

/**
 *  
 *	@ingroup DDSS
 *   
 *	Performs the LU factorization with partial pivoting ( row interchanges )
 *	of a general M-by-N matrix A:
 *
 *		A = P * L * U
 *
 *	where P is a permutation matrix, L is a lower triangular 
 *	( lower trapezoidal if M > N ) matrix with unit diagonal elements and
 *	U is an upper triangular ( upper trapezoidal if M < N ) matrix.
 *	The factorization is done in panels, while updates are performed in tiles
 *	
 *	This version improves the kdppgetrf_ipiv_depend version by removing the
 *	dependencies between LU (dgetrf) and swap tasks. 
**/

/**
 *	
 *	@param[in]
 *	M       int.
 *          M specifies the number of rows of the matrix A. M >= 0.
 *	
 *	@param[in]
 *	N       int.
 *          N specifies the number of columns of the matrix A. N >= 0.
 *
 *	@param[in,out]
 *  A  		double *.
 *	    	A is a pointer to a regular matrix of dimension M-by-N.
 *          On exit, if return value is Success, the matrix A is overwriten by 
 *          the factors L and U. The unit diagonal elements of L are not stored.
 *   
 *	@param[in]
 *  LDA     int.
 * 	    	LDA specifies the number of columns of A ( row-major order ).
 * 	    	LDA must be at least max( 1, N ).
 * 	    	
 * 	@param[out]
 * 	IPIV    int *.
 * 			ipiv is a pointer to an array of dimesion at least 
 * 			max( 1, min ( M, N ) ). 
 * 			ipiv( i ) = j, 1 <= i <= min( M, N ) implies that rows i and j have
 * 			been interchanged.    	
 *
 **/

/**
 *
 *	@retval Success successful exit
 *	@retval NoSuccess unsuccessful exit
 *
 **/

/**
 *
 *	@sa ddss_tile
 *	@sa ddss_flat2tiled
 *	@sa ddss_tiled2flat
 *
 **/ 

enum LASS_RETURN kdppgetrf( int M, int N, double *A, int LDA, int *IPIV )
{

	// Local variables
	int i;
	int mt, nt;
	int mi, mmi, ki, ni;
	int tile_size_m;
	int tile_size_mm;
	int tile_size_n;
	int tile_size_km;
	int tile_size_kn;
	int panel_size_km;
	int panel_offset;
	int offset_A, offset_AA;
	int offset_ipiv, offset_ipiv_end;
	int *local_ipiv_size;
	int *current_position;
	int previous_position;
	enum LASS_RETURN ret;
	int *ipiv_local;
	int num_panels;

	// In the panel version number of tiles in the row directions is 1
	if ( M % TILE_SIZE == 0 )
	{
		mt = M / TILE_SIZE;
	}
	else
	{
		mt = ( M / TILE_SIZE ) + 1;
	}

	if ( N % TILE_SIZE == 0 )
	{
		nt = N / TILE_SIZE;
	}
	else
	{
		nt = ( N / TILE_SIZE ) + 1;
	}

	/***************************
	--Tile A matrix allocation--
	***************************/
		
	double ( *TILE_A )[nt][TILE_SIZE * TILE_SIZE] = malloc (mt * nt * TILE_SIZE * TILE_SIZE * sizeof(double) );
	
	if ( TILE_A == NULL )
	{
		fprintf( stderr, "Failure in kdnpgetrf for matrix TILE_A\n" );
		return NoSuccess;
	}

	/* Allocation of local ipiv array */
	ipiv_local = ( int* ) malloc( MIN( M, N ) * sizeof( int ) );
	
	if ( ipiv_local == NULL )
	{
		fprintf( stderr, "Failure in kdnpgetrf for ipiv_local\n" );
		return NoSuccess;
	}

	/* Array to hold the permutations per panel */
	local_ipiv_size = ( int* ) malloc( MIN( mt, nt ) * sizeof( int ) );

	if ( local_ipiv_size == NULL )
	{
		fprintf( stderr, "Failure in kdnpgetrf for local_ipiv_size\n" );
		return NoSuccess;
	}

	/* Array to hold current position of each panel */
	current_position = ( int* ) malloc( ( MIN( mt, nt ) + 1 ) * sizeof( int ) );

	if ( current_position == NULL )
	{
		fprintf( stderr, "Failure in kdnpgetrf for current_position\n" );
		return NoSuccess;
	}

	/************************************************
	// --From flat data layout to tiled data layout--	
	************************************************/

	/* Convert from flat to tiled layout (except the first panel */
	tile_size_kn = ddss_tile_size( N, 0 );
	offset_A = tile_size_kn;
	if( N > TILE_SIZE ){
		ddss_dflat2tiled_trailing( M, N-tile_size_kn, &A[offset_A], LDA, mt, nt, 0, 1, mt, nt, TILE_A, 0 );
	}

	/* Precompute the lengths of permuatations per panel (local_ipiv_size) and current positions */
	num_panels = MIN( mt, nt );
	for ( i = 0; i < num_panels; i++ ){
		local_ipiv_size[i] = MIN(M-TILE_SIZE*i, ddss_tile_size(N, i));
		if( i == 0 )
			current_position[i] = local_ipiv_size[i];
		else
			current_position[i] = current_position[i-1] + local_ipiv_size[i];
	}	

	i = 0;
	
	/****************
	--DGETRF tile--
	****************/	

	/* For loop over panels */
	for ( ki = 0; ki < num_panels; ki++ )
	{
		// Compute the height of the panel
		tile_size_km = ddss_tile_size( M, ki );
		tile_size_kn = ddss_tile_size( N, ki );
		panel_offset = ki * TILE_SIZE;
		panel_size_km = M - panel_offset;
		offset_AA = ki * TILE_SIZE * LDA + ki * TILE_SIZE;
		offset_ipiv = current_position[ki] - MIN(panel_size_km, tile_size_kn);
		offset_ipiv_end = current_position[ki] - 1;

		if( ki != 0 ){
			/* Convert current panel from tile to flat layout (except in the first pass, i.e. when ki == 0).
 			 * The first panel is already in flat layout */
			#pragma oss task inout( TILE_A[ki:mt][ki] ) \
					priority( num_panels ) \
					firstprivate( mt, nt, ki, tile_size_kn ) \
					label ( tiled2flat_panel )
			ddss_dtiled2flat_trailing_nb( M-ki*TILE_SIZE, tile_size_kn, &A[ki*TILE_SIZE*LDA+ki*TILE_SIZE], LDA, mt, nt, ki, ki, mt, ki+1, TILE_A, num_panels);
			
			/* Bring the row of tiles of the previous panel (U factors) to panel layout. This part will not be requird in further decomposition 
 			 * Use non-blocking function. No coalision with other parts (current submatrix) */
			#pragma oss task in ( TILE_A[ki-1][ki:nt] ) \
					shared ( TILE_A ) \
					priority ( num_panels - ki ) \
					firstprivate ( N, ki, mt, nt ) \
					label ( tiled2flat_row )
			ddss_dtiled2flat_trailing_nb( TILE_SIZE, N-ki*TILE_SIZE, &A[(ki-1)*TILE_SIZE*LDA+ki*TILE_SIZE], LDA, mt, nt, ki-1, ki, ki, nt, TILE_A, num_panels-ki );
		}

		/* Compute LU of the current panel */
		#pragma oss task inout( TILE_A[ki:mt][ki] ) \
			inout( ipiv_local[offset_ipiv:offset_ipiv_end] ) \
			priority( num_panels - ki ) \
			firstprivate( i, ki, mt, panel_size_km, tile_size_kn, offset_AA, offset_ipiv ) \
			label( dgetrf )
		{
			LAPACKE_dgetrf( CblasRowMajor,
					panel_size_km, tile_size_kn,
					&A[offset_AA], LDA, 
					&ipiv_local[offset_ipiv] );

			// Update the global ipiv array
			for ( i = 0; i < local_ipiv_size[ki]; i++ )
			{
				( IPIV + offset_ipiv )[i] =
						( ipiv_local + offset_ipiv )[i] + offset_ipiv;
			}
		}

		/* Convert the current panel A[ki:ki+TILE_SIZE, ki:end] to tile storage layout
		 * In the case of the last panel, then there is no need to convert to tile layout */

		if( (ki + 1 < MIN( mt, nt )) || (M < N) ) 
			
			#pragma oss task inout ( TILE_A[ki:mt][ki] ) \
					priority ( num_panels - ki ) \
					firstprivate ( num_panels, M, ki, mt, nt, tile_size_kn ) \
					label ( flat2tiled_panel ) 
			ddss_dflat2tiled_trailing(M-ki*TILE_SIZE, tile_size_kn, &A[ki*TILE_SIZE*LDA+ki*TILE_SIZE], LDA, mt, nt, ki, ki, mt, ki+1, TILE_A, num_panels-ki);

		/* Apply permutations (swap rows) to the previously computed L factors */
		for ( ni = 0; ni < ki; ni++ )
		{		
			tile_size_n = ddss_tile_size( N, ni );
			offset_A = ki*TILE_SIZE*LDA + ni*tile_size_n;
			offset_ipiv = current_position[ki] - local_ipiv_size[ki];
			offset_ipiv_end = current_position[ki];

			// Swap the column of the tiles to the left of current panel A[ki:mt][ki]
			#pragma oss task in ( TILE_A[ki:mt][ni] ) \
					in( ipiv_local[offset_ipiv:offset_ipiv_end-1] ) \
					priority ( num_panels - ki ) \
					firstprivate( tile_size_n, offset_A, ki, ni, mt, offset_ipiv, offset_ipiv_end ) \
					label( tiled_laswp_left )
			LAPACKE_dlaswp( CblasRowMajor,
					tile_size_n, &A[offset_A], LDA,
					1, local_ipiv_size[ki],
					&ipiv_local[offset_ipiv], 1 );
		}

		/* Apply permutations (swap rows) to the trailing submatrix.
 		 * The permutations are applied to a column of tiles at the time */
		for ( ni = ki + 1; ni < nt; ni++ )
		{
			tile_size_n = ddss_tile_size( N, ni );
			offset_A = ki*TILE_SIZE*LDA + ni*TILE_SIZE;
			offset_ipiv = current_position[ki] - local_ipiv_size[ki];
			offset_ipiv_end =  current_position[ki];

			#pragma oss task inout ( TILE_A[ki:mt][ni] ) \
					in ( ipiv_local[offset_ipiv:offset_ipiv_end-1] ) \
					priority ( num_panels - ki ) \
					firstprivate ( mt, nt, tile_size_n, ki, ni, offset_ipiv, offset_ipiv_end ) \
					label ( tiled_laswp )
			ddss_dtiled_laswp( tile_size_n, ki, ni, mt, nt, TILE_A, 0, local_ipiv_size[ki], &ipiv_local[offset_ipiv], 1 );

		}

		/* Update the tile-row to the right of the current panel's diagonal tile */
		for ( ni = ki + 1; ni < nt; ni++ )
		{
			tile_size_n = ddss_tile_size( N, ni );

			/* Update the tiles to the right  of the current panel's top tile A[ki][ki] */
			#pragma oss task in( TILE_A[ki][ki] ) \
				inout( TILE_A[ki][ni] ) \
				priority ( nt - ni + 1 ) \
				firstprivate( ni, ki, tile_size_km, tile_size_n ) \
				label( dtrsm_right )
			cblas_dtrsm( CblasRowMajor, 
					( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower, 
					( CBLAS_TRANSPOSE ) NoTrans, ( CBLAS_DIAG ) Unit, 
					tile_size_km, tile_size_n, 
					1.0, TILE_A[ki][ki], TILE_SIZE, 
					TILE_A[ki][ni], TILE_SIZE );
	
			/* Update trailing submatrix by tiles */
			for ( mmi = ki + 1; mmi < mt; mmi++ )
			{
				tile_size_mm = ddss_tile_size( M, mmi );
		
				#pragma oss task in( TILE_A[mmi][ki] ) \
					in( TILE_A[ki][ni] ) \
					inout( TILE_A[mmi][ni] ) \
					priority ( nt - ni + 1 ) \
					firstprivate( ni, mmi, ki, tile_size_mm, tile_size_n ) \
					label( dgemm )
				cblas_dgemm( CblasRowMajor, 
							 ( CBLAS_TRANSPOSE ) NoTrans,
							 ( CBLAS_TRANSPOSE ) NoTrans,
							 tile_size_mm, 
							 tile_size_n, 
							 TILE_SIZE,
							 -1.0, TILE_A[mmi][ki], TILE_SIZE,
							 	   TILE_A[ki][ni],  TILE_SIZE,
							  1.0, TILE_A[mmi][ni], TILE_SIZE );			
			}
		}
		
		/* If M < N convert columns from (ki-1)*TILE : N from tiled to flat layout */
		if ( (ki+1 == MIN(mt, nt)) && (M < N) ){

			#pragma oss task in ( TILE_A[ki:mt][ki+1:nt] ) \
					firstprivate ( tile_size_km, N, ki, mt, nt ) \
					label ( tiled2flat_last )
			ddss_dtiled2flat_trailing( tile_size_km, N - (ki+1)*TILE_SIZE, &A[ki*TILE_SIZE*LDA+(ki+1)*TILE_SIZE], LDA, mt, nt, ki, ki+1, mt, nt, TILE_A);
		}

	}

	/* Wait for all tasks to finish */		
	#pragma oss taskwait

	// --Free--
	free( TILE_A );
	free( ipiv_local );
	free( local_ipiv_size );
	free( current_position );	

	return Success;

}
