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

Gemm/Fma kernel on avx512 use new solution, with less permute and shuff instruction

Browse Source
This commit is contained in:
Youmu Liu
2026-06-04 16:43:17 +08:00
parent 6c38bcfade
commit 2690cf9341
3 changed files with 44 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
StgSharp.Native/sn-head/sn_intrinsic.context.x86.h
+4
View File
@@ -1,3 +1,6 @@
#ifndef SN_INTRINSIC_CONTEXT_X86_H
#define SN_INTRINSIC_CONTEXT_X86_H
#include "sn_target.h" #include "sn_target.h"
#if SN_IS_ARCH(SN_ARCH_X86_64) #if SN_IS_ARCH(SN_ARCH_X86_64)
@@ -271,3 +274,4 @@ INTERNAL int SN_DECL index_pair_sse(short const *stram, uint32_t target, int len
typedef enum intrinsic_level { SIMD0 = 0 } intrinsic_level; typedef enum intrinsic_level { SIMD0 = 0 } intrinsic_level;
#endif #endif
#endif
StgSharp.Native/sn-src/mk.fma.x86.c
+23 -27
View File
@@ -142,39 +142,35 @@ SN_MK_PROC_DECL_STD(float, 512, , fma)
{ {
__mk_param_std(TILE, LEFT_RIGHT_ANS, float); __mk_param_std(TILE, LEFT_RIGHT_ANS, float);
SN_MAT_KERNEL_LOCAL *restrict ans_cur = ans + GET_INDEX(common_y, x, y); SN_MAT_KERNEL_LOCAL *restrict ans_cur = ans + GET_INDEX(common_y, x, y);
register __m512 c_0 = _mm512_load_ps((float *)ans_cur); register __m512 c0 = _mm512_setzero_ps();
register __m512 c_1 = c_0; register __m512 c1 = _mm512_setzero_ps();
register __m512 c_2 = c_0; register __m512 c2 = _mm512_setzero_ps();
register __m512 c_3 = c_0; register __m512 c3 = _mm512_setzero_ps();
for (size_t z = 0; z < common_z; z+=4) { for (size_t z = 0; z < common_z; z+=4) {
SN_MAT_KERNEL_LOCAL const *right_ = right + z; const SN_MAT_KERNEL_LOCAL *a_base = left + z;
SN_MAT_KERNEL_LOCAL const *left_0 = left + z + 0; const SN_MAT_KERNEL_LOCAL *b_base = right + z;
SN_MAT_KERNEL_LOCAL const *left_1 = left + z + 1;
SN_MAT_KERNEL_LOCAL const *left_2 = left + z + 2;
SN_MAT_KERNEL_LOCAL const *left_3 = left + z + 3;
// full right load：[kb_col0 | kb_col1 | kb_col2 | kb_col3] register __m512 b0 = _mm512_load_ps((float *)&b_base[0].z[0]);
register __m512 b_0 = _mm512_load_ps((float *)(right+0)); register __m512 b1 = _mm512_load_ps((float *)&b_base[1].z[0]);
register __m512 b_1 = _mm512_load_ps((float *)(right+1)); register __m512 b2 = _mm512_load_ps((float *)&b_base[2].z[0]);
register __m512 b_2 = _mm512_load_ps((float *)(right+2)); register __m512 b3 = _mm512_load_ps((float *)&b_base[3].z[0]);
register __m512 b_3 = _mm512_load_ps((float *)(right+3));
// k=0: broadcast ka_col0 到4个lane，permute_ps 取 kb(0,j) register __m512 a0 = _mm512_load_ps((float *)&a_base[0].z[0]);
register __m512 ak; register __m512 a1 = _mm512_load_ps((float *)&a_base[1].z[0]);
register __m512 ak_1; register __m512 a2 = _mm512_load_ps((float *)&a_base[2].z[0]);
register __m512 ak_2; register __m512 a3 = _mm512_load_ps((float *)&a_base[3].z[0]);
register __m512 ak_3;
SN_MK_FMA_F32_512_CYCLE(0);
SN_MK_FMA_F32_512_CYCLE(1);
SN_MK_FMA_F32_512_CYCLE(2);
SN_MK_FMA_F32_512_CYCLE(3);
SN_MK_FMA_F32_512_PHASE(0);
SN_MK_FMA_F32_512_PHASE(1);
SN_MK_FMA_F32_512_PHASE(2);
SN_MK_FMA_F32_512_PHASE(3);
} }
c_0 = _mm512_add_ps(c_0, c_1); SN_FMA_PACK_COLUMN(c0, (float*)&(ans_cur->x[0]));
c_2= _mm512_add_ps(c_2, c_3); SN_FMA_PACK_COLUMN(c1, (float*)&(ans_cur->x[1]));
c_0 = _mm512_add_ps(c_0, c_2); SN_FMA_PACK_COLUMN(c2, (float*)&(ans_cur->x[2]));
_mm512_store_ps((float *)ans_cur, c_0); SN_FMA_PACK_COLUMN(c3, (float*)&(ans_cur->x[3]));
} }
StgSharp.Native/sn-template/sn_template.fma.x86.h
+16 -18
View File
@@ -2,25 +2,23 @@
#if SN_IS_ARCH(SN_ARCH_X86_64) #if SN_IS_ARCH(SN_ARCH_X86_64)
/* #define SN_MK_FMA_F32_512_CYCLE(i) \
* requested ZMMs: do { \
* c0 = _mm512_fmadd_ps(c0, a0, _mm512_permute_ps(b0, _FMA_SELECT(i))); \
*/ c1 = _mm512_fmadd_ps(c1, a1, _mm512_permute_ps(b1, _FMA_SELECT(i))); \
#define SN_MK_FMA_F32_512_CYCLE(i)\ c2 = _mm512_fmadd_ps(c2, a2, _mm512_permute_ps(b2, _FMA_SELECT(i))); \
do{\ c3 = _mm512_fmadd_ps(c3, a3, _mm512_permute_ps(b3, _FMA_SELECT(i))); \
\ } while (0)
}while(0)
#define SN_MK_FMA_F32_512_PHASE(i) \ #define SN_FMA_PACK_COLUMN(col, target) \
do { \ do { \
ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ register __m256 low = _mm512_castps512_ps256(col); \
c_0 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ register __m256 high = _mm512_extractf32x8_ps(col, 1); \
ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ register __m256 half = _mm256_add_ps(high, low); \
c_1 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ register __m128 high128 = _mm256_extractf32x4_ps(half, 1); \
ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ register __m128 low128 = _mm256_castps256_ps128(half); \
c_2 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ register __m128 half128 = _mm_add_ps(high128, low128); \
ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ _mm_store_ps(target, half128); \
c_3 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \
} while (0) } while (0)
// double // double