tile fma added

This commit is contained in:
Youmu Liu
2026-01-09 18:13:46 +08:00
parent 4ff1d856a7
commit 6827259891
12 changed files with 736 additions and 265 deletions
+68 -192
View File
@@ -180,6 +180,50 @@ DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, 512, , fma)
// double
DECLARE_KER_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_load_pd(ans->f64_x + (col * 2 + 0));
register __m128d c_hi = _mm_load_pd(ans->f64_x + (col * 2 + 1));
/* right matrix column split into two chunks */
register __m128d b_lo = _mm_load_pd(right + (2 * col + 0));
register __m128d b_hi = _mm_load_pd(right + (2 * col + 1));
/* k = 0 */
register __m128d a_lo = _mm_load_pd(left + (0 * 2 + 0));
register __m128d a_hi = _mm_load_pd(left + (0 * 2 + 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 */
a_lo = _mm_load_pd(left + (1 * 2 + 0));
a_hi = _mm_load_pd(left + (1 * 2 + 1));
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 */
a_lo = _mm_load_pd(left + (1 * 2 + 0));
a_hi = _mm_load_pd(left + (1 * 2 + 1));
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 */
a_lo = _mm_load_pd(left + (1 * 2 + 0));
a_hi = _mm_load_pd(left + (1 * 2 + 1));
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));
ans->f64_x[col * 2 + 0] = c_lo;
ans->f64_x[col * 2 + 1] = c_hi;
}
}
DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, avx, , fma)
{
const __m256d a0 = _mm256_load_pd((double *)&left->f64_y[0]);
@@ -201,27 +245,27 @@ DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, avx, , fma)
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b0, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b0, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
/* Column 0 */
/* Column 1 */
ccol = _mm256_load_pd((double *)&ans->f32_y[1]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b1, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b1, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b1, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b1, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
/* Column 0 */
ans->f64_y[1] = ccol;
/* Column 2 */
ccol = _mm256_load_pd((double *)&ans->f32_y[2]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b2, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b2, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b2, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b2, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
/* Column 0 */
ans->f64_y[2] = ccol;
/* Column 3 */
ccol = _mm256_load_pd((double *)&ans->f32_y[3]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b3, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b3, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b3, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b3, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
ans->f64_y[3] = ccol;
}
DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, avx, _fma, fma)
@@ -245,27 +289,27 @@ DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, avx, _fma, fma)
ccol = _mm256_fmadd_pd(a2, _mm256_permute4x64_pd(b0, _FMA_SELECT(2)), ccol);
ccol = _mm256_fmadd_pd(a3, _mm256_permute4x64_pd(b0, _FMA_SELECT(3)), ccol);
ans->f64_y[0] = ccol;
/* Column 0 */
/* Column 1 */
ccol = _mm256_load_pd((double *)&ans->f32_y[1]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b1, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b1, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b1, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b1, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
/* Column 0 */
ccol = _mm256_fmadd_pd(a0, _mm256_permute4x64_pd(b1, _FMA_SELECT(0)), ccol);
ccol = _mm256_fmadd_pd(a1, _mm256_permute4x64_pd(b1, _FMA_SELECT(1)), ccol);
ccol = _mm256_fmadd_pd(a2, _mm256_permute4x64_pd(b1, _FMA_SELECT(2)), ccol);
ccol = _mm256_fmadd_pd(a3, _mm256_permute4x64_pd(b1, _FMA_SELECT(3)), ccol);
ans->f64_y[1] = ccol;
/* Column 2 */
ccol = _mm256_load_pd((double *)&ans->f32_y[2]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b2, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b2, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b2, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b2, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
/* Column 0 */
ccol = _mm256_fmadd_pd(a0, _mm256_permute4x64_pd(b2, _FMA_SELECT(0)), ccol);
ccol = _mm256_fmadd_pd(a1, _mm256_permute4x64_pd(b2, _FMA_SELECT(1)), ccol);
ccol = _mm256_fmadd_pd(a2, _mm256_permute4x64_pd(b2, _FMA_SELECT(2)), ccol);
ccol = _mm256_fmadd_pd(a3, _mm256_permute4x64_pd(b2, _FMA_SELECT(3)), ccol);
ans->f64_y[2] = ccol;
/* Column 3 */
ccol = _mm256_load_pd((double *)&ans->f32_y[3]);
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a0, _mm256_permute4x64_pd(b3, _FMA_SELECT(0))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a1, _mm256_permute4x64_pd(b3, _FMA_SELECT(1))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a2, _mm256_permute4x64_pd(b3, _FMA_SELECT(2))));
ccol = _mm256_add_pd(ccol, _mm256_mul_pd(a3, _mm256_permute4x64_pd(b3, _FMA_SELECT(3))));
ans->f64_y[0] = ccol;
ccol = _mm256_fmadd_pd(a0, _mm256_permute4x64_pd(b3, _FMA_SELECT(0)), ccol);
ccol = _mm256_fmadd_pd(a1, _mm256_permute4x64_pd(b3, _FMA_SELECT(1)), ccol);
ccol = _mm256_fmadd_pd(a2, _mm256_permute4x64_pd(b3, _FMA_SELECT(2)), ccol);
ccol = _mm256_fmadd_pd(a3, _mm256_permute4x64_pd(b3, _FMA_SELECT(3)), ccol);
ans->f64_y[3] = ccol;
}
DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, 512, , fma)
@@ -314,172 +358,4 @@ DECLARE_KER_PROC_LEFT_RIGHT_ANS(double, 512, , fma)
ans->f64_z[1] = c23;
}
DECLARE_KER_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_load_pd(ans->f64_x + (col * 2 + 0));
register __m128d c_hi = _mm_load_pd(ans->f64_x + (col * 2 + 1));
/* right matrix column split into two chunks */
register __m128d b_lo = _mm_load_pd(right + (2 * col + 0));
register __m128d b_hi = _mm_load_pd(right + (2 * col + 1));
/* k = 0 */
register __m128d a_lo = _mm_load_pd(left + (0 * 2, 0));
register __m128d a_hi = _mm_load_pd(left + (0 * 2, 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 */
a_lo = _mm_load_pd(left + (1 * 2, 0));
a_hi = _mm_load_pd(left + (1 * 2, 1));
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 */
a_lo = _mm_load_pd(left + (1 * 2, 0));
a_hi = _mm_load_pd(left + (1 * 2, 1));
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 */
a_lo = _mm_load_pd(left + (1 * 2, 0));
a_hi = _mm_load_pd(left + (1 * 2, 1));
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));
ans->f64_x[col * 2 + 0] = c_lo;
ans->f64_x[col * 2 + 1] = 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