Working on L4, now on swiss. FMA_TILE is re arranged.

This commit is contained in:
Youmu Liu
2026-03-03 18:02:31 +08:00
parent ae1a49d6fb
commit 91e136f81f
12 changed files with 1083 additions and 477 deletions
+68 -106
View File
@@ -12,17 +12,12 @@
// float
SN_MK_PROC_DECL_STD(float, sse, , fma)
SN_MK_PROC_DECL_STD(float, sse, , tiled_fma)
{
__mk_param_std();
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, float);
MAT_KERNEL(float)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m128 c0 = _mm_load_ps((float *)&ans_cur->f32_x[0]);
register __m128 c1 = _mm_load_ps((float *)&ans_cur->f32_x[1]);
@@ -30,11 +25,9 @@ SN_MK_PROC_DECL_STD(float, sse, , fma)
register __m128 c3 = _mm_load_ps((float *)&ans_cur->f32_x[3]);
register __m128 bcol;
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(float)
const *right_cur = right + GET_INDEX(commonK, i, k), // GetKernelAddressUnsafe
const *left_cur =
left + GET_INDEX(ansHeight, k, j); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
const __m128 a0 = _mm_load_ps((float *)&left_cur->f32_x[0]);
const __m128 a1 = _mm_load_ps((float *)&left_cur->f32_x[1]);
@@ -76,18 +69,15 @@ SN_MK_PROC_DECL_STD(float, sse, , fma)
_mm_store_ps((float *)&ans_cur->f32_x[3], c3);
}
SN_MK_PROC_DECL_STD(float, avx, , fma)
SN_MK_PROC_DECL_STD(float, avx, , tiled_fma)
{
static const __m256i b01_shift_source = { .m256i_i32 = { 0, 1, 2, 3, 5, 6, 7, 4 } };
static const __m256i b23_shift_source = { .m256i_i32 = { 2, 3, 0, 1, 7, 4, 5, 6 } };
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, float);
MAT_KERNEL(float)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m256 c01 = _mm256_load_ps((float *)&ans->f32_y[0]);
register __m256 c23 = _mm256_load_ps((float *)&ans->f32_y[1]);
@@ -95,11 +85,9 @@ SN_MK_PROC_DECL_STD(float, avx, , fma)
register __m256i b01_shift = _mm256_load_si256(&b01_shift_source);
register __m256i b23_shift = _mm256_load_si256(&b23_shift_source);
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(float)
const *left_cur = left + GET_INDEX(ansHeight, i, k), // GetKernelAddressUnsafe
const *right_cur =
right + GET_INDEX(ansHeight, j, k); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
register __m256 a01 = _mm256_load_ps((float *)&left_cur->f32_y[0]);
register __m256 a23 = _mm256_load_ps((float *)&left_cur->f32_y[1]);
@@ -139,18 +127,15 @@ SN_MK_PROC_DECL_STD(float, avx, , fma)
_mm256_store_ps((float *)&ans_cur->f32_y[1], c23);
}
SN_MK_PROC_DECL_STD(float, avx, _fma, fma)
SN_MK_PROC_DECL_STD(float, avx, _fma, tiled_fma)
{
static const __m256i b01_shift_source = { .m256i_i32 = { 0, 1, 2, 3, 5, 6, 7, 4 } };
static const __m256i b23_shift_source = { .m256i_i32 = { 2, 3, 0, 1, 7, 4, 5, 6 } };
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, float);
MAT_KERNEL(float)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m256 c01 = _mm256_load_ps((float *)&ans_cur->f32_y[0]);
register __m256 c23 = _mm256_load_ps((float *)&ans_cur->f32_y[1]);
@@ -158,11 +143,9 @@ SN_MK_PROC_DECL_STD(float, avx, _fma, fma)
register __m256i b01_shift = _mm256_load_si256(&b01_shift_source);
register __m256i b23_shift = _mm256_load_si256(&b23_shift_source);
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(float)
const *left_cur = left + GET_INDEX(ansHeight, k, j), // GetKernelAddressUnsafe
const *right_cur =
right + GET_INDEX(commonK, i, k); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
register __m256 a01 = _mm256_load_ps((float *)&left_cur->f32_y[0]);
register __m256 a23 = _mm256_load_ps((float *)&left_cur->f32_y[1]);
@@ -202,7 +185,7 @@ SN_MK_PROC_DECL_STD(float, avx, _fma, fma)
_mm256_store_ps((float *)&ans_cur->f32_y[1], c23);
}
SN_MK_PROC_DECL_STD(float, 512, , fma)
SN_MK_PROC_DECL_STD(float, 512, , tiled_fma)
{
static const __m512i b_shift_source = { .m512i_i32 = {
// shuff order
@@ -211,18 +194,13 @@ SN_MK_PROC_DECL_STD(float, 512, , fma)
10, 11, 8, 9, //col2
15, 12, 13, 14 //col3
} };
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
register __m512i b_shift = _mm512_load_si512(&b_shift_source);
register __m512 c = _mm512_load_ps((float *)(ans + GET_INDEX(ansHeight, i, j)));
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(double) *const left_cur = left + GET_INDEX(ansHeight, k, j),
const *right_cur =
right + GET_INDEX(commonK, i,
k); // GetKernelAddressUnsafe
__mk_param_std(TILE, LEFT_RIGHT_ANS, float);
SN_MAT_KERNEL_LOCAL const *ans_cur = ans + GET_INDEX(common_y, x, y);
register __m512 c = _mm512_load_ps(ans_cur);
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
register __m512 a = _mm512_load_ps((float *)left_cur);
register __m512 b = _mm512_load_ps((float *)right_cur);
@@ -245,7 +223,7 @@ SN_MK_PROC_DECL_STD(float, 512, , fma)
c = _mm512_fmadd_ps(a_perm, b_perm, c);
}
_mm512_store_ps((float *)(ans + GET_INDEX(ansHeight, i, j)), c);
_mm512_store_ps((float *)ans_cur, c);
}
// double
@@ -282,26 +260,25 @@ SN_MK_PROC_DECL_STD(float, 512, , fma)
c##i##_hi = _mm_add_pd(c##i##_hi, _mm_mul_pd(a_hi, b)); \
} while (0)
SN_MK_PROC_DECL_STD(double, sse, , fma)
SN_MK_PROC_DECL_STD(double, sse, , tiled_fma)
{
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, double);
register __m128d c0_lo = _mm_load_pd(ans->f64_x + (0 * 2 + 0));
register __m128d c0_hi = _mm_load_pd(ans->f64_x + (0 * 2 + 1));
register __m128d c1_lo = _mm_load_pd(ans->f64_x + (1 * 2 + 0));
register __m128d c1_hi = _mm_load_pd(ans->f64_x + (1 * 2 + 1));
register __m128d c2_lo = _mm_load_pd(ans->f64_x + (2 * 2 + 0));
register __m128d c2_hi = _mm_load_pd(ans->f64_x + (2 * 2 + 1));
register __m128d c3_lo = _mm_load_pd(ans->f64_x + (3 * 2 + 0));
register __m128d c3_hi = _mm_load_pd(ans->f64_x + (3 * 2 + 1));
MAT_KERNEL(double)
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(double)
const *left_cur = left + GET_INDEX(ansHeight, i, k),
const *right_cur = right + GET_INDEX(commonK, j, k); // GetKernelAddressUnsafe
register __m128d c0_lo = _mm_load_pd(ans_cur->f64_x + (0 * 2 + 0));
register __m128d c0_hi = _mm_load_pd(ans_cur->f64_x + (0 * 2 + 1));
register __m128d c1_lo = _mm_load_pd(ans_cur->f64_x + (1 * 2 + 0));
register __m128d c1_hi = _mm_load_pd(ans_cur->f64_x + (1 * 2 + 1));
register __m128d c2_lo = _mm_load_pd(ans_cur->f64_x + (2 * 2 + 0));
register __m128d c2_hi = _mm_load_pd(ans_cur->f64_x + (2 * 2 + 1));
register __m128d c3_lo = _mm_load_pd(ans_cur->f64_x + (3 * 2 + 0));
register __m128d c3_hi = _mm_load_pd(ans_cur->f64_x + (3 * 2 + 1));
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
DOUBLE_KERNEL_FMA_CYCLE(0);
@@ -312,25 +289,22 @@ SN_MK_PROC_DECL_STD(double, sse, , fma)
DOUBLE_KERNEL_FMA_CYCLE(3);
}
_mm_store_pd(ans->f64_x + (0 * 2 + 0), c0_lo);
_mm_store_pd(ans->f64_x + (0 * 2 + 1), c0_hi);
_mm_store_pd(ans->f64_x + (1 * 2 + 0), c1_lo);
_mm_store_pd(ans->f64_x + (1 * 2 + 1), c1_hi);
_mm_store_pd(ans->f64_x + (2 * 2 + 0), c2_lo);
_mm_store_pd(ans->f64_x + (2 * 2 + 1), c2_hi);
_mm_store_pd(ans->f64_x + (3 * 2 + 0), c3_lo);
_mm_store_pd(ans->f64_x + (3 * 2 + 1), c3_hi);
_mm_store_pd(ans_cur->f64_x + (0 * 2 + 0), c0_lo);
_mm_store_pd(ans_cur->f64_x + (0 * 2 + 1), c0_hi);
_mm_store_pd(ans_cur->f64_x + (1 * 2 + 0), c1_lo);
_mm_store_pd(ans_cur->f64_x + (1 * 2 + 1), c1_hi);
_mm_store_pd(ans_cur->f64_x + (2 * 2 + 0), c2_lo);
_mm_store_pd(ans_cur->f64_x + (2 * 2 + 1), c2_hi);
_mm_store_pd(ans_cur->f64_x + (3 * 2 + 0), c3_lo);
_mm_store_pd(ans_cur->f64_x + (3 * 2 + 1), c3_hi);
}
SN_MK_PROC_DECL_STD(double, avx, , fma)
SN_MK_PROC_DECL_STD(double, avx, , tiled_fma)
{
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, double);
MAT_KERNEL(double)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m256d c0 = _mm256_load_pd((double *)&ans_cur->f64_y[0]);
register __m256d c1 = _mm256_load_pd((double *)&ans_cur->f64_y[1]);
@@ -338,11 +312,9 @@ SN_MK_PROC_DECL_STD(double, avx, , fma)
register __m256d c3 = _mm256_load_pd((double *)&ans_cur->f64_y[3]);
register __m256d bcol;
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(double)
const *right_cur = right + GET_INDEX(commonK, i, k), // GetKernelAddressUnsafe
const *left_cur =
left + GET_INDEX(ansHeight, k, j); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
const __m256d a0 = _mm256_load_pd((double *)&left_cur->f64_y[0]);
const __m256d a1 = _mm256_load_pd((double *)&left_cur->f64_y[1]);
@@ -400,15 +372,12 @@ SN_MK_PROC_DECL_STD(double, avx, , fma)
_mm256_store_pd((double *)&ans_cur->f64_y[3], c3);
}
SN_MK_PROC_DECL_STD(double, avx, _fma, fma)
SN_MK_PROC_DECL_STD(double, avx, _fma, tiled_fma)
{
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, double);
MAT_KERNEL(double)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m256d c0 = _mm256_load_pd((double *)&ans_cur->f64_y[0]);
register __m256d c1 = _mm256_load_pd((double *)&ans_cur->f64_y[1]);
@@ -416,11 +385,9 @@ SN_MK_PROC_DECL_STD(double, avx, _fma, fma)
register __m256d c3 = _mm256_load_pd((double *)&ans_cur->f64_y[3]);
register __m256d bcol;
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(double)
const *right_cur = right + GET_INDEX(commonK, i, k), // GetKernelAddressUnsafe
const *left_cur =
left + GET_INDEX(ansHeight, k, j); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
const __m256d a0 = _mm256_load_pd((double *)&left_cur->f64_y[0]);
const __m256d a1 = _mm256_load_pd((double *)&left_cur->f64_y[1]);
@@ -462,18 +429,15 @@ SN_MK_PROC_DECL_STD(double, avx, _fma, fma)
_mm256_store_pd((double *)&ans_cur->f64_y[3], c3);
}
SN_MK_PROC_DECL_STD(double, 512, , fma)
SN_MK_PROC_DECL_STD(double, 512, , tiled_fma)
{
static const __m512i b01_shift_source = { .m512i_i64 = { 0, 1, 2, 3, 5, 6, 7, 4 } };
static const __m512i b23_shift_source = { .m512i_i64 = { 2, 3, 0, 1, 7, 4, 5, 6 } };
int i = enumeration->m128_i32[0];
int j = enumeration->m128_i32[1];
int commonK = enumeration->m128_i32[2];
int ansHeight = enumeration->m128_i32[3];
__mk_param_std(TILE, LEFT_RIGHT_ANS, double);
MAT_KERNEL(double)
const *ans_cur = ans + GET_INDEX(ansHeight, i, j); // GetKernelAddressUnsafe
const *ans_cur = ans + GET_INDEX(common_y, x, y); // GetKernelAddressUnsafe
register __m512d c01 = _mm512_load_pd((double *)&ans->f64_z[0]);
register __m512d c23 = _mm512_load_pd((double *)&ans->f64_z[1]);
@@ -481,11 +445,9 @@ SN_MK_PROC_DECL_STD(double, 512, , fma)
register __m512i b01_shift = _mm512_load_si512(&b01_shift_source);
register __m512i b23_shift = _mm512_load_si512(&b23_shift_source);
for (size_t k = 0; k < commonK; k++) {
MAT_KERNEL(double)
const *left_cur = left + GET_INDEX(ansHeight, i, k), // GetKernelAddressUnsafe
const *right_cur =
right + GET_INDEX(ansHeight, j, k); // GetKernelAddressUnsafe
for (size_t z = 0; z < common_z; z++) {
SN_MAT_KERNEL_LOCAL const *right_cur = right + GET_INDEX(common_z, x, z),
const *left_cur = left + GET_INDEX(common_y, z, y);
register __m512d a01 = _mm512_load_pd((double *)&left_cur->f64_z[0]);
register __m512d a23 = _mm512_load_pd((double *)&left_cur->f64_z[1]);