From 2690cf934113ae23af7370266a29e81e6beb5e63 Mon Sep 17 00:00:00 2001 From: Youmu Liu <105042599+Nitload-NSI@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:43:17 +0800 Subject: [PATCH] Gemm/Fma kernel on avx512 use new solution, with less permute and shuff instruction --- .../sn-head/sn_intrinsic.context.x86.h | 4 ++ StgSharp.Native/sn-src/mk.fma.x86.c | 52 +++++++++---------- .../sn-template/sn_template.fma.x86.h | 34 ++++++------ 3 files changed, 44 insertions(+), 46 deletions(-) diff --git a/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h b/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h index 9c892e68..bb9287a6 100644 --- a/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h +++ b/StgSharp.Native/sn-head/sn_intrinsic.context.x86.h @@ -1,3 +1,6 @@ +#ifndef SN_INTRINSIC_CONTEXT_X86_H +#define SN_INTRINSIC_CONTEXT_X86_H + #include "sn_target.h" #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; #endif +#endif diff --git a/StgSharp.Native/sn-src/mk.fma.x86.c b/StgSharp.Native/sn-src/mk.fma.x86.c index dbeb1c38..36cb31c8 100644 --- a/StgSharp.Native/sn-src/mk.fma.x86.c +++ b/StgSharp.Native/sn-src/mk.fma.x86.c @@ -142,39 +142,35 @@ SN_MK_PROC_DECL_STD(float, 512, , fma) { __mk_param_std(TILE, LEFT_RIGHT_ANS, float); 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 c_1 = c_0; - register __m512 c_2 = c_0; - register __m512 c_3 = c_0; + register __m512 c0 = _mm512_setzero_ps(); + register __m512 c1 = _mm512_setzero_ps(); + register __m512 c2 = _mm512_setzero_ps(); + register __m512 c3 = _mm512_setzero_ps(); for (size_t z = 0; z < common_z; z+=4) { - SN_MAT_KERNEL_LOCAL const *right_ = right + z; - SN_MAT_KERNEL_LOCAL const *left_0 = left + z + 0; - 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; + const SN_MAT_KERNEL_LOCAL *a_base = left + z; + const SN_MAT_KERNEL_LOCAL *b_base = right + z; - // full right load：[kb_col0 | kb_col1 | kb_col2 | kb_col3] - register __m512 b_0 = _mm512_load_ps((float *)(right+0)); - register __m512 b_1 = _mm512_load_ps((float *)(right+1)); - register __m512 b_2 = _mm512_load_ps((float *)(right+2)); - register __m512 b_3 = _mm512_load_ps((float *)(right+3)); + register __m512 b0 = _mm512_load_ps((float *)&b_base[0].z[0]); + register __m512 b1 = _mm512_load_ps((float *)&b_base[1].z[0]); + register __m512 b2 = _mm512_load_ps((float *)&b_base[2].z[0]); + register __m512 b3 = _mm512_load_ps((float *)&b_base[3].z[0]); + + register __m512 a0 = _mm512_load_ps((float *)&a_base[0].z[0]); + register __m512 a1 = _mm512_load_ps((float *)&a_base[1].z[0]); + register __m512 a2 = _mm512_load_ps((float *)&a_base[2].z[0]); + register __m512 a3 = _mm512_load_ps((float *)&a_base[3].z[0]); + + 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); - // k=0: broadcast ka_col0 到4个lane，permute_ps 取 kb(0,j) - register __m512 ak; - register __m512 ak_1; - register __m512 ak_2; - register __m512 ak_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); - c_2= _mm512_add_ps(c_2, c_3); - c_0 = _mm512_add_ps(c_0, c_2); - _mm512_store_ps((float *)ans_cur, c_0); + SN_FMA_PACK_COLUMN(c0, (float*)&(ans_cur->x[0])); + SN_FMA_PACK_COLUMN(c1, (float*)&(ans_cur->x[1])); + SN_FMA_PACK_COLUMN(c2, (float*)&(ans_cur->x[2])); + SN_FMA_PACK_COLUMN(c3, (float*)&(ans_cur->x[3])); } diff --git a/StgSharp.Native/sn-template/sn_template.fma.x86.h b/StgSharp.Native/sn-template/sn_template.fma.x86.h index 87912bc8..9267f06a 100644 --- a/StgSharp.Native/sn-template/sn_template.fma.x86.h +++ b/StgSharp.Native/sn-template/sn_template.fma.x86.h @@ -2,25 +2,23 @@ #if SN_IS_ARCH(SN_ARCH_X86_64) -/* - * requested ZMMs: - * -*/ -#define SN_MK_FMA_F32_512_CYCLE(i)\ -do{\ - \ -}while(0) +#define SN_MK_FMA_F32_512_CYCLE(i) \ + 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))); \ + c2 = _mm512_fmadd_ps(c2, a2, _mm512_permute_ps(b2, _FMA_SELECT(i))); \ + c3 = _mm512_fmadd_ps(c3, a3, _mm512_permute_ps(b3, _FMA_SELECT(i))); \ + } while (0) -#define SN_MK_FMA_F32_512_PHASE(i) \ - do { \ - ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ - c_0 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ - ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ - c_1 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ - ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ - c_2 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ - ak = _mm512_broadcast_f32x4(left_##i->x[0]); \ - c_3 = _mm512_fmadd_ps(ak, _mm512_permute_ps(b_##i, _FMA_SELECT(0)), c_##i); \ +#define SN_FMA_PACK_COLUMN(col, target) \ + do { \ + register __m256 low = _mm512_castps512_ps256(col); \ + register __m256 high = _mm512_extractf32x8_ps(col, 1); \ + register __m256 half = _mm256_add_ps(high, low); \ + register __m128 high128 = _mm256_extractf32x4_ps(half, 1); \ + register __m128 low128 = _mm256_castps256_ps128(half); \ + register __m128 half128 = _mm_add_ps(high128, low128); \ + _mm_store_ps(target, half128); \ } while (0) // double