#include <u.h>
#include <libn.h>
#include <libmath.h>
#include <vendor/blas/cblas.h>

#include <x86intrin.h>
#include <time.h>

// -----------------------------------------------------------------------
// level one

/*
 * Rotate vector
 * x = cos*x + sin*y
 * y = cos*x - sin*y
 */ 

static
void
rot_kernel8_avx2(int n, double *x, double *y, double cos, double sin)
{
    register int i;
    __m256d x256, y256;
    __m128d cos128, sin128;
    __m256d cos256, sin256;

    cos128 = _mm_load_sd(&cos);
    cos256 = _mm256_broadcastsd_pd(cos128);

    sin128 = _mm_load_sd(&sin);
    sin256 = _mm256_broadcastsd_pd(sin128);

    for (i = 0; i < n; i+=8) {
        x256 = _mm256_loadu_pd(x+i+0);
        y256 = _mm256_loadu_pd(y+i+0);
        _mm256_storeu_pd(x+i+0, cos256 * x256 + sin256 * y256);
        _mm256_storeu_pd(y+i+0, cos256 * y256 - sin256 * x256);

        x256 = _mm256_loadu_pd(x+i+4);
        y256 = _mm256_loadu_pd(y+i+4);
        _mm256_storeu_pd(x+i+4, cos256 * x256 + sin256 * y256);
        _mm256_storeu_pd(y+i+4, cos256 * y256 - sin256 * x256);
    }
}

static
void
rot_kernel8(int n, double *x, double *y, double cos, double sin)
{
    register int i;
    register double tmp;

    for (i = 0; i < n; i+=8) {
        tmp = x[i+0], x[i+0] = cos*x[i+0] + sin*y[i+0], y[i+0] = cos*y[i+0] - sin*tmp;
        tmp = x[i+1], x[i+1] = cos*x[i+1] + sin*y[i+1], y[i+1] = cos*y[i+1] - sin*tmp;
        tmp = x[i+2], x[i+2] = cos*x[i+2] + sin*y[i+2], y[i+2] = cos*y[i+2] - sin*tmp;
        tmp = x[i+3], x[i+3] = cos*x[i+3] + sin*y[i+3], y[i+3] = cos*y[i+3] - sin*tmp;
        tmp = x[i+4], x[i+4] = cos*x[i+4] + sin*y[i+4], y[i+4] = cos*y[i+4] - sin*tmp;
        tmp = x[i+5], x[i+5] = cos*x[i+5] + sin*y[i+5], y[i+5] = cos*y[i+5] - sin*tmp;
        tmp = x[i+6], x[i+6] = cos*x[i+6] + sin*y[i+6], y[i+6] = cos*y[i+6] - sin*tmp;
        tmp = x[i+7], x[i+7] = cos*x[i+7] + sin*y[i+7], y[i+7] = cos*y[i+7] - sin*tmp;
    }
}

void
blas·rot(int len, double *x, double *y, double cos, double sin)
{
    register int n;
    register double tmp;

    n = len & ~7;
    rot_kernel8_avx2(n, x, y, cos, sin);

    for (; n < len; n++) {
        tmp = x[n], x[n] = cos*x[n] + sin*y[n], y[n] = cos*y[n]- sin*tmp;
    }
}

/*
 * compute givens rotate vector
 * --      --
 * |+cos -sin| | a | = | r |
 * |-sin +cos| | b | = | 0 |
 * --      --
 */ 

void
blas·rotg(double *a, double *b, double *cos, double *sin)
{
    double abs_a, abs_b, r, rho, scale, z;

    abs_a = math·abs(*a);
    abs_b = math·abs(*b);
    rho   = abs_a > abs_b ? *a : *b;
    scale = abs_a + abs_b;

    if (scale == 0) {
        *cos = 1, *sin = 0;
        r = 0.;
        z = 0.;
    } else {
        r    = math·sgn(rho) * scale * math·sqrt(math·pow(abs_a/scale, 2) + math·pow(abs_b/scale, 2));
        *cos = *a / r;
        *sin = *b / r;
        if (abs_a > abs_b)
            z = *sin;
        else if (abs_b >= abs_a && *cos != 0)
            z = 1/(*cos);
        else 
            z = 1.;
    }
    *a = r;
    *b = z;
}

/*
 * modified Givens rotation of points in plane
 * operates on len points
 *
 * params = [flag, h11, h12, h21, h22]
 * @flag = -1:
 *      H -> [ [h11, h12], [h21, h22] ]
 * @flag = 0.0:
 *      H -> [ [1, h12], [h21, 1] ]
 * @flag = +1:
 *      H -> [ [h11, 1], [-1, h22] ]
 * @flag = -2:
 *      H -> [ [1, 0], [0, 1] ]
 *
 *  Replaces:
 *      x -> H11 * x + H12 * y
 *      y -> H21 * x + H22 * y
 */

static
void
rotm_kernel8(int n, double *x, double *y, double H[4])
{
    register int i;
    register double tmp;

    for (i = 0; i < n; i+=8) {
        tmp = x[i+0], x[i+0] = H[0]*x[i+0] + H[1]*y[i+0], y[i+0] = H[2]*y[i+0] + H[3]*tmp;
        tmp = x[i+1], x[i+1] = H[0]*x[i+1] + H[1]*y[i+1], y[i+1] = H[2]*y[i+1] + H[3]*tmp;
        tmp = x[i+2], x[i+2] = H[0]*x[i+2] + H[1]*y[i+2], y[i+2] = H[2]*y[i+2] + H[3]*tmp;
        tmp = x[i+3], x[i+3] = H[0]*x[i+3] + H[1]*y[i+3], y[i+3] = H[2]*y[i+3] + H[3]*tmp;
        tmp = x[i+4], x[i+4] = H[0]*x[i+4] + H[1]*y[i+4], y[i+4] = H[2]*y[i+4] + H[3]*tmp;
        tmp = x[i+5], x[i+5] = H[0]*x[i+5] + H[1]*y[i+5], y[i+5] = H[2]*y[i+5] + H[3]*tmp;
        tmp = x[i+6], x[i+6] = H[0]*x[i+6] + H[1]*y[i+6], y[i+6] = H[2]*y[i+6] + H[3]*tmp;
        tmp = x[i+7], x[i+7] = H[0]*x[i+7] + H[1]*y[i+7], y[i+7] = H[2]*y[i+7] + H[3]*tmp;
    }
}

error
blas·rotm(int len, double *x, double *y, double p[5])
{
    int    n, flag;
    double tmp, H[4];

    flag = math·round(p[0]);
    switch (flag) {
        case  0: H[0] = p[1], H[1] = p[2], H[2] = p[3], H[3] = p[4]; break;
        case -1: H[0] = +1,   H[1] = p[2], H[2] = p[3], H[3] = +1;   break;
        case +1: H[0] = p[1], H[1] = +1,   H[2] = -1,   H[3] = p[4]; break;
        case -2: H[0] = +1,   H[1] = 0,    H[2] = 0,    H[3] = +1;   break;
        default:
            errorf("rotm: flag '%d' unrecognized", flag);
            return 1;
    }

    n = len & ~7;
    rotm_kernel8(n, x, y, H);

    for (; n < len; n++) {
        tmp = x[n], x[n] = H[0]*x[n] + H[1]*y[n], y[n] = H[2]*y[n] + H[3]*tmp;
    }

    return 0;
}


/*
 * Scale vector
 * x = ax
 */ 

static
void
scale_kernel8_avx2(int n, double *x, double a)
{
    __m128d a128;
    __m256d a256;
    register int i;

    a128 = _mm_load_sd(&a);
    a256 = _mm256_broadcastsd_pd(a128);
    for (i = 0; i < n; i += 8) {
        _mm256_storeu_pd(x+i+0, a256 * _mm256_loadu_pd(x+i+0));
        _mm256_storeu_pd(x+i+4, a256 * _mm256_loadu_pd(x+i+4));
    }
}

static
void
scale_kernel8(int n, double *x, double a)
{
    register int i;
    for (i = 0; i < n; i += 8) {
        x[i+0] *= a;
        x[i+1] *= a;
        x[i+2] *= a;
        x[i+3] *= a;
        x[i+4] *= a;
        x[i+5] *= a;
        x[i+6] *= a;
        x[i+7] *= a;
    }
}

void
blas·scale(int len, double *x, double a)
{
    int n;

    n = len & ~7;
    scale_kernel8_avx2(n, x, a);

    for (; n < len; n++) {
        x[n] *= a;
    }
}

/*
 * Daxpy
 * y = ax + y
 */
 

static
void
daxpy_kernel8_avx2(int n, double a, double *x, double *y)
{
    __m128d a128;
    __m256d a256;
    register int i;

    a128 = _mm_load_sd(&a);
    a256 = _mm256_broadcastsd_pd(a128);
    for (i = 0; i < n; i += 8) {
        _mm256_storeu_pd(y+i+0, _mm256_loadu_pd(y+i+0) + a256 * _mm256_loadu_pd(x+i+0));
        _mm256_storeu_pd(y+i+4, _mm256_loadu_pd(y+i+4) + a256 * _mm256_loadu_pd(x+i+4));
    }
}

static
void
daxpy_kernel8(int n, double a, double *x, double *y)
{
    register int i;
    for (i = 0; i < n; i += 8) {
        y[i+0] += a*x[i+0];
        y[i+1] += a*x[i+1];
        y[i+2] += a*x[i+2];
        y[i+3] += a*x[i+3];
        y[i+4] += a*x[i+4];
        y[i+5] += a*x[i+5];
        y[i+6] += a*x[i+6];
        y[i+7] += a*x[i+7];
    }
}

void
blas·daxpy(int len, double a, double *x, double *y)
{
    int n;

    n = len & ~7;
    daxpy_kernel8_avx2(n, a, x, y);

    for (; n < len; n++) {
        y[n] += a*x[n];
    }
}

/************************************************
 * Dot product
 * x·y
 ***********************************************/

static
double
dot_kernel8_avx2(int len, double *x, double *y)
{
    register int i;
    __m256d sum[4];
    __m128d res;

    for (i = 0; i < arrlen(sum); i++) {
        sum[i] = _mm256_setzero_pd();
    }

    for (i = 0; i < len; i += 16) {
        sum[0] += _mm256_loadu_pd(x+i+0)  * _mm256_loadu_pd(y+i+0);
        sum[1] += _mm256_loadu_pd(x+i+4)  * _mm256_loadu_pd(y+i+4);
        sum[2] += _mm256_loadu_pd(x+i+8)  * _mm256_loadu_pd(y+i+8);
        sum[3] += _mm256_loadu_pd(x+i+12) * _mm256_loadu_pd(y+i+12);
    }

    sum[0] += sum[1] + sum[2] + sum[3];

    res = _mm_add_pd(_mm256_extractf128_pd(sum[0], 0), _mm256_extractf128_pd(sum[0], 1));
    res = _mm_hadd_pd(res, res);

    return res[0];
}

static
double
dot_kernel8_fma3(int len, double *x, double *y)
{
    register int i;
    __m256d sum[4];
    __m128d res;

    for (i = 0; i < arrlen(sum); i++) {
        sum[i] = _mm256_setzero_pd();
    }

    for (i = 0; i < len; i += 16) {
        sum[0] = _mm256_fmadd_pd(_mm256_loadu_pd(x+i+0),  _mm256_loadu_pd(y+i+0), sum[0]);
        sum[1] = _mm256_fmadd_pd(_mm256_loadu_pd(x+i+4),  _mm256_loadu_pd(y+i+4), sum[1]);
        sum[2] = _mm256_fmadd_pd(_mm256_loadu_pd(x+i+8),  _mm256_loadu_pd(y+i+8), sum[2]);
        sum[3] = _mm256_fmadd_pd(_mm256_loadu_pd(x+i+12), _mm256_loadu_pd(y+i+12), sum[3]);
    }

    sum[0] += sum[1] + sum[2] + sum[3];

    res = _mm_add_pd(_mm256_extractf128_pd(sum[0], 0), _mm256_extractf128_pd(sum[0], 1));
    res = _mm_hadd_pd(res, res);

    return res[0];
}

static
double
dot_kernel8(int len, double *x, double *y)
{
    double res;
    register int i;

    for (i = 0; i < len; i += 8) {
        res += x[i]   * y[i]   +
               x[i+1] * y[i+1] +
               x[i+2] * y[i+2] +
               x[i+3] * y[i+3] +
               x[i+4] * y[i+4] +
               x[i+5] * y[i+5] +
               x[i+6] * y[i+6] +
               x[i+7] * y[i+7];
    }

    return res;
}

double
blas·dot(int len, double *x, double *y)
{
    int i, n;
    double res;

    n = len & ~15; // neat trick
    res = dot_kernel8_fma3(n, x, y);

    for (i = n; i < len; i++) {
        res += x[i] * y[i];
    }

    return res;
}

// -----------------------------------------------------------------------
// level two

/* 
 * Notation: (number of rows) x (number of columns) _ unroll factor
 *            N => variable we sum over
 */


/* 
 * Affine transformation
 * y = aMx + by
 */
static
void
gemv_kernel4xN_4_avx2(int ncol, double **row, double *x, double *y)
{
    int c;
    __m128d hr;
    __m256d x256, r256[4];

    for (c = 0; c < 4; c++) {
        r256[c] = _mm256_setzero_pd();
    }

    for (c = 0; c < ncol; c += 4) {
        x256     = _mm256_loadu_pd(x+c);
        r256[0] += x256 * _mm256_loadu_pd(row[0] + c);
        r256[1] += x256 * _mm256_loadu_pd(row[1] + c);
        r256[2] += x256 * _mm256_loadu_pd(row[2] + c);
        r256[3] += x256 * _mm256_loadu_pd(row[3] + c);
    }

    for (c = 0; c < 4; c++) {
        hr   = _mm_add_pd(_mm256_extractf128_pd(r256[c], 0), _mm256_extractf128_pd(r256[c], 1));
        hr   = _mm_hadd_pd(hr, hr);
        y[c] = hr[0];
    }
}

static
void
gemv_kernel4xN_4(int ncol, double **row, double *x, double *y)
{
    int c;
    double res[4];

    res[0] = 0.;
    res[1] = 0.;
    res[2] = 0.;
    res[3] = 0.;

    for (c = 0; c < ncol; c += 4) {
        res[0] += row[0][c+0]*x[c+0] + row[0][c+1]*x[c+1] + row[0][c+2]*x[c+2] + row[0][c+3]*x[c+3];
        res[1] += row[1][c+0]*x[c+0] + row[1][c+1]*x[c+1] + row[1][c+2]*x[c+2] + row[1][c+3]*x[c+3];
        res[2] += row[2][c+0]*x[c+0] + row[2][c+1]*x[c+1] + row[2][c+2]*x[c+2] + row[2][c+3]*x[c+3];
        res[3] += row[3][c+0]*x[c+0] + row[3][c+1]*x[c+1] + row[3][c+2]*x[c+2] + row[3][c+3]*x[c+3];
    }

    y[0] = res[0];
    y[1] = res[1];
    y[2] = res[2];
    y[3] = res[3];
}

static
void
gemv_kernel1xN_4(int ncol, double *row, double *x, double *y)
{
    int c;
    double res;

    res = 0.;
    for (c = 0; c < ncol; c += 4) {
        res += row[c+0]*x[c+0] + row[c+1]*x[c+1] + row[c+2]*x[c+2] + row[c+3]*x[c+3];
    }

    y[0] = res;
}

error
blas·gemv(int nrow, int ncol, double a, double *m, double *x, double b, double *y) 
{
    int c, r, nr, nc;
    double *row[4], res[4];

    nr = nrow & ~3;
    nc = ncol & ~3;
    for (r = 0; r < nrow; r += 4) {
        row[0] = m + (r * (ncol+0));
        row[1] = m + (r * (ncol+1));
        row[2] = m + (r * (ncol+2));
        row[3] = m + (r * (ncol+3));

        gemv_kernel4xN_4_avx2(ncol, row, x + r, res);

        for (c = nc; c < ncol; c++) {
            res[0] += row[0][c];
            res[1] += row[1][c];
            res[2] += row[2][c];
            res[3] += row[3][c];
        }

        y[r+0] = res[0] + b*y[r+0];
        y[r+1] = res[1] + b*y[r+1];
        y[r+2] = res[2] + b*y[r+2];
        y[r+3] = res[3] + b*y[r+3];
    }

    for (; r < nrow; r++) {
        gemv_kernel1xN_4(nrow, m + (r * ncol), x + r, res);
        y[r] = res[0] + b*y[r];
    }

    return 0;
}

/* 
 * rank one addition
 * M = ax(y^T) + M
 * TODO: vectorize kernel
 */

void
blas·ger(int nrow, int ncol, double a, double *x, double *y, double *m)
{
    int i, j;

    for (i = 0; i < nrow; i++) {
        for (j = 0; j < ncol; j++) {
            m[i+ncol*j] += a * x[i] * y[j];
        }
    }
}

/* 
 * symmetric rank one addition
 * M = ax(x^T) + M
 * TODO: vectorize kernel
 */

void
blas·syr(int nrow, int ncol, double a, double *x, double *m)
{
    int i, j;

    for (i = 0; i < nrow; i++) {
        for (j = 0; j < ncol; j++) {
            m[i+ncol*j] += a * x[i] * x[j];
        }
    }
}

// -----------------------------------------------------------------------
// level three

/* 
 * matrix multiplication
 * m3 = a(m1 * m2) + b(m3)
 * einstein notation:
 * m3_{ij} = a m1_{ik} m2_{kj} + b m3_{ij}
 *
 * n1 = # rows of m1 = # rows of m3
 * n2 = # cols of m2 = # cols of m3
 * n3 = # cols of m1 = # rows of m2
 *
 * TODO: Right now we are 2x slower than OpenBLAS.
 *       This is because we use a very simple algorithm.
 */

void
blas·gemm(int n1, int n2, int n3, double a, double *m1, double *m2, double b, double *m3)
{
    int i, j, k, len;

    // TODO: Is there anyway this computation can be integrated into the one below?
    for (i = 0; i < n1; i++) {
        for (j = 0; j < n2; j++) {
            m3[i + n2*j] *= b;
        }
    }

    len = n1 & ~7;
    for (j = 0; j < n2; j++) {
        for (k = 0; k < n3; k++) {
            daxpy_kernel8_avx2(len, a * m2[k + n2*j], m1 + n3*k, m3 + n2*j);

            // remainder
            for (i = len; i < n1; i++) {
                m3[i + n2*j] += a * m1[i + n3*k] * m2[k + n2*j]; 
            }
        }
    }
}

#define NITER 5000
#define NCOL  10007
#define NROW  10007

error
test·gemm()
{
    int i, n;
    clock_t t;

    double *x, *y, *m[3];
    double res[2], tprof[2];

    // openblas_set_num_threads(1);

    x    = malloc(sizeof(*x)*NCOL);
    y    = malloc(sizeof(*x)*NCOL);
    m[0] = malloc(sizeof(*x)*NCOL*NROW);
    m[1] = malloc(sizeof(*x)*NCOL*NROW);
    m[2] = malloc(sizeof(*x)*NCOL*NROW);

    tprof[0] = 0;
    tprof[1] = 0;

    for (n = 0; n < NITER; n++) {
        for (i = 0; i < NCOL; i++) {
            x[i] = i*i+1;
            y[i] = i+1;
        }

        for (i = 0; i < NCOL*NROW; i++) {
            m[0][i] = i/(NCOL*NROW);
            m[1][i] = i*i/(NCOL*NROW*NCOL*NROW);
            m[2][i] = 1;
        }

        t = clock();
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, NCOL, NROW, NROW, 1.2, m[0], NROW, m[1], NROW, 2.8, m[2], NROW);
        t = clock() - t;
        tprof[1] += 1000.*t/CLOCKS_PER_SEC;
        res[1]    = cblas_ddot(NROW*NCOL, m[2], 1, m[2], 1);

        for (i = 0; i < NCOL; i++) {
            x[i] = i*i+1;
            y[i] = i+1;
        }

        for (i = 0; i < NCOL*NROW; i++) {
            m[0][i] = i/(NCOL*NROW);
            m[1][i] = i*i/(NCOL*NROW*NCOL*NROW);
            m[2][i] = 1;
        }

        t = clock();
        blas·gemm(NROW, NROW, NROW, 1.2, m[0], m[1], 2.8, m[2]);
        t = clock() - t;
        tprof[0] += 1000.*t/CLOCKS_PER_SEC;
        res[0]    = blas·dot(NROW*NCOL, m[2], m[2]);
    }
    printf("mean time/iteration (naive): %fms\n", tprof[0]/NITER);
    printf("--> result (naive): %f\n", res[0]);
    printf("mean time/iteration (oblas): %fms\n", tprof[1]/NITER);
    printf("--> result (oblas): %f\n", res[1]);

    return 0;
}

void
print·array(int n, double *x)
{
    double *end;
    printf("[");
    for (end=x+n; x != end; ++x) {
        printf("%f,", *x);
    }
    printf("]\n");
}

error
main()
{
    int i, n;
    double *x, *y;
    double tprof[2];
    clock_t t;

    x = malloc(sizeof(*x)*NCOL);
    y = malloc(sizeof(*x)*NCOL);

    for (n = 0; n < NITER; n++) {
        for (i = 0; i < NCOL; i++) {
            x[i] = i*i+1;
            y[i] = i+1;
        }

        t = clock();
        blas·rot(NCOL, x, y, .707, .707); 
        t = clock() - t;
        tprof[0] += 1000.*t/CLOCKS_PER_SEC;

        for (i = 0; i < NCOL; i++) {
            x[i] = i*i+1;
            y[i] = i+1;
        }

        t = clock();
        cblas_drot(NCOL, x, 1, y, 1, .707, .707);
        t = clock() - t;
        tprof[1] += 1000.*t/CLOCKS_PER_SEC;
    }

    printf("mean time/iteration (naive): %fms\n", tprof[0]/NITER);
    printf("mean time/iteration (oblas): %fms\n", tprof[1]/NITER);

    double a, b, c, s;

    a = 10.234, b = 2.;
    cblas_drotg(&a, &b, &c, &s);
    printf("%f, %f, %f, %f\n", a, b, c, s);

    a = 10.234, b = 2.;
    blas·rotg(&a, &b, &c, &s);
    printf("%f, %f, %f, %f\n", a, b, c, s);
}