youmuliu/stgsharp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

make a trial on packing GEMM kernel

Browse Source
This commit is contained in:
Youmu Liu
2026-06-11 20:03:02 +08:00
parent 6b449a341e
commit 06a119cc6f
4 changed files with 105 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
StgSharp.Native/sn-src/mk.fma.x86.c
+71 -12
View File
@@ -11,29 +11,88 @@
#define _FMA_SELECT(i) _MM_SHUFFLE(i, i, i, i)
#define GET_INDEX(col_len, x, y) (col_len * x + y)
// float
SN_MK_PROC_DECL_STD(float, sse, , fma_pack)
/*
* pack left matrix data to L4 buffer
* scalar pack carries column length and offset data:
* [column length, buffer requires, empty, empty]
*/
SN_MK_PROC_DECL_STD(float, sse, , fma_pack_left)
{
__mk_param_std(RIGHT_ANS, float);
memcpy(ans, right+offset_1, sizeof(MAT_KERNEL(float)) * count_1);
__mk_param_std(RIGHT_ANS_SCALAR, float);
// for SSE branch, no need for rounding, a direct copy is sufficient as long as the buffer is large enough to hold the packed data
memcpy(ans, right + offset_1, sizeof(MAT_KERNEL(float)) * count_2);
}
SN_MK_PROC_DECL_STD(float, avx, , fma_pack)
SN_MK_PROC_DECL_STD(float, avx, , fma_pack_left)
{
__mk_param_std(RIGHT_ANS, float);
__mk_param_std(RIGHT_ANS_SCALAR, float);
size_t i=0;
for(; i <= count_1; i+=2) {
register __m256 ker_0_left = _mm256_load_ps((float *)&right[offset_1 + i].y[0]);
register __m256 ker_0_right = _mm256_load_ps((float *)&right[offset_1 + i].y[1]);
register __m256 ker_1_left = _mm256_load_ps((float *)&right[offset_1 + i].y[0]);
register __m256 ker_1_right = _mm256_load_ps((float *)&right[offset_1 + i].y[1]);
register __m256 col_0 = _mm256_unpacklo_ps(ker_0_left, ker_1_left);
register __m256 col_1 = _mm256_unpackhi_ps(ker_0_left, ker_1_left);
register __m256 col_2 = _mm256_unpacklo_ps(ker_0_right, ker_1_right);
register __m256 col_3 = _mm256_unpackhi_ps(ker_0_right, ker_1_right);
//memcpy(ans, right, sizeof(MAT_KERNEL(float)) * common_y * common_z);
_mm256_store_ps((float *)&ans[i].y[0], col_0);
_mm256_store_ps((float *)&ans[i].y[1], col_1);
_mm256_store_ps((float *)&ans[i].y[2], col_2);
_mm256_store_ps((float *)&ans[i].y[3], col_3);
}
if(count_1 == i) {
register __m256 ker_0_left = _mm256_load_ps((float *)&right[offset_1 + i].y[0]);
register __m256 ker_0_right = _mm256_load_ps((float *)&right[offset_1 + i].y[1]);
register __m256 zero = _mm256_setzero_ps();
register __m256 col_0 = _mm256_unpacklo_ps(ker_0_left, zero);
register __m256 col_1 = _mm256_unpackhi_ps(ker_0_left, zero);
register __m256 col_2 = _mm256_unpacklo_ps(ker_0_right, zero);
register __m256 col_3 = _mm256_unpackhi_ps(ker_0_right, zero);
_mm256_store_ps((float *)&ans[i].y[0], col_0);
_mm256_store_ps((float *)&ans[i].y[1], col_1);
_mm256_store_ps((float *)&ans[i].y[2], col_2);
_mm256_store_ps((float *)&ans[i].y[3], col_3);
}
}
SN_MK_PROC_DECL_STD(float, 512, , fma_pack)
SN_MK_PROC_DECL_STD(float, 512, , fma_pack_left)
{
__mk_param_std(RIGHT_ANS, float);
__mk_param_std(RIGHT_ANS_SCALAR, float);
size_t i=0;
for(; i <= count_1; i+=2) {
register __m512 ker_0_left = _mm512_load_ps((float *)&right[offset_1 + i].z[0]);
register __m512 ker_0_right = _mm512_load_ps((float *)&right[offset_1 + i].z[1]);
register __m512 ker_1_left = _mm512_load_ps((float *)&right[offset_1 + i].z[0]);
register __m512 ker_1_right = _mm512_load_ps((float *)&right[offset_1 + i].z[1]);
register __m512 col_0 = _mm256_unpacklo_ps(ker_0_left, ker_1_left);
register __m512 col_1 = _mm256_unpackhi_ps(ker_0_left, ker_1_left);
register __m512 col_2 = _mm256_unpacklo_ps(ker_0_right, ker_1_right);
register __m512 col_3 = _mm256_unpackhi_ps(ker_0_right, ker_1_right);
//memcpy(ans, right, sizeof(MAT_KERNEL(float)) * common_y * common_z);
_mm256_store_ps((float *)&ans[i].y[0], col_0);
_mm256_store_ps((float *)&ans[i].y[1], col_1);
_mm256_store_ps((float *)&ans[i].y[2], col_2);
_mm256_store_ps((float *)&ans[i].y[3], col_3);
}
if(count_1 == i) {
register __m256 ker_0_left = _mm256_load_ps((float *)&right[offset_1 + i].y[0]);
register __m256 ker_0_right = _mm256_load_ps((float *)&right[offset_1 + i].y[1]);
register __m256 zero = _mm256_setzero_ps();
register __m256 col_0 = _mm256_unpacklo_ps(ker_0_left, zero);
register __m256 col_1 = _mm256_unpackhi_ps(ker_0_left, zero);
register __m256 col_2 = _mm256_unpacklo_ps(ker_0_right, zero);
register __m256 col_3 = _mm256_unpackhi_ps(ker_0_right, zero);
_mm256_store_ps((float *)&ans[i].y[0], col_0);
_mm256_store_ps((float *)&ans[i].y[1], col_1);
_mm256_store_ps((float *)&ans[i].y[2], col_2);
_mm256_store_ps((float *)&ans[i].y[3], col_3);
}
}
SN_MK_PROC_DECL_STD(float, sse, , fma)