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.Mathematics/Numeric/Runtime/MatrixParallelTask.cs
+4 -4
View File
@@ -44,11 +44,11 @@ namespace StgSharp.Mathematics.Numeric.Runtime
}
[StructLayout(LayoutKind.Explicit)]
[StructLayout(LayoutKind.Explicit, Size = 128)]
internal unsafe struct MatrixParallelTask
{
[FieldOffset(104)] private ulong reversed;
[FieldOffset(104)] private ulong _reserved0;
[FieldOffset(0)] public L4.Handle Buffer0;
[FieldOffset(16)] public L4.Handle Buffer1;
@@ -56,8 +56,8 @@ namespace StgSharp.Mathematics.Numeric.Runtime
[FieldOffset(48)] public L4.Handle Buffer3;
[FieldOffset(112)] public M128 ScalarPack;
[FieldOffset(64)] public M512 Profile;
[FieldOffset(64)] public MatrixElementType ElementType;
[FieldOffset(68)] public MatrixIntrinsicHandle ComputeHandle;
[FieldOffset(64)] public MatrixElementType ElementType; //记录了类型掩码的类型位宽的union，大小4byte
[FieldOffset(68)] public MatrixIntrinsicHandle ComputeHandle; //记录了函数指针在类型组中的索引，实际就是int，但是为每一个索引起了名字，于是变成了enum，大小4byte
public MatrixParallelTask()
{
StgSharp.Native/sn-head/sn_intrinsic.context.x86.h
+18 -6
View File
@@ -170,9 +170,21 @@ SN_MK_PROC_DECL_STD(double, 512, , fill);
SN_MK_PROC_DECL_STD(float, sse, , quality);
SN_MK_PROC_DECL_STD(float, sse, , fma_pack);
SN_MK_PROC_DECL_STD(float, avx, , fma_pack);
SN_MK_PROC_DECL_STD(float, 512, , fma_pack);
/*
* pack left matrix data to L4 buffer
*
* CONVENSIONS:
*
* buffer_1 is pointer to original colum, offset_1 is the offset to the original column, count_1 is the count can be packed
* C# shceduler is responsible for calculating the offset and count based on the actual buffer size and matrix dimension,
* and make sure that count_1 will not exceeds lenght of the column
*
* scalar pack carries column length and offset data:
* [column length, offset, buffer requires, empty]
*/
SN_MK_PROC_DECL_STD(float, sse, , fma_pack_left);
SN_MK_PROC_DECL_STD(float, avx, , fma_pack_left);
SN_MK_PROC_DECL_STD(float, 512, , fma_pack_left);
SN_MK_PROC_DECL_STD(float, sse, , fma);
SN_MK_PROC_DECL_STD(float, avx, , fma);
@@ -183,9 +195,9 @@ SN_MK_PROC_DECL_STD(float, sse, , fma_unpack);
SN_MK_PROC_DECL_STD(float, avx, , fma_unpack);
SN_MK_PROC_DECL_STD(float, 512, , fma_unpack);
SN_MK_PROC_DECL_STD(double, sse, , fma_pack);
SN_MK_PROC_DECL_STD(double, avx, , fma_pack);
SN_MK_PROC_DECL_STD(double, 512, , fma_pack);
SN_MK_PROC_DECL_STD(double, sse, , fma_pack_left);
SN_MK_PROC_DECL_STD(double, avx, , fma_pack_left);
SN_MK_PROC_DECL_STD(double, 512, , fma_pack_left);
SN_MK_PROC_DECL_STD(double, sse, , fma);
SN_MK_PROC_DECL_STD(double, avx, , fma);
StgSharp.Native/sn-head/sn_intrinsic.h
+12 -9
View File
@@ -117,15 +117,18 @@ typedef struct matrix_task {
sn_L4_handle buffer_3;
union {
struct {
int offset_0;
int count_0;
int offset_1;
int count_1;
int offset_2;
int count_2;
int offset_3;
int count_3;
uint64_t reverse;
int32_t reverse_0;
int32_t reverse_1;
int32_t offset_0;
int32_t count_0;
int32_t offset_1;
int32_t count_1;
int32_t offset_2;
int32_t count_2;
int32_t offset_3;
int32_t count_3;
int32_t reverse_2;
int32_t reverse_3;
M128 scalar_pack;
};
M512 profile;
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)