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

/**
 *
 * 	@file kdsyrk.c
 *
 * 	@brief LASs-DDSs kdsyrk 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 Boro Sofranac       boro.sofranac@bsc.es
 * 	@date 2018-02-13
 * 	@reviewer 
 * 	@modified 
 *
 **/

/**
 *  
 *	@ingroup DDSS
 *   
 *	Performs one of the symmetric rank k operations:
 *
 *		C = ALPHA * A * op( A ) + BETA * C	or
 *		C = ALPHA * op( A ) * A + BETA * C
 *	
 *	where op( X ) is:
 *
 *		op( X ) = X**T
 *
 *	ALPHA and BETA are scalars, C is an N by N symmetric matrix and A is an
 *	N by K matrix in the first case and a K by N matrix in the second case.
 *
 **/

/**
 *	@param[in]
 *	UPLO	enum DDSS_UPLO.
 *	    	UPLO specifies the form in which C is stored:
 *       	- Lower: Lower triangular part of C is stored. The upper traingular
 *       	part is not referenced.
 *      	- Upper: Upper triangular part of C is stored. The lower triangular
 *      	part is not referenced.
 *  
 *  @param[in]
 *	TRANS_A	enum DDSS_TRANS.
 *	        TRANS_A specifies the operation to be performed as follows:  
 *      	- NoTrans:   C = ALPHA * A * A**T + BETA * C
 *       	- Trans:     C = ALPHA * A**T * A + BETA * C
 *       	
 *	@param[in]
 *	N       int.
 *			N specifies the order of matrix C. N must be at least zero.
 *
 *	@param[in]
 *	K       int.
 *          With TRANS_A = NoTrans, K specifies the number of columns of the
 *          matrix A, and with TRANS_A = Trans, K specifies the number of rows 
 *			of the matrix A. K must be at least zero.
 *    
 *	@param[in]
 *  ALPHA   double.
 *   		ALPHA specifies the scalar alpha.
 *   		
 *	@param[in]
 *  A  		double *.
 *          A is a pointer to a matrix of dimension Na ( rows ) by Ka  
 *			( columns ), where Na is N and Ka is K when TRANS_A = NoTrans, 
 *			and Na is K and Ka is N otherwise.
 *
 *  @param[in]
 *  LDA     int.
 *          LDA specifies the number of columns of A ( row-major order ).
 *          When TRANS_A = NoTrans then LDA must be at least  max( 1, K ),
 *          otherwise LDA must be at least max( 1, N ).
 *          
 *	@param[in]
 *  BETA    double.
 *   		BETA specifies the scalar beta.
 *   
 *	@param[in, out]
 *  C  		double *.
 * 		    C is a pointer to a matrix of dimension N by N. 
 *			When UPLO = Uppper the strictly lower triangular part of C is not
 *			referenced. On exit, the upper triangular part of C is overwritten
 *			by the upper triangular part of the updated solution matrix C.
 *			When UPLO = Lower the strictly upper triangular part of C is not
 *			referenced. On exit, the lower triangular part of C is overwritten
 *			by the lower triangular part of the updated solution matrix C.
 *   
 *	@param[in]
 *  LDC     int.
 * 		    LDC specifies the number of columns of C ( row-major order ).
 *          LDC must be at least max( 1, N ).
 *
 **/


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

/**
 *
 *	@sa ddss_tile
 *	@sa ddss_flat2tiled
 *	@sa ddss_dsymflat2tiled
 *	@sa ddss_dsymtiled2flat
 *
 **/

enum LASS_RETURN kdsyrk( enum DDSS_UPLO UPLO, enum DDSS_TRANS TRANS_A,
		int N, int K,
		const double ALPHA, double *A, int LDA,
		const double BETA,  double *C, int LDC ) 
{
	
	// Local variables
	int nt, kt;
	int mi, ni, ki;
	int Am, An;
	int Cm, Cn;
	int tile_size_m;
	int tile_size_n;
	int tile_size_k;
	int ka;
	double beta;

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

	if ( K % TILE_SIZE == 0 )
	{
		kt = K / TILE_SIZE;
	}
	else
	{
		kt = ( K / TILE_SIZE ) + 1;
	}

	/****************************
	 --Tile matrices declaration--
	 ****************************/

	if ( TRANS_A == NoTrans )
	{
		Am = nt;
		An = kt;
		ka = N;
	}
	else
	{
		Am = kt;
		An = nt;
		ka = K;
	}
	
	Cm = Cn = nt;
	
	/***************************
	 --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 ddss_dtile_alloc for matrix TILE_A\n" );
		return NoSuccess;
	}

	double ( *TILE_C )[Cn][TILE_SIZE * TILE_SIZE] = malloc(
			Cm * Cn * TILE_SIZE * TILE_SIZE * sizeof( double ) );

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

	/*********************************************
	 --From flat data layout to tiled data layout--	
	 *********************************************/
	
	// From flat matrix A to tile matrix TILE_A
	ddss_dflat2tiled( ka, LDA, A, LDA, Am, An, TILE_A );
	
	// From flat matrix C to tile matrix TILE_C
	ddss_dsymflat2tiled( N, N, C, LDC, Cm, Cn, TILE_C, UPLO );

	/*************
	 --DSYRK tile--
	 *************/
	
	// --TRANS_A = NoTrans & UPLO = Upper-- 
	if ( TRANS_A == NoTrans )
	{
		if ( UPLO == Upper )
		{
			for ( mi = 0; mi < nt; mi++ )
			{
				tile_size_m = ddss_tile_size( N, mi );
				for ( ni = 0; ni < kt; ni++ )
				{
					tile_size_n = ddss_tile_size( K, ni );

					if ( ni == 0 )
					{
						beta = BETA;
					}
					else
					{
						beta = 1.0;
					}

					#pragma oss task in( TILE_A[mi][ni] ) \
						inout( TILE_C[mi][mi] ) \
						shared( TILE_A, TILE_C ) \
						firstprivate( mi, ni ) \
						label( dsyrk )
					cblas_dsyrk( CblasRowMajor,
							     ( CBLAS_UPLO ) UPLO, 
								 ( CBLAS_TRANSPOSE ) TRANS_A,
								 tile_size_m,
								 tile_size_n,
								 ALPHA, TILE_A[mi][ni], tile_size_n,
							     beta,  TILE_C[mi][mi], tile_size_m );

					for ( ki = mi + 1; ki < nt; ki++ )
					{
						tile_size_k = ddss_tile_size( N, ki );

						#pragma oss task in( TILE_A[mi][ni] ) \
							in( TILE_A[ki][ni] ) \
							inout( TILE_C[mi][ki] ) \
							shared( TILE_A, TILE_C ) \
							firstprivate( mi, ni, ki ) \
							label( dgemm )
						cblas_dgemm( CblasRowMajor,
									 CblasNoTrans, CblasTrans,
								     tile_size_m,
									 tile_size_k,
									 tile_size_n,
									 ALPHA, TILE_A[mi][ni], tile_size_n,
									 	    TILE_A[ki][ni], tile_size_n,
									 beta,  TILE_C[mi][ki], tile_size_k );
					}

				}
			}
		}
		// --TRANS_A = NoTrans & UPLO = Lower-- 
		else
		{
			for ( mi = 0; mi < nt; mi++ ) 
			{
				tile_size_m = ddss_tile_size( N, mi );
				for ( ni = 0; ni < kt; ni++ )
				{
					tile_size_n = ddss_tile_size( K, ni );

					if ( ni == 0 )
					{
						beta = BETA;
					} 
					else 
					{
						beta = 1.0;
					}

					#pragma oss task in( TILE_A[mi][ni] ) \
						inout( TILE_C[mi][mi] ) \
						shared( TILE_A, TILE_C ) \
						firstprivate( mi, ni ) \
						label( dsyrk )
					cblas_dsyrk( CblasRowMajor,
							     ( CBLAS_UPLO ) UPLO, 
								 ( CBLAS_TRANSPOSE ) TRANS_A,
								 tile_size_m,
								 tile_size_n,
								 ALPHA, TILE_A[mi][ni], tile_size_n,
								 beta,  TILE_C[mi][mi], tile_size_m );

					for ( ki = mi + 1; ki < nt; ki++ )
					{
						tile_size_k = ddss_tile_size( N, ki );

						#pragma oss task in( TILE_A[ki][ni] ) \
							in( TILE_A[mi][ni] ) \
							inout( TILE_C[ki][mi] ) \
							shared( TILE_A, TILE_C ) \
							firstprivate( mi, ni, ki ) \
							label( dgemm )
						cblas_dgemm( CblasRowMajor,
									 CblasNoTrans, CblasTrans,
									 tile_size_k,
									 tile_size_m,
									 tile_size_n,
									 ALPHA, TILE_A[ki][ni], tile_size_n,
									 	 	TILE_A[mi][ni], tile_size_n,
									 beta,  TILE_C[ki][mi], tile_size_m );
					}

				}
			}
		}
	} 
	// --TRANS_A = Trans & UPLO = Upper-- 
	else
	{
		if ( UPLO == Upper )
		{
			for ( mi = 0; mi < kt; mi++ )
			{
				tile_size_m = ddss_tile_size( K, mi );
				for ( ni = 0; ni < nt; ni++ )
				{
					tile_size_n = ddss_tile_size( N, ni );

					if ( mi == 0 )
					{
						beta = BETA;
					}
					else
					{
						beta = 1.0;
					}

					#pragma oss task in( TILE_A[mi][ni] ) \
						inout( TILE_C[ni][ni] ) \
						shared( TILE_A, TILE_C ) \
						firstprivate( mi, ni ) \
						label( dsyrk )
					cblas_dsyrk( CblasRowMajor,
								 ( CBLAS_UPLO ) UPLO, 
								 ( CBLAS_TRANSPOSE ) TRANS_A,
								 tile_size_n,
								 tile_size_m,
								 ALPHA, TILE_A[mi][ni], tile_size_n,
								 beta,  TILE_C[ni][ni], tile_size_n );

					for ( ki = ni + 1; ki < nt; ki++ ) 
					{
						tile_size_k = ddss_tile_size( N, ki );

						#pragma oss task in( TILE_A[mi][ni] ) \
							in( TILE_A[mi][ki] ) \
							inout( TILE_C[ni][ki] ) \
							shared( TILE_A, TILE_C ) \
							firstprivate( mi, ni, ki ) \
							label( dgemm )
						cblas_dgemm( CblasRowMajor,
									 CblasTrans, CblasNoTrans,
									 tile_size_n,
									 tile_size_k,
									 tile_size_m,
									 ALPHA, TILE_A[mi][ni], tile_size_n,
									 	 	TILE_A[mi][ki], tile_size_k,
									 beta,  TILE_C[ni][ki], tile_size_k );
					}
				}
			}
		} 
		// --TRANS_A = Trans & UPLO = Lower-- 
		else 
		{
			for ( mi = 0; mi < kt; mi++ )
			{
				tile_size_m = ddss_tile_size( K, mi );
				for ( ni = 0; ni < nt; ni++ ) 
				{
					tile_size_n = ddss_tile_size( N, ni );

					if ( mi == 0 )
					{
						beta = BETA;
					} 
					else
					{
						beta = 1.0;
					}

					#pragma oss task in( TILE_A[mi][ni] ) \
						inout( TILE_C[ni][ni] ) \
						shared( TILE_A, TILE_C ) \
						firstprivate( mi, ni ) \
						label( dsyrk )
					cblas_dsyrk( CblasRowMajor,
								 ( CBLAS_UPLO ) UPLO, 
								 ( CBLAS_TRANSPOSE ) TRANS_A,
								 tile_size_n,
								 tile_size_m,
								 ALPHA, TILE_A[mi][ni], tile_size_n,
								 beta,  TILE_C[ni][ni], tile_size_n );

					for ( ki = ni + 1; ki < nt; ki++ )
					{
						tile_size_k = ddss_tile_size( N, ki );

						#pragma oss task in( TILE_A[mi][ki] ) \
							in( TILE_A[mi][ni] ) \
							inout( TILE_C[ki][ni] ) \
							shared( TILE_A, TILE_C ) \
							firstprivate( mi, ni, ki ) \
							label( dgemm )
						cblas_dgemm( CblasRowMajor,
									 CblasTrans, CblasNoTrans,
									 tile_size_k,
									 tile_size_n,
									 tile_size_m,
									 ALPHA, TILE_A[mi][ki], tile_size_k,
									 	    TILE_A[mi][ni], tile_size_n,
									 beta,  TILE_C[ki][ni], tile_size_n );
					}
				}
			}
		}
	}

	/*********************************************
	 --From tiled data layout to flat data layout--	
	 *********************************************/
	
	ddss_dsymtiled2flat( N, N, C, LDC, Cm, Cn, TILE_C, UPLO );

	// --Tile matrices free--
	free( TILE_A );
	free( TILE_C );

	return Success;

}
