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

/**
 *
 * 	@file kdposv.c
 *
 * 	@brief LASs-DDSs kdposv routine.
 *
 * 	LASs-DDSs is a software package provided by:
 * 	Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
 *
 * 	@author Pedro Valero-Lara pedro.valero@bsc.es
 * 	@date 2018-04-05
 * 	@reviewer 
 * 	@modified 
 *
 **/

/**
 *  
 *	@ingroup DDSS
 *   
 *	Solves a system of linear equations A \times X = B, where A is a M-by-M
 *	simmetric positive definite matrix and X and B are M-by-NRHS matrices.  
 *	The matrix A is descomposed as:
 *
 *		A = L \times L^T
 *		or
 *		A = U^T \times U
 *
 *	where L is a lower triangular matrix and U is an upper triangular matrix.
 *
**/

/**
 *	
 * 	@param[in]
 *	UPLO	enum DDSS_UPLO.
 *	        UPLO specifies the form of A is stored:
 *       	- Lower: Lower triangle of A is stored. The upper traingular part is
 *       	not referenced.
 *      	- Upper: Upper triangle of A is stored. The lower triangular part is
 *      	not referenced.
 *	
 *	@param[in]
 *	N       int.
 *          N specifies the order of the square matrix A. N >= 0.
 *    
 *	@param[in]
 *	NRHS    int.
 *          NRHS specifies the number of right-hand-sides (number 
 *			of columns of B). NRHS >= 0.
 *
 *	@param[in,out]
 *  A  		double *.
 *	    	A is a pointer to a positive definite matrix of dimension N by LDA. 
 *			On exit, if return value is Success, the matrix A
 * 			is overwriten by the factor U or L.
 *   
 *	@param[in]
 *  LDA     int.
 * 	    	LDA specifies the number of columns of A ( row-major order ).
 * 	    	LDA must be at least max( 1, N ).
 *
 *	@param[in,out]
 *  B  		double *.
 *	    	B is a pointer to a matrix of dimension N by NRHS, which stores the 
 *			right-hand-sides of the systems of linear equations.  
 *			(row-major order). On exit, if return value is Success, the matrix B
 * 			is overwriten by the solution matrix X.
 *   
 *	@param[in]
 *  LDB     int.
 * 	    	LDB specifies the number of columns of B ( row-major order ).
 * 	    	LDB must be at least max( 1, NRHS ).
 *
 **/

/**
 *
 *	@retval Success sucessful exit
 *	@retval NoSuccess unsucessful exit
 *
 **/

/**
 *
 *	@sa ddss_dpotrf
 *	@sa ddss_tile
 *	@sa ddss_symflat2tiled
 *	@sa ddss_symtiled2flat
 *
 **/ 

enum LASS_RETURN kdposv( enum DDSS_UPLO UPLO,
        int N, int NRHS,
        double *A, int LDA,
        double *B, int LDB )
{

	// Local variables
	int nt, nrhst;
	int mi, mmi, ki, ni, nni;
	int Am, An;
	int Bm, Bn;
	int tile_size_m;
	int tile_size_mm;
	int tile_size_n;
	int tile_size_nn;
	int tile_size_k;

	// Number of tiles
	if ( N % TILE_SIZE == 0 )
	{
		nt = N / TILE_SIZE;
	}
	else
	{
		nt = ( N / TILE_SIZE ) + 1;
	}
	if ( NRHS % TILE_SIZE == 0 )
	{
		nrhst = NRHS / TILE_SIZE;
	}
	else
	{
		nrhst = ( NRHS / TILE_SIZE ) + 1;
	}
	
	
    /****************************
	--Tile A matrix declaration--
	****************************/

	Am = nt;
	An = nt;

	Bm = nt;
	Bn = nrhst;	
	
	/***************************
	--Tile matrices allocation--
	***************************/
		
	double (*TILE_A)[An][TILE_SIZE * TILE_SIZE] = malloc ( Am * An * 
		TILE_SIZE * TILE_SIZE * sizeof( double ) );
	
	if ( TILE_A == NULL)
	{
		fprintf( stderr, "Failure in kdposv for matrix TILE_A\n" );
		return NoSuccess;
	}
	
	double ( *TILE_B )[Bn][TILE_SIZE * TILE_SIZE] = malloc(
			 Bm * Bn * TILE_SIZE * TILE_SIZE * sizeof( double ) );

	if ( TILE_B == NULL ) 
	{
		fprintf( stderr, "Failure in kdposv for matrix TILE_B\n" );
		return NoSuccess;
	}

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

	// From flat matrix A to tile matrix TILE_A
	ddss_dsymflat2tiled( N, N, A, LDA, nt, nt, TILE_A, UPLO );
	
	// From flat matrix B to tile matrix TILE_B 
	ddss_dflat2tiled( N, NRHS, B, LDB, nt, nrhst, TILE_B );

	/**************
	--DPOSV tile--
	**************/	

	if ( UPLO == Lower )
	{
		// Cholesky descomposition	
		// --UPLO = Lower--
		for ( ki = 0; ki < nt; ki++ )
		{
			tile_size_k = ddss_tile_size( N, ki );
	
			#pragma oss task inout( TILE_A[ki][ki] ) \
				firstprivate( ki ) \
				label( dpotrf )
			LAPACKE_dpotrf_work( LAPACK_ROW_MAJOR, 
							 'L',
        					 tile_size_k, 
							 TILE_A[ki][ki], tile_size_k );

			for ( mi = ki + 1; mi < nt; mi++ )
			{
				tile_size_m = ddss_tile_size( N, mi );
			
				#pragma oss task in( TILE_A[ki][ki] ) \
					inout(TILE_A[mi][ki]) \
					firstprivate( TILE_A ) \
					firstprivate( mi, ki ) \
					label( dpotrf_dtrsm )
				cblas_dtrsm( CblasRowMajor, 
					     	 ( CBLAS_SIDE ) Right, ( CBLAS_UPLO ) Lower, 
					     	 ( CBLAS_TRANSPOSE ) Trans, ( CBLAS_DIAG ) NonUnit, 
						 	 tile_size_m, TILE_SIZE, 
						 	 1.0, TILE_A[ki][ki], TILE_SIZE, 
							  	  TILE_A[mi][ki], TILE_SIZE );
			}
			for ( mmi = ki + 1; mmi < nt; mmi++ )
			{
				tile_size_mm = ddss_tile_size( N, mmi );
			
				#pragma oss task in( TILE_A[mmi][ki] ) \
					inout( TILE_A[mmi][mmi] ) \
					firstprivate( mmi, ki ) \
					label( dpotrf_dsyrk )
				cblas_dsyrk( CblasRowMajor, 
					         ( CBLAS_UPLO ) Lower, ( CBLAS_TRANSPOSE ) NoTrans, 
						     tile_size_mm, TILE_SIZE, 
						     -1.0,  TILE_A[mmi][ki], TILE_SIZE, 
						      1.0, TILE_A[mmi][mmi], tile_size_mm );
			
				for ( ni = ki + 1; ni < mmi; ni++ )
				{
				
					#pragma oss task in( TILE_A[mmi][ki] ) \
                		in( TILE_A[ni][ki] ) \
                    	inout( TILE_A[mmi][ni] ) \
 						firstprivate( ni, mmi, ki ) \
						label( dpotrf_dgemm )
					cblas_dgemm( CblasRowMajor, 
								 ( CBLAS_TRANSPOSE ) NoTrans,
								 ( CBLAS_TRANSPOSE ) Trans,
								 TILE_SIZE, 
								 TILE_SIZE, 
								 TILE_SIZE,
							 	 -1.0, TILE_A[mmi][ki], TILE_SIZE,
							 		    TILE_A[ni][ki], TILE_SIZE,
							 	  1.0, TILE_A[mmi][ni], TILE_SIZE );
				}
			}
		}
	
		/*********************************************
	 	--From tiled data layout to flat data layout--	
	 	*********************************************/
		ddss_dsymtiled2flat_nb( N, N, A, LDA, nt, nt, TILE_A, UPLO );
		
		// Triangular solve
		// --SIDE = Left & UPLO = Lower & TRANS_A = NoTrans--  
		for ( ki = 0; ki < nt; ki++ )
		{
			tile_size_k = ddss_tile_size( N, ki );

			for ( ni = 0; ni < nrhst; ni++ )
			{
				tile_size_n = ddss_tile_size( NRHS, ni );
						
				#pragma oss task in( TILE_A[ki][ki] ) \
					inout(  TILE_B[ki][ni] ) \
					shared( TILE_A, TILE_B ) \
					firstprivate( ki, ni ) \
					label( dtrsm_dtrsm1 )
				cblas_dtrsm( CblasRowMajor,
							 ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower,
							 ( CBLAS_TRANSPOSE ) NoTrans,
							 ( CBLAS_DIAG ) NonUnit,
							 tile_size_k,
							 tile_size_n,
							 1.0, TILE_A[ki][ki], tile_size_k,
								  TILE_B[ki][ni], tile_size_n );
			}
			for ( nni = 0; nni < nrhst; nni++ ) 
			{
				tile_size_nn = ddss_tile_size( NRHS, nni );
				
				for ( mi = ki + 1; mi < nt; mi++ )
				{
					
					#pragma oss task in( TILE_A[mi][ki] ) \
                    	in( TILE_B[ki][nni] ) \
                        inout( TILE_B[mi][nni] ) \
                        shared( TILE_A, TILE_B ) \
                        firstprivate( ki, mi, nni ) \
                        label( dtrsm_dgemm1 )
					cblas_dgemm( CblasRowMajor,
								 CblasNoTrans, CblasNoTrans,
								 tile_size_k,
								 tile_size_nn,
								 tile_size_k,
								 -1.0,  TILE_A[mi][ki],  tile_size_k,
									   TILE_B[ki][nni], tile_size_nn,
								  1.0, TILE_B[mi][nni], tile_size_nn );
				}
			}
		}

		// Triangular solve
		// --SIDE = Left & UPLO = Lower & TRANS_A = Trans-- 
		for ( ki = 0; ki < nt; ki++ )
        {
            tile_size_k = ddss_tile_size( N, ki );

            for ( ni = 0; ni < nrhst; ni++ )
            {
                tile_size_n = ddss_tile_size( NRHS, ni );

                #pragma oss task in( TILE_A[ki][ki] ) \
                    inout( TILE_B[ki][ni] ) \
                    shared( TILE_A, TILE_B ) \
                    firstprivate( ki, ni ) \
                    label( dtrsm_dtrsm2 )
                cblas_dtrsm( CblasRowMajor,
                             ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower,
                             ( CBLAS_TRANSPOSE ) Trans,
                             ( CBLAS_DIAG ) NonUnit,
                             tile_size_k,
                             tile_size_n,
                             1.0, 	TILE_A[ki][ki], tile_size_k,
                            		TILE_B[ki][ni], tile_size_n );
            }

            for ( nni = 0; nni < nrhst; nni++ )
            {
                tile_size_nn = ddss_tile_size( NRHS, nni );
            	
				for ( mi = ki - 1; mi >= 0; mi-- )
            	{
            		tile_size_m = ddss_tile_size( N, mi );

                    #pragma oss task in( TILE_A[ki][mi] ) \
                    	in( TILE_B[ki][nni] ) \
                        inout(  TILE_B[mi][nni] ) \
                        shared( TILE_A, TILE_B ) \
                        firstprivate( ki, mi, nni ) \
                        label( dtrsm_dgemm2 )
                    cblas_dgemm( CblasRowMajor,
                                 CblasTrans, CblasNoTrans,
                                 tile_size_m,
                                 tile_size_nn,
                                 tile_size_k,
                                 -1.0,  TILE_A[ki][mi],  tile_size_m,
                                       TILE_B[ki][nni], tile_size_nn,
                                  1.0, TILE_B[mi][nni], tile_size_nn );
                }
            }
        }
	}
	else
	{
		// Cholesky descomposition	
		// --UPLO = Upper--	
		for ( ki = 0; ki < nt; ki++ )
		{
			tile_size_k = ddss_tile_size( N, ki );
			
			#pragma oss task inout( TILE_A[ki][ki] ) \
				shared( TILE_A ) \
				firstprivate( ki ) \
				label( dpotrf )
			LAPACKE_dpotrf_work( LAPACK_ROW_MAJOR, 
							 'U',
        					 tile_size_k, TILE_A[ki][ki], 
							 tile_size_k );

			for ( mi = ki + 1; mi < nt; mi++ )
			{
				tile_size_m = ddss_tile_size( N, mi );
			
				#pragma oss task in( TILE_A[ki][ki] ) \
					inout(TILE_A[ki][mi]) \
					shared( TILE_A ) \
					firstprivate( mi, ki ) \
					label( dpotrf_dtrsm )
				cblas_dtrsm( CblasRowMajor, 
					     	 ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Upper, 
					     	 ( CBLAS_TRANSPOSE ) Trans, ( CBLAS_DIAG ) NonUnit, 
						 	 TILE_SIZE, tile_size_m, 
						 	 1.0, TILE_A[ki][ki], TILE_SIZE, 
							  	  TILE_A[ki][mi], tile_size_m );
			}
			for ( mmi = ki + 1; mmi < nt; mmi++ )
			{
				tile_size_mm = ddss_tile_size( N, mmi );
			
				#pragma oss task in( TILE_A[ki][mmi] ) \
					inout( TILE_A[mmi][mmi] ) \
					shared( TILE_A ) \
					firstprivate( mmi, ki ) \
					label( dpotrf_dsyrk )
				cblas_dsyrk( CblasRowMajor, 
					         ( CBLAS_UPLO ) Upper, ( CBLAS_TRANSPOSE ) Trans, 
						     tile_size_mm, TILE_SIZE, 
						     -1.0,  TILE_A[ki][mmi], tile_size_mm, 
						      1.0, TILE_A[mmi][mmi], tile_size_mm );
			
				for ( ni = ki + 1; ni < mmi; ni++ )
				{
					#pragma oss task in( TILE_A[ki][ni] ) \
                		in( TILE_A[ki][mmi] ) \
                    	inout( TILE_A[ni][mmi] ) \
						shared( TILE_A ) \
 						firstprivate( ni, mmi, ki ) \
						label( dpotrf_dgemm )
					    cblas_dgemm( CblasRowMajor, 
								     ( CBLAS_TRANSPOSE ) Trans,
								     ( CBLAS_TRANSPOSE ) NoTrans,
								     TILE_SIZE, 
								     TILE_SIZE, 
								     TILE_SIZE,
							 	     -1.0, TILE_A[ki][ni],  TILE_SIZE,
							 		       TILE_A[ki][mmi], TILE_SIZE,
							 	      1.0, TILE_A[ni][mmi], TILE_SIZE );
				}
			}
		}
		/*********************************************
	 	--From tiled data layout to flat data layout--	
	 	 *********************************************/
		ddss_dsymtiled2flat_nb( N, N, A, LDA, nt, nt, TILE_A, UPLO );
		// Triangular solve
		// --SIDE = Left & UPLO = Upper & TRANS_A = Trans-- 
		for ( ki = 0; ki < nt; ki++ )
        {
        	tile_size_k = ddss_tile_size( N, ki );
                                
			for ( ni = 0; ni < nrhst; ni++ )
            {
            	tile_size_n = ddss_tile_size( NRHS, ni );

                #pragma oss task in( TILE_A[ki][ki] ) \
                	inout( TILE_B[ki][ni] ) \
                    shared( TILE_A, TILE_B ) \
                    firstprivate( ki, ni ) \
                    label( dtrsm_dtrsm1 )
                cblas_dtrsm( CblasRowMajor,
                             ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower,
                             ( CBLAS_TRANSPOSE ) Trans,
                             ( CBLAS_DIAG ) NonUnit,
                             tile_size_k,
                             tile_size_n,
                             1.0, TILE_A[ki][ki], tile_size_k,
                                  TILE_B[ki][ni], tile_size_n );
            }

            for ( mi = ki + 1; mi < nt; mi++ )
            {
            	tile_size_m = ddss_tile_size( N, mi );
                        
				for ( nni = 0; nni < nrhst; nni++ )
                {
                    tile_size_nn = ddss_tile_size( NRHS, nni );

                    #pragma oss task in( TILE_A[ki][mi]) \
                    	in( TILE_B[ki][nni] ) \
                        inout(  TILE_B[mi][nni] ) \
                        shared( TILE_A, TILE_B ) \
                        firstprivate( ki, mi, nni ) \
                        label( dtrsm_dgemm1 )
                    cblas_dgemm( CblasRowMajor,
                                 CblasTrans, CblasNoTrans,
                                 tile_size_m,
                                 tile_size_nn,
                                 tile_size_k,
                                 -1.0,  TILE_A[ki][mi],  tile_size_m,
                                       TILE_B[ki][nni], tile_size_nn,
                                  1.0, TILE_B[mi][nni], tile_size_nn );
                }
            }
        }
		// Triangular solve
		// --SIDE = Left & UPLO = Upper & TRANS_A = NoTrans-- 
		for ( ki = 0; ki < nt; ki++ )  
        { 
            tile_size_k = ddss_tile_size( N, ki ); 
 
        	for ( ni = 0; ni < nrhst; ni++ )  
            { 
                tile_size_n = ddss_tile_size( NRHS, ni ); 
 
                #pragma oss task in( TILE_A[ki][ki] ) \
					inout( TILE_B[ki][ni] ) \
					shared( TILE_A, TILE_B ) \
					firstprivate( ki, ni ) \
					label( dtrsm_dtrsm2 ) 
                cblas_dtrsm( CblasRowMajor, 
                             ( CBLAS_SIDE ) Left, ( CBLAS_UPLO ) Lower, 
                             ( CBLAS_TRANSPOSE ) NoTrans, 
                             ( CBLAS_DIAG ) NonUnit, 
                             tile_size_k,  
                             tile_size_n, 
                             1.0, TILE_A[ki][ki], tile_size_k, 
                                  TILE_B[ki][ni], tile_size_n ); 
            } 
 
			for ( nni = 0; nni < nrhst; nni++ )  
            { 
                tile_size_nn = ddss_tile_size( NRHS, nni ); 
            
				for ( mi = ki - 1; mi >= 0; mi-- )  
            	{ 
                	tile_size_m = ddss_tile_size( N, mi ); 
 
                    #pragma oss task in( TILE_A[mi][ki] ) \
						in( TILE_B[ki][nni] ) \
						inout( TILE_B[mi][nni] ) \
						shared( TILE_A, TILE_B ) \
						firstprivate( ki, mi, nni ) \
						label( dtrsm_dgemm2 ) 
                    cblas_dgemm( CblasRowMajor, 
                            	 CblasNoTrans, CblasNoTrans, 
                                 tile_size_m,  
                                 tile_size_nn,  
                                 tile_size_k, 
                                 -1.0,  TILE_A[mi][ki],  tile_size_k, 
                                 	   TILE_B[ki][nni], tile_size_nn, 
                                  1.0, TILE_B[mi][nni], tile_size_nn ); 
                } 
            } 
        }
	}

	/*********************************************
	--From tiled data layout to flat data layout--	
	*********************************************/
	ddss_dtiled2flat( N, NRHS, B, LDB, nt, nrhst, TILE_B );
	
	// --Tile A and B matrices free--
	free( TILE_A );
	free( TILE_B );

	return Success;

}
