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

/**
 *
 * 	@file kdtrmm.c
 *
 * 	@brief LASs-DDSs kdtrmm 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-03-19
 * 	@reviewer 
 * 	@modified 
 *
 **/

/**
 *  
 *	@ingroup DDSS
 *   
 *	Performs one of the matrix-matrix operations:
 *
 *		B = ALPHA * op( A ) * B, or B = ALPHA * B * op( A )
 *
 *	where op( A ) is one of:
 *
 *		op( A ) = A      or
 *		op( A ) = A**T 
 *	
 *	ALPHA is a scalar, B is a M by N matrix and A is a unit, or non-unit,
 *	upper or lower triangular matrix.
 *
 **/

/**
 *	
 * 	@param[in]
 *	SIDE	enum DDSS_SIDE.
 *	        SIDE specifies the position of the triangular A matrix in 
 *			the operations:
 *      	- Left:      B = ALPHA * op( A ) * B.
 *       	- Right:     B = ALPHA * B * op( A ).
 *	
 *	@param[in]
 *	UPLO	enum DDSS_UPLO.
 *	    	UPLO specifies the form in which 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]
 *	TRANS_A	enum DDSS_TRANS.
 *	        TRANS_A specifies the form of op( A ) to be used:  
 *      	- NoTrans:   op( A ) = A.
 *       	- Trans:     op( A ) = A**T.
 *       	
 *  @param[in]
 *	DIAG	enum DDSS_DIAG.
 *	        DIAG specifies whether or not A is unit triangular as follows:
 *      	- Unit:      A is assumed to be unit triangular.
 *       	- NonUnit:   A is not assumed to be unit triangular.
 *       	
 *	@param[in]
 *	M       int.
 *			M specifies the number of rows of B. M must be at least zero.
 *
 *	@param[in]
 *	N       int.
 *          N specifies the number of columns of B.  N 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 K by K, where K is M when 
 *          SIDE = Left and is N otherwise.
 *			When UPLO = Uppper the strictly lower triangular part of A is not
 *			referenced and when UPLO = Lower the strictly upper triangular part
 *			of A is not referenced.
 *          Note that when DIAG = Unit, the diagonal elements of A are not 
 *          referenced either, but are assumed to be unity.
 *
 *  @param[in]
 *  LDA     int.
 *          LDA specifies the number of columns of A ( row-major order ).
 *          When SIDE = Left then LDA must be at least  max( 1, M ), otherwise
 *          LDA must be at least max( 1, N ).
 *   
 *	@param[in, out]
 *  B  		double *.
 * 		    B is a pointer to a matrix of dimension M by N. 
 *			On exit the matrix B is overwritten by the transformed matrix.
 *   
 *	@param[in]
 *  LDB     int.
 * 		    LDB specifies the number of columns of B ( row-major order ).
 *          LDB 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_tiled2flat
 *
 **/

enum LASS_RETURN kdtrmm( enum DDSS_SIDE SIDE, enum DDSS_UPLO UPLO,
		enum DDSS_TRANS TRANS_A, enum DDSS_DIAG DIAG,
		int M, int N,
		const double ALPHA, double *A, int LDA,
						    double *B, int LDB )
{

	// Local variables
	int k;
	int mt, nt;
	int ki, ni, mi, nni, mmi;
	int Am, An;
	int Bm, Bn;
	int tile_size_m;
	int tile_size_n;
	int tile_size_k;
	
	// Number of tiles
	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 matrices declaration--
	 ****************************/

	if ( SIDE == Left ) 
	{
		Am = An = mt;
		k = M;
	} 
	else 
	{
		Am = An = nt;
		k = N;
	}

	Bm = mt;
	Bn = 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_B )[Bn][TILE_SIZE * TILE_SIZE] = malloc(
			 Bm * Bn * TILE_SIZE * TILE_SIZE * sizeof( double ) );

	if ( TILE_B == NULL ) 
	{
		fprintf( stderr, "Failure in ddss_dtile_alloc 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( k, k, A, LDA, Am, An, TILE_A, UPLO );

	// From flat matrix B to tile matrix TILE_B 
	ddss_dflat2tiled( M, N, B, LDB, Bm, Bn, TILE_B );
	
	/*************
	 --DTRMM tile--
	 *************/

	if ( SIDE == Left ) 
	{
		if ( UPLO == Upper ) 
		{
			// --SIDE = Left & UPLO = Upper & TRANS_A = NoTrans-- 
			if ( TRANS_A == NoTrans ) 
			{
				for ( ki = 0; ki < mt; ki++ ) 
				{
					tile_size_k = ddss_tile_size( M, ki );
					
					for ( ni = 0; ni < nt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( N, ni );
						
						#pragma oss task in( TILE_A[ki][ki] ) \
							    inout( TILE_B[ki][ni] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_k, 
								     tile_size_n,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ki][ni], tile_size_n );
						
						for ( mi = ki + 1; mi < mt; mi++ ) 
						{
							tile_size_m = ddss_tile_size( M, mi );
						
							#pragma oss task in( TILE_A[ki][mi] ) \
									in( TILE_B[mi][ni] ) \
									inout( TILE_B[ki][ni] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasNoTrans,
										 tile_size_k, 
										 tile_size_n, 
										 tile_size_m,
										 ALPHA, TILE_A[ki][mi],  tile_size_m,
												TILE_B[mi][ni],  tile_size_n,
										 1.0,   TILE_B[ki][ni],  tile_size_n );
						}
					}
				}
			}
			// --SIDE = Left & UPLO = Upper & TRANS_A = Trans-- 
			else 
			{	
				for ( ki = mt-1; ki > -1; ki-- ) 
				{
					tile_size_k = ddss_tile_size( M, ki );
					
					for ( ni = 0; ni < nt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( N, ni );
					
						#pragma oss task in( TILE_A[ki][ki] ) \
								inout( TILE_B[ki][ni] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_k, 
								     tile_size_n,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ki][ni], tile_size_n );
						
						for ( mi = 0; mi < ki; mi++ ) 
						{
							tile_size_m = ddss_tile_size( M, mi );
							
							#pragma oss task in( TILE_A[mi][ki] ) \
									in( TILE_B[mi][ni] ) \
									inout( TILE_B[ki][ni] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									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_B[mi][ni],  tile_size_n,
										 1.0,   TILE_B[ki][ni],  tile_size_n );
						}
					}
				}
				
			}
		}
		// --SIDE = Left & UPLO = Lower & TRANS_A = NoTrans--  
		else 
		{
			if ( TRANS_A == NoTrans ) 
			{
				for ( ki = mt-1; ki > -1; ki-- ) 
				{
					tile_size_k = ddss_tile_size( M, ki );
					
					for ( ni = 0; ni < nt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( N, ni );
						
						#pragma oss task in( TILE_A[ki][ki] ) \
								inout( TILE_B[ki][ni] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_k, 
								     tile_size_n,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ki][ni], tile_size_n );
						
						for ( mi = 0; mi < ki; mi++ ) 
						{	
							tile_size_m = ddss_tile_size( M, mi );
							
							#pragma oss task in( TILE_A[ki][mi] ) \
									in( TILE_B[mi][ni] ) \
									inout( TILE_B[ki][ni] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasNoTrans,
										 tile_size_k, 
										 tile_size_n, 
										 tile_size_m,
										 ALPHA, TILE_A[ki][mi],  tile_size_m,
												TILE_B[mi][ni],  tile_size_n,
										 1.0,   TILE_B[ki][ni],  tile_size_n );
						}
					}
				}
				
			}
			// --SIDE = Left & UPLO = Lower & TRANS_A = Trans--    
			else 
			{		
				for ( ki = 0; ki < mt; ki++ ) 
				{
					tile_size_k = ddss_tile_size( M, ki );
					
					for ( ni = 0; ni < nt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( N, ni );
											
						#pragma oss task in( TILE_A[ki][ki] ) \
				    			inout( TILE_B[ki][ni] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_k, 
								     tile_size_n,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ki][ni], tile_size_n );
						
						for ( mi = ki + 1; mi < mt; mi++ ) 
						{	
							tile_size_m = ddss_tile_size( M, mi );
							
							#pragma oss task in( TILE_A[mi][ki] ) \
									in( TILE_B[mi][ni] ) \
									inout( TILE_B[ki][ni] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									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_B[mi][ni],  tile_size_n,
										 1.0,   TILE_B[ki][ni],  tile_size_n );
						}
					}
				}
			}
		}
	}
	// --SIDE = Right & UPLO = Upper & TRANS_A = NoTrans--   
	else
	{
		if ( UPLO == Upper )
		{
			if ( TRANS_A == NoTrans )
			{	
				for ( ki = nt-1; ki > -1; ki-- ) 
				{
					tile_size_k = ddss_tile_size( N, ki );
					
					for ( ni = 0; ni < mt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( M, ni );
							
						#pragma oss task in( TILE_A[ki][ki] ) \
				    			inout( TILE_B[ni][ki] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_n, 
								     tile_size_k,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ni][ki], tile_size_k );
						
						for ( mi = 0; mi < ki; mi++ ) 
						{	
							tile_size_m = ddss_tile_size( N, mi );
							
							#pragma oss task in( TILE_B[ni][mi] ) \
									in( TILE_A[mi][ki] ) \
									inout( TILE_B[ni][ki] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasNoTrans,
										 tile_size_n, 
										 tile_size_k, 
										 tile_size_m,
										 ALPHA, TILE_B[ni][mi],  tile_size_m,
												TILE_A[mi][ki],  tile_size_k,
										 1.0,   TILE_B[ni][ki],  tile_size_k );
						}
					}
				}
			}
			// --SIDE = Right & UPLO = Upper & TRANS_A = Trans-- 
			else
			{
				for ( ki = 0; ki < nt; ki++ ) 
				{
					tile_size_k = ddss_tile_size( N, ki );
					
					for ( ni = 0; ni < mt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( M, ni );
						
						#pragma oss task in( TILE_A[ki][ki] ) \
				    			inout( TILE_B[ni][ki] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_n, 
								     tile_size_k,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ni][ki], tile_size_k );
						
						for ( mi = ki+1; mi < nt; mi++ ) 
						{	
							tile_size_m = ddss_tile_size( N, mi );
							
							#pragma oss task in( TILE_B[ni][mi] ) \
									in( TILE_A[ki][mi] ) \
									inout( TILE_B[ni][ki] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasTrans,
										 tile_size_n, 
										 tile_size_k, 
										 tile_size_m,
										 ALPHA, TILE_B[ni][mi],  tile_size_m,
												TILE_A[ki][mi],  tile_size_m,
										 1.0,   TILE_B[ni][ki],  tile_size_k );	
						}
					}
				}
			}
		}
		// --SIDE = Right & UPLO = Lower & TRANS_A = NoTrans--  
		else
		{
			if ( TRANS_A == NoTrans )
			{
				for ( ki = 0; ki < nt; ki++ ) 
				{
					tile_size_k = ddss_tile_size( N, ki );
					
					for ( ni = 0; ni < mt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( M, ni );
						
						#pragma oss task in( TILE_A[ki][ki] ) \
				    			inout( TILE_B[ni][ki] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_n, 
								     tile_size_k,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ni][ki], tile_size_k );
						
						for ( mi = ki+1; mi < nt; mi++ ) 
						{
							tile_size_m = ddss_tile_size( N, mi );
							
							#pragma oss task in( TILE_B[ni][mi] ) \
									in( TILE_A[mi][ki] ) \
									inout( TILE_B[ni][ki] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasNoTrans,
										 tile_size_n, 
										 tile_size_k, 
										 tile_size_m,
										 ALPHA, TILE_B[ni][mi],  tile_size_m,
												TILE_A[mi][ki],  tile_size_k,
										 1.0,   TILE_B[ni][ki],  tile_size_k );
						}
					}
				}
			}
			// --SIDE = Right & UPLO = Lower & TRANS_A = Trans--  
			else 
			{
				for ( ki = nt-1; ki > -1; ki-- ) 
				{
					tile_size_k = ddss_tile_size( N, ki );
					
					for ( ni = 0; ni < mt; ni++ ) 
					{
						tile_size_n = ddss_tile_size( M, ni );
						
						#pragma oss task in( TILE_A[ki][ki] ) \
				    			inout( TILE_B[ni][ki] ) \
								shared( TILE_A, TILE_B ) \
								firstprivate( ki, ni ) \
								label( dtrmm )
						cblas_dtrmm( CblasRowMajor,
								     ( CBLAS_SIDE ) SIDE, ( CBLAS_UPLO ) UPLO,
								     ( CBLAS_TRANSPOSE ) TRANS_A,
								     ( CBLAS_DIAG ) DIAG,
								     tile_size_n, 
								     tile_size_k,
								     ALPHA, TILE_A[ki][ki], tile_size_k,
								            TILE_B[ni][ki], tile_size_k );
						
						for ( mi = 0; mi < ki; mi++ ) 
						{
							tile_size_m = ddss_tile_size( N, mi );
							
							#pragma oss task in( TILE_B[ni][mi] ) \
									in( TILE_A[ki][mi] ) \
									inout( TILE_B[ni][ki] ) \
									shared( TILE_A, TILE_B ) \
									firstprivate( ki, mi, ni ) \
									label( dgemm )
							cblas_dgemm( CblasRowMajor,
										 CblasNoTrans, CblasTrans,
										 tile_size_n, 
										 tile_size_k, 
										 tile_size_m,
										 ALPHA, TILE_B[ni][mi],  tile_size_m,
												TILE_A[ki][mi],  tile_size_m,
										 1.0,   TILE_B[ni][ki],  tile_size_k );
						}
					}
				}
			}
		}
	}

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

	ddss_dtiled2flat( M, N, B, LDB, Bm, Bn, TILE_B );

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

	return Success;

}

