macro gymnastics to implement storage kernel, buffer op and complex compute panel

This commit is contained in:
Youmu Liu
2025-11-13 19:09:27 +08:00
parent 2dfb4ae86a
commit 1417785726
18 changed files with 373 additions and 383 deletions
+4 -13
View File
@@ -3,16 +3,9 @@
#ifdef _WIN32
// clang-format off
BOOL APIENTRY DllMain(
HMODULE hModule,
DWORD ul_reason_for_call,
LPVOID lpReserved
)
BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
switch (ul_reason_for_call)
{
switch (ul_reason_for_call) {
case DLL_PROCESS_ATTACH:
return 1;
case DLL_THREAD_ATTACH:
@@ -21,9 +14,9 @@ BOOL APIENTRY DllMain(
return 1;
case DLL_PROCESS_DETACH:
return 1;
break;
break;
}
return TRUE;
return TRUE;
}
#else
attribute((constructor)) void DllMain()
@@ -31,5 +24,3 @@ attribute((constructor)) void DllMain()
}
#endif // _WIN32
// clang-format on
+52 -58
View File
@@ -1,29 +1,34 @@
#include "sn_intrinsic.h"
INTERNAL void SN_DECL f32_fill_sse(MAT_KERNEL(float) *restrict dest, scalar_pack *scaler)
DECLARE_BUF_PROC_ANS_SCALAR(float, sse, fill)
{
register __m128 vec = _mm_set_ps1(((float *)(scaler->data))[0]);
dest->xmm[0] = vec;
dest->xmm[1] = vec;
dest->xmm[2] = vec;
dest->xmm[3] = vec;
register __m128 vec = _mm_set_ps1(((float *)(scalar->data))[0]);
for (size_t i = 0; i < count; i++) {
ans[i].xmm[0] = vec;
ans[i].xmm[1] = vec;
ans[i].xmm[2] = vec;
ans[i].xmm[3] = vec;
}
}
INTERNAL void SN_DECL f32_fill_avx(MAT_KERNEL(float) *restrict dest, scalar_pack *scaler)
DECLARE_BUF_PROC_ANS_SCALAR(float, avx, fill)
{
register __m256 vec = _mm256_broadcast_ss(((float *)(scaler->data)));
dest->ymm[0] = vec;
dest->ymm[1] = vec;
register __m256 vec = _mm256_broadcast_ss(((float *)(scalar->data)));
for (size_t i = 0; i < count; i++) {
ans[i].ymm[0] = vec;
ans[i].ymm[1] = vec;
}
}
INTERNAL void SN_DECL f32_fill_512(MAT_KERNEL(float) *restrict dest, scalar_pack *scaler)
DECLARE_BUF_PROC_ANS_SCALAR(float, 512, fill)
{
register __m512 vec = _mm512_broadcastss_ps(((__m128 *)(scaler->data))[0]);
dest->zmm[0] = vec;
register __m512 vec = _mm512_broadcastss_ps(((__m128 *)(scalar->data))[0]);
for (size_t i = 0; i < count; i++) {
ans[i].zmm[0] = vec;
}
}
INTERNAL void SN_DECL f32_add_sse(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, sse, add)
{
register __m128 c0 = _mm_add_ps(left->xmm[0], right->xmm[0]);
register __m128 c1 = _mm_add_ps(left->xmm[1], right->xmm[1]);
@@ -36,8 +41,7 @@ INTERNAL void SN_DECL f32_add_sse(MAT_KERNEL(float) const *left, MAT_KERNEL(floa
ans->xmm[3] = c3;
}
INTERNAL void SN_DECL f32_add_avx(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, avx, add)
{
register __m256 c0 = _mm256_add_ps(left->ymm[0], right->ymm[0]);
register __m256 c1 = _mm256_add_ps(left->ymm[1], right->ymm[1]);
@@ -45,15 +49,13 @@ INTERNAL void SN_DECL f32_add_avx(MAT_KERNEL(float) const *left, MAT_KERNEL(floa
ans->ymm[1] = c1;
}
INTERNAL void SN_DECL f32_add_512(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, 512, add)
{
register __m512 c0 = _mm512_add_ps(left->zmm[0], right->zmm[0]);
ans->zmm[0] = c0;
}
INTERNAL void SN_DECL f32_sub_sse(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, sse, sub)
{
register __m128 c0 = _mm_sub_ps(left->xmm[0], right->xmm[0]);
register __m128 c1 = _mm_sub_ps(left->xmm[1], right->xmm[1]);
@@ -66,8 +68,7 @@ INTERNAL void SN_DECL f32_sub_sse(MAT_KERNEL(float) const *left, MAT_KERNEL(floa
ans->xmm[3] = c3;
}
INTERNAL void SN_DECL f32_sub_avx(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, avx, sub)
{
register __m256 c0 = _mm256_sub_ps(left->ymm[0], right->ymm[0]);
register __m256 c1 = _mm256_sub_ps(left->ymm[1], right->ymm[1]);
@@ -75,37 +76,34 @@ INTERNAL void SN_DECL f32_sub_avx(MAT_KERNEL(float) const *left, MAT_KERNEL(floa
ans->ymm[1] = c1;
}
INTERNAL void SN_DECL f32_sub_512(MAT_KERNEL(float) const *left, MAT_KERNEL(float) const *right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, 512, sub)
{
register __m512 c0 = _mm512_sub_ps(left->zmm[0], right->zmm[0]);
ans->zmm[0] = c0;
}
// Matrix transpose implementations
INTERNAL void SN_DECL f32_transpose_sse(MAT_KERNEL(float) const *source,
MAT_KERNEL(float) *restrict target)
DECLARE_KER_PROC_RIGHT_ANS(float, sse, transpose)
{
register __m128 col0 = source->xmm[0];
register __m128 col1 = source->xmm[1];
register __m128 col2 = source->xmm[2];
register __m128 col3 = source->xmm[3];
register __m128 col0 = right->xmm[0];
register __m128 col1 = right->xmm[1];
register __m128 col2 = right->xmm[2];
register __m128 col3 = right->xmm[3];
register __m128 tmp0 = _mm_unpacklo_ps(col0, col1);
register __m128 tmp1 = _mm_unpackhi_ps(col0, col1);
register __m128 tmp2 = _mm_unpacklo_ps(col2, col3);
register __m128 tmp3 = _mm_unpackhi_ps(col2, col3);
target->xmm[0] = _mm_movelh_ps(tmp0, tmp2);
target->xmm[1] = _mm_movehl_ps(tmp2, tmp0);
target->xmm[2] = _mm_movelh_ps(tmp1, tmp3);
target->xmm[3] = _mm_movehl_ps(tmp3, tmp1);
ans->xmm[0] = _mm_movelh_ps(tmp0, tmp2);
ans->xmm[1] = _mm_movehl_ps(tmp2, tmp0);
ans->xmm[2] = _mm_movelh_ps(tmp1, tmp3);
ans->xmm[3] = _mm_movehl_ps(tmp3, tmp1);
}
INTERNAL void SN_DECL f32_transpose_avx(MAT_KERNEL(float) const *source,
MAT_KERNEL(float) *restrict target)
DECLARE_KER_PROC_RIGHT_ANS(float, avx, transpose)
{
register __m256 col_low = source->ymm[0];
register __m256 col_high = source->ymm[1];
register __m256 col_low = right->ymm[0];
register __m256 col_high = right->ymm[1];
register __m256 unpack_low = _mm256_unpacklo_ps(col_low, col_high);
register __m256 unpack_high = _mm256_unpackhi_ps(col_low, col_high);
@@ -113,30 +111,28 @@ INTERNAL void SN_DECL f32_transpose_avx(MAT_KERNEL(float) const *source,
col_low = _mm256_permute2f128_ps(unpack_low, unpack_high, 0x20);
col_high = _mm256_permute2f128_ps(unpack_low, unpack_high, 0x31);
target->ymm[0] = _mm256_shuffle_ps(col_high, col_low, _MM_SHUFFLE(2, 0, 2, 0));
target->ymm[1] = _mm256_shuffle_ps(col_high, col_low, _MM_SHUFFLE(3, 1, 3, 1));
ans->ymm[0] = _mm256_shuffle_ps(col_high, col_low, _MM_SHUFFLE(2, 0, 2, 0));
ans->ymm[1] = _mm256_shuffle_ps(col_high, col_low, _MM_SHUFFLE(3, 1, 3, 1));
}
INTERNAL void SN_DECL f32_transpose_512(MAT_KERNEL(float) const *source,
MAT_KERNEL(float) *restrict target)
DECLARE_KER_PROC_RIGHT_ANS(float, 512, transpose)
{
register __m512 matrix = source->zmm[0];
register __m512 matrix = right->zmm[0];
const __m512i transpose_indices =
_mm512_setr_epi32(0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15);
target->zmm[0] = _mm512_permutexvar_ps(transpose_indices, matrix);
ans->zmm[0] = _mm512_permutexvar_ps(transpose_indices, matrix);
}
INTERNAL void SN_DECL f32_scalar_mul_sse(MAT_KERNEL(float) const *matrix,
MAT_KERNEL(float) *restrict ans, scalar_pack const *scalar)
DECLARE_KER_PROC_RIGHT_ANS_SCALAR(float, sse, scalar_mul)
{
register __m128 scalar_vec = ((__m128 *)(scalar->data))[0];
register __m128 c0 = _mm_mul_ps(matrix->xmm[0], scalar_vec);
register __m128 c1 = _mm_mul_ps(matrix->xmm[1], scalar_vec);
register __m128 c2 = _mm_mul_ps(matrix->xmm[2], scalar_vec);
register __m128 c3 = _mm_mul_ps(matrix->xmm[3], scalar_vec);
register __m128 c0 = _mm_mul_ps(right->xmm[0], scalar_vec);
register __m128 c1 = _mm_mul_ps(right->xmm[1], scalar_vec);
register __m128 c2 = _mm_mul_ps(right->xmm[2], scalar_vec);
register __m128 c3 = _mm_mul_ps(right->xmm[3], scalar_vec);
ans->xmm[0] = c0;
ans->xmm[1] = c1;
@@ -144,24 +140,22 @@ INTERNAL void SN_DECL f32_scalar_mul_sse(MAT_KERNEL(float) const *matrix,
ans->xmm[3] = c3;
}
INTERNAL void SN_DECL f32_scalar_mul_avx(MAT_KERNEL(float) const *matrix,
MAT_KERNEL(float) *restrict ans, scalar_pack const *scalar)
DECLARE_KER_PROC_RIGHT_ANS_SCALAR(float, avx, scalar_mul)
{
register __m256 scalar_vec = ((__m256 *)(scalar->data))[0];
register __m256 c0 = _mm256_mul_ps(matrix->ymm[0], scalar_vec);
register __m256 c1 = _mm256_mul_ps(matrix->ymm[1], scalar_vec);
register __m256 c0 = _mm256_mul_ps(right->ymm[0], scalar_vec);
register __m256 c1 = _mm256_mul_ps(right->ymm[1], scalar_vec);
ans->ymm[0] = c0;
ans->ymm[1] = c1;
}
INTERNAL void SN_DECL f32_scalar_mul_512(MAT_KERNEL(float) const *matrix,
MAT_KERNEL(float) *restrict ans, scalar_pack const *scalar)
DECLARE_KER_PROC_RIGHT_ANS_SCALAR(float, 512, scalar_mul)
{
register __m512 scalar_vec = ((__m512 *)(scalar->data))[0];
register __m512 c0 = _mm512_mul_ps(matrix->zmm[0], scalar_vec);
register __m512 c0 = _mm512_mul_ps(right->zmm[0], scalar_vec);
ans->zmm[0] = c0;
}
+1 -50
View File
@@ -1,14 +1,8 @@
#include "sn_intrinsic.h"
INTERNAL void SN_DECL f32_fma_sse(MAT_PANEL_SSE(int) * left)
{
}
// ans = left * right + ans
// right is not transposed
INTERNAL void SN_DECL f32_fma_sse(MAT_KERNEL(float) const *restrict left,
MAT_KERNEL(float) const *restrict right,
MAT_KERNEL(float) *restrict ans)
DECLARE_KER_PROC_LEFT_RIGHT_ANS(float, sse, fma)
{
register __m128 source0 = _mm_load_ps(&left->xmm[0]);
register __m128 source1 = _mm_load_ps(&left->xmm[1]);
@@ -39,46 +33,3 @@ INTERNAL void SN_DECL f32_fma_sse(MAT_KERNEL(float) const *restrict left,
_mm_store_ps(&ans->xmm[2], ans2);
_mm_store_ps(&ans->xmm[3], ans3);
}
// ans = left * right + ans
INTERNAL void SN_DECL f32_fma_avx(MAT_KERNEL(float) const *restrict left,
MAT_KERNEL(float) const *restrict right,
MAT_KERNEL(float) *restrict ans)
{
register __m256 source0 = _mm256_load_ps(&left->ymm[0]);
register __m256 source1 = _mm256_load_ps(&left->ymm[1]);
register __m256 ans0 = _mm256_load_ps(&ans->ymm[0]);
register __m256 ans1 = _mm256_load_ps(&ans->ymm[1]);
register __m128 bs0, bs1;
register __m256 broadcast;
for (size_t i = 0; i < 4; i++) {
bs0 = _mm_load_ps(&right->xmm[i]);
__m128 lo = _mm_shuffle_ps(bs0, bs0, 0x00);
__m128 hi = _mm_shuffle_ps(bs0, bs0, 0x55);
broadcast = _mm256_castps128_ps256(lo);
broadcast = _mm256_insertf128_ps(broadcast, bs0, 1);
ans0 = _mm256_add_ps(ans0, _mm256_mul_ps(source0, broadcast));
lo = _mm_shuffle_ps(bs0, bs0, 0xAA);
hi = _mm_shuffle_ps(bs0, bs0, 0xFF);
broadcast = _mm256_castps128_ps256(lo);
broadcast = _mm256_insertf128_ps(broadcast, bs0, 1);
ans1 = _mm256_add_ps(ans1, _mm256_mul_ps(source1, broadcast));
}
_mm256_store_ps(&ans->ymm[0], ans0);
_mm256_store_ps(&ans->ymm[1], ans1);
}
// ans = left * right + ans
INTERNAL void SN_DECL f32_fma_512(MAT_KERNEL(float) const *restrict left,
MAT_KERNEL(float) const *restrict right,
MAT_KERNEL(float) *restrict ans)
{
/*
* to be implemented
*/
}
+15 -15
View File
@@ -65,11 +65,11 @@ void load_intrinsic_sse(sn_intrinsic *context)
context->city_hash_simplify = city_hash_simplify_sse;
context->factorial_simd = factorial_simd_sse;
context->index_pair = index_pair_sse;
context->f32_add = (MAT_PROC_LRA)f32_add_sse;
context->f32_fma = (MAT_PROC_LRA)f32_fma_sse;
context->f32_scalar_mul = f32_scalar_mul_sse;
context->f32_sub = (MAT_PROC_LRA)f32_sub_sse;
context->f32_transpose = (MAT_PROC_LA)f32_transpose_sse;
context->f32_kernel_add = KER_PROC(float, sse, add);
context->f32_panel_fma = KER_PROC(float, sse, fma);
context->f32_kernel_scalar_mul = KER_PROC(float, sse, scalar_mul);
context->f32_kernel_sub = KER_PROC(float, sse, sub);
context->f32_kernel_transpose = KER_PROC(float, sse, transpose);
context->f32_normalize_3 = f32_normalize;
}
@@ -78,11 +78,11 @@ void load_intrinsic_avx(sn_intrinsic *context)
context->city_hash_simplify = city_hash_simplify_sse;
context->factorial_simd = factorial_simd_sse;
context->index_pair = index_pair_sse;
context->f32_add = (MAT_PROC_LRA)f32_add_avx;
context->f32_fma = (MAT_PROC_LRA)f32_fma_avx;
context->f32_scalar_mul = f32_scalar_mul_avx;
context->f32_sub = (MAT_PROC_LRA)f32_sub_avx;
context->f32_transpose = (MAT_PROC_LA)f32_transpose_avx;
context->f32_kernel_add = KER_PROC(float, avx, add);
context->f32_panel_fma = KER_PROC(float, sse, fma);
context->f32_kernel_scalar_mul = KER_PROC(float, avx, scalar_mul);
context->f32_kernel_sub = KER_PROC(float, avx, sub);
context->f32_kernel_transpose = KER_PROC(float, avx, transpose);
context->f32_normalize_3 = f32_normalize;
}
@@ -91,10 +91,10 @@ void load_intrinsic_512(sn_intrinsic *context)
context->city_hash_simplify = city_hash_simplify_sse;
context->factorial_simd = factorial_simd_sse;
context->index_pair = index_pair_sse;
context->f32_add = (MAT_PROC_LRA)f32_add_512;
context->f32_fma = (MAT_PROC_LRA)f32_fma_512;
context->f32_scalar_mul = f32_scalar_mul_512;
context->f32_sub = (MAT_PROC_LRA)f32_sub_512;
context->f32_transpose = (MAT_PROC_LA)f32_transpose_512;
context->f32_kernel_add = KER_PROC(float, 512, add);
context->f32_panel_fma = KER_PROC(float, sse, fma);
context->f32_kernel_scalar_mul = KER_PROC(float, 512, scalar_mul);
context->f32_kernel_sub = KER_PROC(float, 512, sub);
context->f32_kernel_transpose = KER_PROC(float, 512, transpose);
context->f32_normalize_3 = f32_normalize;
}