/ani/mrses

#include <stdio.h>
#include <string.h>
#include <assert.h>
#include <blas_s.h>
#include <simdmath.h>
#include <libvector.h>

#include <simdmath/sqrtf4.h>
#include <simdmath/fmaxf4.h>
#include <simdmath/divf4.h>
#include <simdmath/divf4_fast.h>
#include <simdmath/rsqrtf4.h>
#include <simdmath/recipf4.h>
#include <simdmath/recipf4_fast.h>
#include <simdmath/logf4.h>

#ifdef __SPU__
#include <spu_intrinsics.h>
#else /* not __SPU__ */
#include <altivec.h>
#endif


static vector float zero = {0, 0, 0, 0};
static vector float two = {2, 2, 2, 2};
static vector float four = {4, 4, 4, 4};
static vector float eight = {8, 8, 8, 8};


// fast versions does not bring significant performance benefits
// fmaxf between nan and zero is zero
#define spu_max _fmaxf4

// divide by zero: nan, nan; in fast mode -0, -0
#define spu_div _divf4
//#define spu_div _divf4_fast

// sqrt of negative: 0; in fast mode: undefined, sqrt of nan: big num, ?
#define spu_sqrt(a) _sqrtf4(spu_max(a, zero))
//#define spu_sqrt(a) sqrtf4_fast(spu_max(a, zero))

#define spu_recip _recipf4
//#define spu_recip _recipf4_fast

#define spu_rsqrt _rsqrtf4

#define spu_log _logf4

#define VECTOR_PRINT(var, val) \
    printf("Vector %10s is: {% 6.4e, % 6.4e, % 6.4e, % 6.4e}\n", var, spu_extract(val, 0), spu_extract(val, 1), spu_extract(val, 2), spu_extract(val, 3));

#define iVECTOR_PRINT(var, val) \
    printf("Vector %10s is: {%i, %i, %i, %i}\n", var, spu_extract(val, 0), spu_extract(val, 1), spu_extract(val, 2), spu_extract(val, 3));

#define VECTOR_PRINT_MATRIX(num, fmt, MATRIX, width, n, m) \
{ \
    int i, j; \
    for (i = 0; i < n; i++) { \
	for (j = 0; j < m; j++) \
	    printf(fmt, spu_extract(MATRIX[i * width + j], num)); \
	printf("\n"); \
    } \
    printf("\n"); \
}


inline static vector unsigned int vec_potrfU_4(vector float *A, const short int lda)
{
    vector float *pA1=A+lda, *pA2=pA1+lda, *pA3=pA2+lda;
    vector float L11 = *A, L21 = *pA1, L31 = *pA2, L41 = *pA3;
    vector float L22 = pA1[1], L32 = pA2[1], L42 = pA3[1];
    vector float L33 = pA2[2], L43 = pA3[2];
    vector float L44 = pA3[3];

    vector unsigned int errors;

    errors = spu_cmpgt(L11, zero);

    L11 = spu_rsqrt(L11);
    *A = spu_recip(L11);

    L21 = spu_mul(L21, L11);
    L31 = spu_mul(L31, L11);
    L41 = spu_mul(L41, L11);

    *pA1 = L21;
    *pA2 = L31;
    *pA3 = L41;

    L22 = spu_nmsub(L21, L21, L22);					// L22 -= L21*L21;

    errors = spu_and(spu_cmpgt(L22, zero), errors);

    L22 = spu_rsqrt(L22);
    pA1[1] = spu_recip(L22);

    L32 = spu_mul(L22, spu_nmsub(L31, L21, L32));			// (L32 - L31*L21) * L22;
    L42 = spu_mul(L22, spu_nmsub(L41, L21, L42));			// (L42 - L41*L21) * L22;
    L33 = spu_nmsub(L32, L32, spu_nmsub(L31, L31, L33));		// (L33 - L31*L31) - L32*L32;

    pA2[1] = L32;
    pA3[1] = L42;

    errors = spu_and(spu_cmpgt(L33, zero), errors);

    L33 = spu_rsqrt(L33);
    pA2[2] = spu_recip(L33);

    L43 = spu_mul(L33, spu_nmsub(L42, L32, spu_nmsub(L41, L31, L43)));		// ((L43 - L41*L31) - L42*L32) * L33;
    L44 = spu_nmsub(L43, L43, spu_nmsub(L42, L42, spu_nmsub(L41, L41, L44)));	// (((L44 - L41*L41) - L42*L42) - L43*L43);
    pA3[2] = L43;

    errors = spu_and(spu_cmpgt(L44, zero), errors);
    pA3[3] = spu_sqrt(L44);

    return errors;
}


inline static vector unsigned int vec_potrfU_3(vector float *A, const short int lda)
{
    vector float *pA1=A+lda, *pA2=pA1+lda;
    register vector float L11 = *A, L21 = *pA1, L31 = *pA2;
    register vector float L22=pA1[1], L32=pA2[1];
    register vector float L33=pA2[2];

    vector unsigned int errors;

    errors = spu_cmpgt(L11, zero);

    L11 = spu_rsqrt(L11);
    *A = spu_recip(L11);

    L21 = spu_mul(L21, L11);
    L31 = spu_mul(L31, L11);

    *pA1 = L21;
    *pA2 = L31;

    L22 = spu_nmsub(L21, L21, L22);

    errors = spu_and(spu_cmpgt(L22, zero), errors);

    L22 = spu_rsqrt(L22);
    L32 = spu_mul( spu_nmsub(L31, L21, L32), L22);  		// (L32 - L31*L21) / sqrt(L22);

    L33 = spu_nmsub(L32, L32, spu_nmsub(L31, L31, L33)); 	// (L33 - L31*L31) - L32 * L32;

    errors = spu_and(spu_cmpgt(L33, zero), errors);

    pA1[1] = spu_recip(L22);
    pA2[1] = L32;
    pA2[2] = spu_sqrt(L33);

    return errors;
}


inline static vector unsigned int vec_potrfU_2(vector float *A, const short int lda)
{
    vector float *pA1 = A + lda;
    register vector float L11 = *A, L21 = *pA1, L22 = pA1[1];

    register vector unsigned int errors;

    errors = spu_cmpgt(L11, zero);

    L11 = spu_rsqrt(L11);
    *A = spu_recip(L11);
    *pA1 = L21 = spu_mul(L21, L11);	// division by zero

    L22 = spu_nmsub(L21, L21, L22);	// L22 -= L21*L21;

    errors = spu_and(spu_cmpgt(L22, zero), errors);
    pA1[1] = spu_sqrt(L22);

    return errors;
}

inline static void vec_strsv_rlnn (short int N, const vector float *A, const int lda, vector float *X) {
    register short int i, j;

    for (i = 0; i < N; i++) {
        vector float tmp = X[i];

        for (j = 0; j < i; j++) {
            tmp = spu_nmsub(A[lda * i + j], X[j], tmp);			// tmp -= A[lda * i + j] * X[j];
        }
        X[i] = spu_div(tmp, A[(lda + 1) * i]);				// X[i] = tmp / A[lda * i + i];
    }
}

#define BLOCK_X 4
#define BLOCK_Y 8

#include <sum_across_float4.h>
#include "vec_potrf_mtxmul.h"


#define HSUM(r1,r2,r3,r4) \
	tmp1 = spu_rlqwbyte(r1, 8); \
	tmp2 = spu_rlqwbyte(r3, 8); \
 \
	tmp3 = spu_sel(r1, r3, mask1); \
	tmp4 = spu_sel(tmp1, tmp2, mask1); \
	tmp5 = spu_add(tmp3, tmp4); \
 \
	tmp1 = spu_rlqwbyte(r2, 8); \
	tmp2 = spu_rlqwbyte(r4, 8); \
 \
	tmp3 = spu_sel(r2, r4, mask1); \
	tmp4 = spu_sel(tmp1, tmp2, mask1); \
	tmp6 = spu_add(tmp3, tmp4); \
 \
	tmp1 = spu_rlqwbyte(tmp5, 4); \
	tmp2 = spu_rlqwbyte(tmp6, -4); \
 \
	tmp3 = spu_sel(tmp5, tmp6, mask2); \
	tmp4 = spu_sel(tmp1, tmp2, mask2); \
 \
	tmp5 = spu_add(tmp3, tmp4);


#define DO_MATRIX(size) \
    DECLARE_R##size \
    for (k = 0; k < K; k += BLOCK_X) { \
	const vector float *B = A + k; \
 \
	DECLARE_T##size(B) \
	COMPUTE_T##size(C) \
    } \
    SUM_T##size(C)


#include "tools.h"
void vec_ssyrk_rln_11 (short int N, short int K, const vector float *A, short int lda, float *C, short int ldc)
{
    register short int k;

    register vector float tmp1, tmp2, tmp3, tmp4, tmp5, tmp6;
    register vector unsigned int mask1 = {0, 0, 0xFFFFFFFF, 0xFFFFFFFF};
    register vector unsigned int mask2 = {0, 0xFFFFFFFF, 0, 0xFFFFFFFF};

    assert((K%16)==0);

    K >>= 2;
    lda >>= 2;

    if (N < 5) {
        if (N < 3) {
            if (N < 2) {
                DO_MATRIX(1)
            } else {
                DO_MATRIX(2)
            }
        } else {
            if (N < 4) {
                DO_MATRIX(3)
            } else {
                DO_MATRIX(4)
            }
        }
    } else if (N < 9) {
        if (N < 7) {
            if (N < 6) {
                DO_MATRIX(5)
            } else {
                DO_MATRIX(6)
            }
        } else {
            if (N < 8) {
                DO_MATRIX(7)
            } else {
                DO_MATRIX(8)
            }
        }
    } else if (N < 12) {
        if (N < 11) {
            if (N < 10) {
                DO_MATRIX(9)
            } else {
                DO_MATRIX(10)
            }
        } else {
            if (N < 12) {
                DO_MATRIX(11);
            } else {
                DO_MATRIX(12);
            }
        }
    } else {
//        assert((N%16)==0);
	    // ldc is walloc, we can have N=ralloc smaller for memory optimization, but still should use walloc here to have multiple of 16
        memset(C, 0, N * ldc * sizeof(float));
        ssyrk_spu((float*)A, C, 1, /*N*/ldc, K * 4, lda * 4, ldc);
    }

    /*
        int i,j;
        for (i = 0; i < N; i++) {
          for (j = 0; j <= i; j++) {
            register vector float temp = zero;
    	for (k = 0; k < K; k++) {
              temp = spu_madd(A[i * lda + k], A[j * lda + k], temp);
            }
            C[i * ldc + j] = spu_extract(temp, 0) +  spu_extract(temp, 1) +  spu_extract(temp, 2) +  spu_extract(temp, 3);
          }
        }
    */
}


#define DO_ROW(size) \
    DECLARE_VR##size \
    for (k = 0; k < K; k += BLOCK_X) { \
	const vector float *B = A + row * lda + k; \
    \
	DECLARE_VX(B) \
	COMPUTE_V##size((A + l * lda + k)) \
    } \
    SUM_V##size(((vector float*)(C+l)));


void vec_update_row (short int row, short int N, short int K, const vector float *A, short int lda, float *C) {
    int k, l;

    register vector float tmp1, tmp2, tmp3, tmp4, tmp5, tmp6;
    register vector unsigned int mask1 = {0, 0, 0xFFFFFFFF, 0xFFFFFFFF};
    register vector unsigned int mask2 = {0, 0xFFFFFFFF, 0, 0xFFFFFFFF};
//    vector float temp;

    assert((K%16)==0);

    K >>= 2;
    lda >>= 2;

    for (l = 0; l < N; l+=16) {
        short int rem = min(16, N - l);

        if (rem < 9) {
            if (rem < 5) {
                if (rem < 3) {
                    if (rem < 2) {
                        DO_ROW(1)
                    } else {
                        DO_ROW(2)
                    }
                } else {
                    if (rem < 4) {
                        DO_ROW(3)
                    } else {
                        DO_ROW(4)
                    }
                }
            } else
                if (rem < 7) {
                    if (rem < 6) {
                        DO_ROW(5)
                    } else {
                        DO_ROW(6)
                    }
                } else {
                    if (rem < 8) {
                        DO_ROW(7)
                    } else {
                        DO_ROW(8)
                    }
                }
        } else {
            if (rem < 12) {
                if (rem < 11) {
                    if (rem < 10) {
                        DO_ROW(9)
                    } else {
                        DO_ROW(10)
                    }
                } else {
                    if (rem < 12) {
                        DO_ROW(11);
                    } else {
                        DO_ROW(12);
                    }
                }
            } else {
                if (rem < 15) {
                    if (rem < 14) {
                        DO_ROW(13)
                    } else {
                        DO_ROW(14)
                    }
                } else {
                    if (rem < 16) {
                        DO_ROW(15);
                    } else {
                        DO_ROW(16);
                    }
                }
            }
        }
    }


    /*
        DECLARE_VR5
        for (k = 0; k < K; k+= BLOCK_X) {
    	const vector float *B = A + row * lda + k;

    	DECLARE_VX(B)
    	COMPUTE_V5((A + l * lda + k))
        }
        SUM_V5(((vector float*)(C+l)));
    */


    /*
        vector float temp[N];
        memset(temp, 0, sizeof(vector float)*N);
        for (k = 0; k < K; k++) {

    	register vector float Arow = A[row * lda + k];
    	for (l = 0; l < N; l++) {
    	    temp[l] = spu_madd(Arow, A[l * lda + k], temp[l]);
    	    C[l] = spu_extract(temp[l], 0) +  spu_extract(temp[l], 1) +  spu_extract(temp[l], 2) +  spu_extract(temp[l], 3);
    	}
    	if ((k%4)==3) {
    	    printf("Step: %i\n",k);
        	    PRINT_MATRIX("% 6.4f ", C, 1, 1, 16);
    	}
        }
        for (l = 0; l < N; l++) {
    	C[l] = spu_extract(temp[l], 0) +  spu_extract(temp[l], 1) +  spu_extract(temp[l], 2) +  spu_extract(temp[l], 3);
        }
        PRINT_MATRIX("% 6.4f ", C, 1, 1, 16);
    */

    /*
        for (l = 0; l < N; l++) {
    	if (l < row) {
    	    i = l; j = row;
    	} else {
    	    i = row; j = l;
    	}

            register vector float tmp = zero;
    	for (k = 0; k < K; k++) {
              tmp = spu_madd(A[i * lda + k], A[j * lda + k], tmp);
            }
            C[l] = spu_extract(tmp, 0) +  spu_extract(tmp, 1) +  spu_extract(tmp, 2) +  spu_extract(tmp, 3);
        }

        PRINT_MATRIX("% 6.4f ", C, 1, 1, 16);
    */
}


inline static void vec_spotrf_step(const short int M, const short int N, vector float *A, const short int lda) {
    register short int i, j, k;

    vector float *B = A + M * lda;
    vector float *C = B + M;


    for (i = 0; i < N; i++) {
        for (j = 0; j < M; j++) {
            vector float sum = {0, 0, 0, 0};

            for (k = 0; k < j; k++) {
                sum = spu_madd(A[j * lda + k],  B[i * lda + k], sum);
            }

            B[i * lda + j] =  spu_div(spu_sub(B[i * lda + j], sum), A[j * lda + j]);
        }

        for (j = 0; j <= i; j++) {
            vector float temp = {0, 0, 0, 0};

            for (k = 0; k < M; k++) {
                temp = spu_madd(B[i * lda + k],  B[j * lda + k], temp);
            }

            C[i * lda + j] = spu_sub(C[i * lda + j], temp);
        }
    }
}


// we ignore errors if it will lead to 0 determinant
vector unsigned int vec_spotrf_u(short int N, vector float *A, short int lda) {
    short int Nleft, Nright;

    vector unsigned int errors;

    //PRINT_MULTIMATRIX(0, "% 7.4f  ", ((float*)A), 5, 5, 5);

    if (N > 4) {
        Nleft = N >> 1;
        Nright = N - Nleft;

//    printf("passed matrix\n");
//    VECTOR_PRINT_MATRIX(0, "% 8.4f  ", A, 5, 5, 5);

        errors = vec_spotrf_u(Nleft, A, lda);

//    printf("left\n");
//    VECTOR_PRINT_MATRIX(0, "% 8.4f  ", A, 5, 5, 5);

        vec_spotrf_step(Nleft, Nright, A, lda);
//    printf("step\n");
//    VECTOR_PRINT_MATRIX(0, "% 8.4f  ", A, 5, 5, 5);

        spu_and(vec_spotrf_u(Nright, A + (lda + 1) * Nleft, lda), errors);
//    printf("right\n");
//    VECTOR_PRINT_MATRIX(0, "% 8.4f  ", A, 5, 5, 5);

    }
    else if (N==4) errors = vec_potrfU_4(A, lda);
    else if (N==3) errors = vec_potrfU_3(A, lda);
    else if (N==2) errors = vec_potrfU_2(A, lda);
    else if (N==1) {
        errors = spu_cmpgt(*A, zero);
        *A = spu_sqrt(*A);
    } else return spu_cmpgt(zero, zero); // anything

    return errors;
}

vector float vec_rcorr(short int N, vector float *C, vector float *Ca, vector float *Cb) {
    short int i;
    short int step = N + 1;
    short int end = N * N;
    vector float detC, detAB;

//    VECTOR_PRINT_MATRIX(1, "%6.4f", C, 5, 5, 5);

    detC = C[0];
    detAB = Ca[0] * Cb[0];
    for (i = step; i < end; i+= step) {
        detAB = spu_mul(detAB, spu_mul(Ca[i], Cb[i]));
        detC = spu_mul(detC, C[i]);
    }

    /*
        VECTOR_PRINT("detAB", detAB);
        VECTOR_PRINT("detC", detC);
        VECTOR_PRINT("detC^2", spu_mul(detC, detC));
        VECTOR_PRINT("C/AB", spu_div(spu_mul(detC, detC), detAB));
        VECTOR_PRINT("log", spu_log(spu_div(spu_mul(detC, detC), detAB)));
    */

    return spu_mul(two, spu_log(spu_div(spu_mul(detC, detC), detAB)));
}

vector float vec_rmahal(short int N, vector float *C, vector float *X) {
    register short int i;
    vector float sum = zero;

    vec_strsv_rlnn(N, C, N, X);
    for (i = 0; i < N; i++) {
        vector float val = X[i];
        sum = spu_madd(val, val, sum);
    }

    return sum;
}

vector float vec_bhata(vector float rcorr, vector float rmahal) {
    return spu_add(spu_div(rmahal, eight), spu_div(rcorr, four));
}