AVX512加速矩阵乘法
最近打PKU的HPCGAME用的代码,这里只用上了20个zmm寄存器,改变block的大小应该还能优化一下速度。
代码只考虑了方阵,其他非2^n次方阵要自己改代码。具体原理很简单,看看代码就差不多知道。
const int BLOCK_SIZE = 1024;
const int BLOCK_SIZE2 = 256;
inline static void block_avx512_32x4( // AVX256效果不好,硬着头皮上吧(汇编上看好像还有12个寄存器没用上,还有优化空间)
int n, int K, //方阵大小
double* A, double* B, double* C)
{
__m512d c0000_0700,c0800_1500, c1600_2300, c2400_3100,
c0001_0701, c0801_1501, c1601_2301, c2401_3101,
c0002_0702, c0802_1502, c1602_2302, c2402_3102,
c0003_0703, c0803_1503, c1603_2303, c2403_3103;
__m512d a0x_7x, a8x_15x, a16x_23x, a24x_31x,
bx0, bx1, bx2, bx3;
double* c0001_0701_ptr = C + n;
double* c0002_0702_ptr = C + n * 2;
double* c0003_0703_ptr = C + n * 3;
c0000_0700 = _mm512_load_pd(C);
c0800_1500 = _mm512_load_pd(C + 8);
c1600_2300 = _mm512_load_pd(C + 16);
c2400_3100 = _mm512_load_pd(C + 24);
c0001_0701 = _mm512_load_pd(c0001_0701_ptr);
c0801_1501 = _mm512_load_pd(c0001_0701_ptr + 8);
c1601_2301 = _mm512_load_pd(c0001_0701_ptr + 16);
c2401_3101 = _mm512_load_pd(c0001_0701_ptr + 24);
c0002_0702 = _mm512_load_pd(c0002_0702_ptr);
c0802_1502 = _mm512_load_pd(c0002_0702_ptr + 8);
c1602_2302 = _mm512_load_pd(c0002_0702_ptr + 16);
c2402_3102 = _mm512_load_pd(c0002_0702_ptr + 24);
c0003_0703 = _mm512_load_pd(c0003_0703_ptr);
c0803_1503 = _mm512_load_pd(c0003_0703_ptr + 8);
c1603_2303 = _mm512_load_pd(c0003_0703_ptr + 16);
c2403_3103 = _mm512_load_pd(c0003_0703_ptr + 24);
for (int x = 0; x < K; ++x)
{
a0x_7x = _mm512_load_pd(A);
a8x_15x = _mm512_load_pd(A + 8);
a16x_23x = _mm512_load_pd(A + 16);
a24x_31x = _mm512_load_pd(A + 24);
A += 32;
bx0 = _mm512_broadcastsd_pd(_mm_load_sd(B++));
bx1 = _mm512_broadcastsd_pd(_mm_load_sd(B++));
bx2 = _mm512_broadcastsd_pd(_mm_load_sd(B++));
bx3 = _mm512_broadcastsd_pd(_mm_load_sd(B++));
c0000_0700 = _mm512_add_pd(_mm512_mul_pd(a0x_7x, bx0), c0000_0700);
c0800_1500 = _mm512_add_pd(_mm512_mul_pd(a8x_15x, bx0), c0800_1500);
c1600_2300 = _mm512_add_pd(_mm512_mul_pd(a16x_23x, bx0), c1600_2300);
c2400_3100 = _mm512_add_pd(_mm512_mul_pd(a24x_31x, bx0), c2400_3100);
c0001_0701 = _mm512_add_pd(_mm512_mul_pd(a0x_7x, bx1), c0001_0701);
c0801_1501 = _mm512_add_pd(_mm512_mul_pd(a8x_15x, bx1), c0801_1501);
c1601_2301 = _mm512_add_pd(_mm512_mul_pd(a16x_23x, bx1), c1601_2301);
c2401_3101 = _mm512_add_pd(_mm512_mul_pd(a24x_31x, bx1), c2401_3101);
c0002_0702 = _mm512_add_pd(_mm512_mul_pd(a0x_7x, bx2), c0002_0702);
c0802_1502 = _mm512_add_pd(_mm512_mul_pd(a8x_15x, bx2), c0802_1502);
c1602_2302 = _mm512_add_pd(_mm512_mul_pd(a16x_23x, bx2), c1602_2302);
c2402_3102 = _mm512_add_pd(_mm512_mul_pd(a24x_31x, bx2), c2402_3102);
c0003_0703 = _mm512_add_pd(_mm512_mul_pd(a0x_7x, bx3), c0003_0703);
c0803_1503 = _mm512_add_pd(_mm512_mul_pd(a8x_15x, bx3), c0803_1503);
c1603_2303 = _mm512_add_pd(_mm512_mul_pd(a16x_23x, bx3), c1603_2303);
c2403_3103 = _mm512_add_pd(_mm512_mul_pd(a24x_31x, bx3), c2403_3103);
}
_mm512_storeu_pd(C, c0000_0700);
_mm512_storeu_pd(C + 8, c0800_1500);
_mm512_storeu_pd(C + 16, c1600_2300);
_mm512_storeu_pd(C + 24, c2400_3100);
_mm512_storeu_pd(c0001_0701_ptr, c0001_0701);
_mm512_storeu_pd(c0001_0701_ptr + 8, c0801_1501);
_mm512_storeu_pd(c0001_0701_ptr + 16, c1601_2301);
_mm512_storeu_pd(c0001_0701_ptr + 24, c2401_3101);
_mm512_storeu_pd(c0002_0702_ptr, c0002_0702);
_mm512_storeu_pd(c0002_0702_ptr + 8, c0802_1502);
_mm512_storeu_pd(c0002_0702_ptr + 16, c1602_2302);
_mm512_storeu_pd(c0002_0702_ptr + 24, c2402_3102);
_mm512_storeu_pd(c0003_0703_ptr, c0003_0703);
_mm512_storeu_pd(c0003_0703_ptr + 8, c0803_1503);
_mm512_storeu_pd(c0003_0703_ptr + 16, c1603_2303);
_mm512_storeu_pd(c0003_0703_ptr + 24, c2403_3103);
}
static inline void copy_avx512_b(int lda, const int K, double* b_src, double* b_dest) {
double* b_ptr0, * b_ptr1, * b_ptr2, * b_ptr3;
b_ptr0 = b_src;
b_ptr1 = b_ptr0 + lda;
b_ptr2 = b_ptr1 + lda;
b_ptr3 = b_ptr2 + lda;
for (int i = 0; i < K; ++i)
{
*b_dest++ = *b_ptr0++;
*b_dest++ = *b_ptr1++;
*b_dest++ = *b_ptr2++;
*b_dest++ = *b_ptr3++;
}
}
static inline void copy_avx512_a(int lda, const int K, double* a_src, double* a_dest) {
for (int i = 0; i < K; ++i)
{
memcpy(a_dest, a_src, 32 * 8);
a_dest += 32;
a_src += lda;
}
}
static inline void do_block_avx512(int lda, int M, int N, int K, double* A, double* B, double* C)
{
double* A_block, * B_block;
A_block = (double*)_mm_malloc(M * K * sizeof(double), 64);
B_block = (double*)_mm_malloc(K * N * sizeof(double), 64);
double* a_ptr, * b_ptr, * c;
const int Nmax = N - 3;
int Mmax = M - 32;
int i = 0, j = 0, p = 0;
for (j = 0; j < Nmax; j += 4)
{
b_ptr = &B_block[j * K];
copy_avx512_b(lda, K, B + j * lda, b_ptr); // 将 B 展开
for (i = 0; i < Mmax; i += 32) {
a_ptr = &A_block[i * K];
if (j == 0) copy_avx512_a(lda, K, A + i, a_ptr); // 将 A 展开
c = C + i + j * lda;
block_avx512_32x4(lda, K, a_ptr, b_ptr, c);
}
}
_mm_free(A_block);
_mm_free(B_block);
}
void gemm_avx512(int lda, double* A, double* B, double* C)
{
#pragma omp parallel for
for (int j = 0; j < lda; j += BLOCK_SIZE) { // j i k 序 内存读写更快
for (int i = 0; i < lda; i += BLOCK_SIZE) {
for (int k = 0; k < lda; k += BLOCK_SIZE) {
// 大分块里小分块
for (int jj = j; jj < j + BLOCK_SIZE; jj += BLOCK_SIZE2)
for (int ii = i; ii < i + BLOCK_SIZE; ii += BLOCK_SIZE2)
for (int kk = k; kk < k + BLOCK_SIZE; kk += BLOCK_SIZE2)
do_block_avx512(lda, BLOCK_SIZE2, BLOCK_SIZE2, BLOCK_SIZE2, A + ii + kk * lda, B + kk + jj * lda, C + ii + jj * lda);
}
}
}
}
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 微软正式发布.NET 10 Preview 1:开启下一代开发框架新篇章
· 没有源码,如何修改代码逻辑?
· PowerShell开发游戏 · 打蜜蜂
· 在鹅厂做java开发是什么体验
· WPF到Web的无缝过渡:英雄联盟客户端的OpenSilver迁移实战
2022-02-01 32位程序获取32/64位进程文件地址通用方法