From 06a119cc6f1d455ae6a566ae9fb70fc0fa727c55 Mon Sep 17 00:00:00 2001 From: Youmu Liu <105042599+Nitload-NSI@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:03:02 +0800 Subject: [PATCH] make a trial on packing GEMM kernel --- .../Numeric/Runtime/MatrixParallelTask.cs | 8 +- .../sn-head/sn_intrinsic.context.x86.h | 24 ++++-- StgSharp.Native/sn-head/sn_intrinsic.h | 21 +++-- StgSharp.Native/sn-src/mk.fma.x86.c | 83 ++++++++++++++++--- 4 files changed, 105 insertions(+), 31 deletions(-) diff --git a/StgSharp.Mathematics/Numeric/Runtime/MatrixParallelTask.cs b/StgSharp.Mathematics/Numeric/Runtime/MatrixParallelTask.cs index 3b715ba7..bab82edf 100644 --- a/StgSharp.Mathematics/Numeric/Runtime/MatrixParallelTask.cs +++ b/StgSharp.Mathematics/Numeric/Runtime/MatrixParallelTask.cs @@ -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() { diff --git a/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h b/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h index e0e1b8fb..7a8b6baa 100644 --- a/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h +++ b/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h @@ -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); diff --git a/StgSharp.Native/sn-head/sn_intrinsic.h b/StgSharp.Native/sn-head/sn_intrinsic.h index 1924617b..5abcd85a 100644 --- a/StgSharp.Native/sn-head/sn_intrinsic.h +++ b/StgSharp.Native/sn-head/sn_intrinsic.h @@ -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; diff --git a/StgSharp.Native/sn-src/mk.fma.x86.c b/StgSharp.Native/sn-src/mk.fma.x86.c index e4924db2..f1e1c1a8 100644 --- a/StgSharp.Native/sn-src/mk.fma.x86.c +++ b/StgSharp.Native/sn-src/mk.fma.x86.c @@ -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)