#include "sn_target.h"

#if SN_IS_ARCH(SN_ARCH_X86_64)

#include "sn_intrinsic.h"
#include "sn_intrincontext.x86.h"

#define _FMA_SELECT(i) _MM_SHUFFLE(i, i, i, i)

// ans = left * right + ans
// right is not transposed
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, sse, , fma)
{
        const __m128 a0 = _mm_load_ps((float *)&left->f32_x[0]);
        const __m128 a1 = _mm_load_ps((float *)&left->f32_x[1]);
        const __m128 a2 = _mm_load_ps((float *)&left->f32_x[2]);
        const __m128 a3 = _mm_load_ps((float *)&left->f32_x[3]);

        __m128 c0 = _mm_load_ps((float *)&ans->f32_x[0]);
        __m128 c1 = _mm_load_ps((float *)&ans->f32_x[1]);
        __m128 c2 = _mm_load_ps((float *)&ans->f32_x[2]);
        __m128 c3 = _mm_load_ps((float *)&ans->f32_x[3]);

        __m128 bcol;

        /* Column 0 */
        bcol = _mm_load_ps((float *)&right->f32_x[0]);
        c0 = _mm_add_ps(c0, _mm_mul_ps(a0, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(0))));
        c0 = _mm_add_ps(c0, _mm_mul_ps(a1, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(1))));
        c0 = _mm_add_ps(c0, _mm_mul_ps(a2, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(2))));
        c0 = _mm_add_ps(c0, _mm_mul_ps(a3, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(3))));

        /* Column 1 */
        bcol = _mm_load_ps((float *)&right->f32_x[1]);
        c1 = _mm_add_ps(c1, _mm_mul_ps(a0, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(0))));
        c1 = _mm_add_ps(c1, _mm_mul_ps(a1, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(1))));
        c1 = _mm_add_ps(c1, _mm_mul_ps(a2, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(2))));
        c1 = _mm_add_ps(c1, _mm_mul_ps(a3, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(3))));

        /* Column 2 */
        bcol = _mm_load_ps((float *)&right->f32_x[2]);
        c2 = _mm_add_ps(c2, _mm_mul_ps(a0, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(0))));
        c2 = _mm_add_ps(c2, _mm_mul_ps(a1, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(1))));
        c2 = _mm_add_ps(c2, _mm_mul_ps(a2, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(2))));
        c2 = _mm_add_ps(c2, _mm_mul_ps(a3, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(3))));

        /* Column 3 */
        bcol = _mm_load_ps((float *)&right->f32_x[3]);
        c3 = _mm_add_ps(c3, _mm_mul_ps(a0, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(0))));
        c3 = _mm_add_ps(c3, _mm_mul_ps(a1, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(1))));
        c3 = _mm_add_ps(c3, _mm_mul_ps(a2, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(2))));
        c3 = _mm_add_ps(c3, _mm_mul_ps(a3, _mm_shuffle_ps(bcol, bcol, _FMA_SELECT(3))));

        ans->f32_x[0] = c0;
        ans->f32_x[1] = c1;
        ans->f32_x[2] = c2;
        ans->f32_x[3] = c3;
}

DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, avx, _fma, fma)
{
        register __m256 a01 = _mm256_load_ps((float *)&left->f32_y[0]);
        register __m256 a23 = _mm256_load_ps((float *)&left->f32_y[1]);

        register __m256 b01 = _mm256_load_ps((float *)&right->f32_y[0]);
        register __m256 b23 = _mm256_load_ps((float *)&right->f32_y[1]);

        register __m256 c01 = _mm256_load_ps((float *)&ans->f32_y[0]);
        register __m256 c23 = _mm256_load_ps((float *)&ans->f32_y[1]);

        //group 0
        register __m256 b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(0));
        register __m256 b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(0));
        c01 = _mm256_fmadd_ps(a01, b01_perm, c01);
        c23 = _mm256_fmadd_ps(a23, b23_perm, c23);

        //group 2
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(2));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(2));
        c01 = _mm256_fmadd_ps(a23, b01_perm, c01);
        c23 = _mm256_fmadd_ps(a01, b23_perm, c23);

        //group 1
        register __m256 a12 = _mm256_permute2f128_ps(a01, a23, 0x21);
        register __m256 a03 = _mm256_permute2f128_ps(a01, a23, 0x30);
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(1));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(1));
        c01 = _mm256_fmadd_ps(a12, b01_perm, c01);
        c23 = _mm256_fmadd_ps(a03, b23_perm, c23);

        //group 3
        a12 = _mm256_permute2f128_ps(a01, a23, 0x21);
        a03 = _mm256_permute2f128_ps(a01, a23, 0x30);
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(3));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(3));
        c01 = _mm256_fmadd_ps(a03, b01_perm, c01);
        c23 = _mm256_fmadd_ps(a12, b23_perm, c23);

        ans->f32_y[0] = c01;
        ans->f32_y[1] = c23;
}

DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, avx, , fma)
{
        register __m256 a01 = _mm256_load_ps((float *)&left->f32_y[0]);
        register __m256 a23 = _mm256_load_ps((float *)&left->f32_y[1]);

        register __m256 b01 = _mm256_load_ps((float *)&right->f32_y[0]);
        register __m256 b23 = _mm256_load_ps((float *)&right->f32_y[1]);

        register __m256 c01 = _mm256_load_ps((float *)&ans->f32_y[0]);
        register __m256 c23 = _mm256_load_ps((float *)&ans->f32_y[1]);

        //group 0
        register __m256 b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(0));
        register __m256 b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(0));
        c01 = _mm256_add_ps(_mm256_mul_ps(a01, b01_perm), c01);
        c23 = _mm256_add_ps(_mm256_mul_ps(a23, b23_perm), c23);

        //group 2
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(2));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(2));
        c01 = _mm256_add_ps(_mm256_mul_ps(a23, b01_perm), c01);
        c23 = _mm256_add_ps(_mm256_mul_ps(a01, b23_perm), c23);

        //group 1
        register __m256 a12 = _mm256_permute2f128_ps(a01, a23, 0x21);
        register __m256 a03 = _mm256_permute2f128_ps(a01, a23, 0x30);
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(1));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(1));
        c01 = _mm256_add_ps(_mm256_mul_ps(a12, b01_perm), c01);
        c23 = _mm256_add_ps(_mm256_mul_ps(a03, b23_perm), c23);

        //group 3
        a12 = _mm256_permute2f128_ps(a01, a23, 0x21);
        a03 = _mm256_permute2f128_ps(a01, a23, 0x30);
        b01_perm = _mm256_shuffle_ps(b01, b01, _FMA_SELECT(3));
        b23_perm = _mm256_shuffle_ps(b23, b23, _FMA_SELECT(3));
        c01 = _mm256_add_ps(_mm256_mul_ps(a03, b01_perm), c01);
        c23 = _mm256_add_ps(_mm256_mul_ps(a12, b23_perm), c23);

        ans->f32_y[0] = c01;
        ans->f32_y[1] = c23;
}

DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, 512, , fma)
{
        register __m512 a = _mm512_load_ps((float *)left);
        register __m512 b = _mm512_load_ps((float *)right);
        register __m512 c = _mm512_load_ps((float *)ans);

        register __m512 a_perm = _mm512_shuffle_f32x4(a, a, _MM_SHUFFLE(3, 2, 1, 0));
        register __m512 b_perm = _mm512_permute_ps(b, _FMA_SELECT(0));
        c = _mm512_fmadd_ps(a, b_perm, c);

        a_perm = _mm512_shuffle_f32x4(a, a, _MM_SHUFFLE(2, 1, 0, 3));
        b_perm = _mm512_permute_ps(b, _FMA_SELECT(1));
        c = _mm512_fmadd_ps(a, b_perm, c);

        a_perm = _mm512_shuffle_f32x4(a, a, _MM_SHUFFLE(1, 0, 3, 2));
        b_perm = _mm512_permute_ps(b, _FMA_SELECT(2));
        c = _mm512_fmadd_ps(a, b_perm, c);

        a_perm = _mm512_shuffle_f32x4(a, a, _MM_SHUFFLE(0, 3, 2, 1));
        b_perm = _mm512_permute_ps(b, _FMA_SELECT(3));
        c = _mm512_fmadd_ps(a, b_perm, c);

        _mm512_store_ps((float *)ans, c);
}

DECLARE_PNL_PROC_LEFT_RIGHT_ANS(double, sse, , fma)
{
        for (int col = 0; col < 4; ++col) {
                /* current answer column (two chunks: rows 0-1 and 2-3) */
                register __m128d c_lo = _mm_castps_pd(PANEL_COL_VEC((*ans), col, 0));
                register __m128d c_hi = _mm_castps_pd(PANEL_COL_VEC((*ans), col, 1));

                /* right matrix column split into two chunks */
                register __m128d b_lo = _mm_castps_pd(PANEL_COL_VEC((*right), col, 0));
                register __m128d b_hi = _mm_castps_pd(PANEL_COL_VEC((*right), col, 1));

                /* k = 0 */
                {
                        register __m128d a_lo = _mm_castps_pd(PANEL_COL_VEC((*left), 0, 0));
                        register __m128d a_hi = _mm_castps_pd(PANEL_COL_VEC((*left), 0, 1));
                        register __m128d b = _mm_shuffle_pd(b_lo, b_lo, 0x0);
                        c_lo = _mm_add_pd(c_lo, _mm_mul_pd(a_lo, b));
                        c_hi = _mm_add_pd(c_hi, _mm_mul_pd(a_hi, b));
                }

                /* k = 1 */
                {
                        register __m128d a_lo = _mm_castps_pd(PANEL_COL_VEC((*left), 1, 0));
                        register __m128d a_hi = _mm_castps_pd(PANEL_COL_VEC((*left), 1, 1));
                        register __m128d b = _mm_shuffle_pd(b_lo, b_lo, 0x3);
                        c_lo = _mm_add_pd(c_lo, _mm_mul_pd(a_lo, b));
                        c_hi = _mm_add_pd(c_hi, _mm_mul_pd(a_hi, b));
                }

                /* k = 2 */
                {
                        register __m128d a_lo = _mm_castps_pd(PANEL_COL_VEC((*left), 2, 0));
                        register __m128d a_hi = _mm_castps_pd(PANEL_COL_VEC((*left), 2, 1));
                        register __m128d b = _mm_shuffle_pd(b_hi, b_hi, 0x0);
                        c_lo = _mm_add_pd(c_lo, _mm_mul_pd(a_lo, b));
                        c_hi = _mm_add_pd(c_hi, _mm_mul_pd(a_hi, b));
                }

                /* k = 3 */
                {
                        register __m128d a_lo = _mm_castps_pd(PANEL_COL_VEC((*left), 3, 0));
                        register __m128d a_hi = _mm_castps_pd(PANEL_COL_VEC((*left), 3, 1));
                        register __m128d b = _mm_shuffle_pd(b_hi, b_hi, 0x3);
                        c_lo = _mm_add_pd(c_lo, _mm_mul_pd(a_lo, b));
                        c_hi = _mm_add_pd(c_hi, _mm_mul_pd(a_hi, b));
                }

                PANEL_COL_VEC((*ans), col, 0) = _mm_castpd_ps(c_lo);
                PANEL_COL_VEC((*ans), col, 1) = _mm_castpd_ps(c_hi);
        }
}

DECLARE_PNL_PROC_LEFT_RIGHT_ANS(double, avx, , fma)
{
        register __m256d a0 = _mm256_load_pd((double *)&left->vec[0][0]);
        register __m256d a1 = _mm256_load_pd((double *)&left->vec[1][0]);
        register __m256d a2 = _mm256_load_pd((double *)&left->vec[2][0]);
        register __m256d a3 = _mm256_load_pd((double *)&left->vec[3][0]);

        register __m256d c0 = _mm256_load_pd((double *)&ans->vec[0][0]);
        register __m256d c1 = _mm256_load_pd((double *)&ans->vec[1][0]);
        register __m256d c2 = _mm256_load_pd((double *)&ans->vec[2][0]);
        register __m256d c3 = _mm256_load_pd((double *)&ans->vec[3][0]);

        register __m256d bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c0 = _mm256_add_pd(c0, _mm256_mul_pd(a0, _mm256_permute_pd(bcol, 0x0)));
        c0 = _mm256_add_pd(c0, _mm256_mul_pd(a1, _mm256_permute_pd(bcol, 0x3)));
        c0 = _mm256_add_pd(c0, _mm256_mul_pd(a2, _mm256_permute_pd(bcol, 0x0)));
        c0 = _mm256_add_pd(c0, _mm256_mul_pd(a3, _mm256_permute_pd(bcol, 0x3)));
        _mm256_store_pd((double *)&ans->vec[0][0], c0);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c1 = _mm256_add_pd(c1, _mm256_mul_pd(a0, _mm256_permute_pd(bcol, 0x0)));
        c1 = _mm256_add_pd(c1, _mm256_mul_pd(a1, _mm256_permute_pd(bcol, 0x3)));
        c1 = _mm256_add_pd(c1, _mm256_mul_pd(a2, _mm256_permute_pd(bcol, 0x0)));
        c1 = _mm256_add_pd(c1, _mm256_mul_pd(a3, _mm256_permute_pd(bcol, 0x3)));
        _mm256_store_pd((double *)&ans->vec[1][0], c1);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c2 = _mm256_add_pd(c2, _mm256_mul_pd(a0, _mm256_permute_pd(bcol, 0x0)));
        c2 = _mm256_add_pd(c2, _mm256_mul_pd(a1, _mm256_permute_pd(bcol, 0x3)));
        c2 = _mm256_add_pd(c2, _mm256_mul_pd(a2, _mm256_permute_pd(bcol, 0x0)));
        c2 = _mm256_add_pd(c2, _mm256_mul_pd(a3, _mm256_permute_pd(bcol, 0x3)));
        _mm256_store_pd((double *)&ans->vec[2][0], c2);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c3 = _mm256_add_pd(c3, _mm256_mul_pd(a0, _mm256_permute_pd(bcol, 0x0)));
        c3 = _mm256_add_pd(c3, _mm256_mul_pd(a1, _mm256_permute_pd(bcol, 0x3)));
        c3 = _mm256_add_pd(c3, _mm256_mul_pd(a2, _mm256_permute_pd(bcol, 0x0)));
        c3 = _mm256_add_pd(c3, _mm256_mul_pd(a3, _mm256_permute_pd(bcol, 0x3)));
        _mm256_store_pd((double *)&ans->vec[3][0], c3);
}

DECLARE_PNL_PROC_LEFT_RIGHT_ANS(double, avx, _fma, fma)
{
        register __m256d a0 = _mm256_load_pd((double *)&left->vec[0][0]);
        register __m256d a1 = _mm256_load_pd((double *)&left->vec[1][0]);
        register __m256d a2 = _mm256_load_pd((double *)&left->vec[2][0]);
        register __m256d a3 = _mm256_load_pd((double *)&left->vec[3][0]);

        register __m256d c0 = _mm256_load_pd((double *)&ans->vec[0][0]);
        register __m256d c1 = _mm256_load_pd((double *)&ans->vec[1][0]);
        register __m256d c2 = _mm256_load_pd((double *)&ans->vec[2][0]);
        register __m256d c3 = _mm256_load_pd((double *)&ans->vec[3][0]);

        register __m256d bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c0 = _mm256_fmadd_pd(a0, _mm256_permute_pd(bcol, 0x0), c0);
        c0 = _mm256_fmadd_pd(a1, _mm256_permute_pd(bcol, 0x3), c0);
        c0 = _mm256_fmadd_pd(a2, _mm256_permute_pd(bcol, 0x0), c0);
        c0 = _mm256_fmadd_pd(a3, _mm256_permute_pd(bcol, 0x3), c0);
        _mm256_store_pd((double *)&ans->vec[0][0], c0);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c1 = _mm256_fmadd_pd(a0, _mm256_permute_pd(bcol, 0x0), c1);
        c1 = _mm256_fmadd_pd(a1, _mm256_permute_pd(bcol, 0x3), c1);
        c1 = _mm256_fmadd_pd(a2, _mm256_permute_pd(bcol, 0x0), c1);
        c1 = _mm256_fmadd_pd(a3, _mm256_permute_pd(bcol, 0x3), c1);
        _mm256_store_pd((double *)&ans->vec[1][0], c1);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c2 = _mm256_fmadd_pd(a0, _mm256_permute_pd(bcol, 0x0), c2);
        c2 = _mm256_fmadd_pd(a1, _mm256_permute_pd(bcol, 0x3), c2);
        c2 = _mm256_fmadd_pd(a2, _mm256_permute_pd(bcol, 0x0), c2);
        c2 = _mm256_fmadd_pd(a3, _mm256_permute_pd(bcol, 0x3), c2);
        _mm256_store_pd((double *)&ans->vec[2][0], c2);

        bcol = _mm256_load_pd((double *)&right->vec[0][0]);
        c3 = _mm256_fmadd_pd(a0, _mm256_permute_pd(bcol, 0x0), c3);
        c3 = _mm256_fmadd_pd(a1, _mm256_permute_pd(bcol, 0x3), c3);
        c3 = _mm256_fmadd_pd(a2, _mm256_permute_pd(bcol, 0x0), c3);
        c3 = _mm256_fmadd_pd(a3, _mm256_permute_pd(bcol, 0x3), c3);
        _mm256_store_pd((double *)&ans->vec[3][0], c3);
}

// ans = left * right + ans
// right is not transposed
DECLARE_PNL_PROC_LEFT_RIGHT_ANS(double, 512, , fma)
{
        const __m512i ones = _mm512_set1_epi32(1);

        register __m512d a0 = _mm512_load_pd((double *)&left->vec[0][0]);
        register __m512d a1 = _mm512_load_pd((double *)&left->vec[1][0]);
        register __m512d a2 = _mm512_load_pd((double *)&left->vec[2][0]);
        register __m512d a3 = _mm512_load_pd((double *)&left->vec[3][0]);
        register __m512d a4 = _mm512_load_pd((double *)&left->vec[4][0]);
        register __m512d a5 = _mm512_load_pd((double *)&left->vec[5][0]);
        register __m512d a6 = _mm512_load_pd((double *)&left->vec[6][0]);
        register __m512d a7 = _mm512_load_pd((double *)&left->vec[7][0]);

        for (int col = 0; col < 8; ++col) {
                register __m512d bcol = _mm512_load_pd((double *)&right->vec[col][0]);
                register __m512d ccol = _mm512_load_pd((double *)&ans->vec[col][0]);
                register __m512i idx = _mm512_setzero_si512();
                register __m512i increament = _mm512_load_si512(&ones);

                ccol = _mm512_fmadd_pd(a0, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a1, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a2, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a3, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a4, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a5, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a6, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);
                ccol = _mm512_fmadd_pd(a7, _mm512_permutexvar_pd(idx, bcol), ccol);
                idx = _mm512_add_epi64(idx, increament);

                _mm512_store_pd((double *)&ans->vec[col][0], ccol);
        }
}

#endif