基于AVX-512的列主序矩阵乘法优化及cblas_sgemm适配问询

我用AVX512 intrinsics实现了浮点矩阵乘法，代码如下：

for (int i=200; i<400; i++) {
        for (int k=1200; k<1400; k++) {
            tmp=val[440000+ (k-1200)*200 + (i-200)];
            sv=_mm512_set1_ps(tmp);
            for (int j=0; j<512/16; j++) {
                cv = _mm512_loadu_ps(&y[i*512 + 16*j]);
                bv = _mm512_loadu_ps(&x[k*512 + 16*j]);
                cv = _mm512_fmadd_ps(sv, bv, cv);
                _mm512_storeu_ps(&y[512*i+16*j], cv);
            }
        }
    }

其中val以列主序存储，i、k、j为矩阵维度索引。编译参数为：gcc -O3 -march=native -mavx -mprefer-vector-width=512 -ffast-math。

我尝试了循环展开与合并，得到优化后的代码：

for (int i=200; i<400; i++) {
    for (int k=1400; k<1600; k+=2) {
        sv[0] = _mm512_set1_ps(val[520000 + (k+-1400) * 200 + (i-200)]);
        sv[1] = _mm512_set1_ps(val[520000 + (k+-1399) * 200 + (i-200)]);
        for (int j=0; j<512/16; j++){
            cv = _mm512_loadu_ps(&y[i*512 + 16*j]);
            bv[0] = _mm512_loadu_ps(&x[(k+0)*512 + 16*j]);
            cv = _mm512_fmadd_ps(sv[0], bv[0], cv);
            bv[1] = _mm512_loadu_ps(&x[(k+1)*512 + 16*j]);
            cv = _mm512_fmadd_ps(sv[1], bv[1], cv);
            _mm512_storeu_ps(&y[512*i+16*j], cv);
        }
    }
}

此外还尝试过j循环展开、j循环分块、转置val数组。现提出两个问题：

1. 除并行化外，还有哪些方法可进一步优化这些循环？
2. 是否可以直接改用cblas_sgemm实现？

1. 除并行化外的其他优化方向 #

• 优化内存访问模式：
  • 若x实际为列主序存储，当前按k*512 + 16*j的访问方式会导致非连续内存访问，缓存命中率低。可提前转置x为行主序，或调整循环顺序让内存访问连续。
  • val是列主序，当前索引方式下k递增时访问同一列连续元素，这部分没问题，但可尝试用_mm512_load_ps结合广播指令代替标量加载后广播，减少标量操作开销。
• 减少寄存器读写与存储开销：
  • 避免每次j循环都加载、存储y，可将cv寄存器累积多次计算后再一次性存回内存，降低内存IO次数。
  • 提升k维度的循环展开倍数（比如从2到4或8），利用AVX512的多寄存器能力增加指令级并行，让CPU执行单元充分利用。
• 优化指令选择：
  • 确保x、y数组内存地址64字节对齐，改用_mm512_load_ps/_mm512_store_ps代替非对齐版本，避免额外开销。可通过__attribute__((aligned(64)))或posix_memalign实现对齐。
  • 用_mm512_broadcastss_ps代替_mm512_set1_ps，前者可直接从内存广播到向量寄存器，减少中间标量指令。
• 循环顺序与分块优化：
  • 调整循环顺序（如改为k→i→j），提升内存访问的空间局部性。同时采用分块（Tiling）策略，将大矩阵拆分为适配CPU缓存大小的小块（如32x32、64x64），让块内计算全程利用缓存，减少缓存miss。
• 编译器优化辅助：
  • 给x、y、val数组添加__restrict__关键字，告知编译器数组无内存重叠，允许更激进的优化。
  • 明确添加-mavx512f编译参数，确保编译器生成最优的AVX512指令，而非依赖-march=native的自动检测。

2. 能否直接改用cblas_sgemm实现？ #

完全可以，且这通常是更优选择，原因如下：

• 高度优化：CBLAS（如OpenBLAS、MKL等实现）的sgemm由专业团队优化，不仅充分利用AVX512等SIMD指令，还针对不同CPU架构调整了循环分块、内存调度等细节，性能普遍优于手动编写的intrinsics代码。
• 简洁易维护：调用cblas_sgemm仅需几行代码，无需手动处理循环和指令细节，降低出错概率，提升代码可维护性。
• 架构适配性强：CBLAS库会自动检测CPU架构，启用对应指令集优化，无需手动调整编译参数或代码。

使用时需注意：

• 根据x、y、val的实际存储顺序，设置cblas_sgemm的行/列主序参数（CblasRowMajor或CblasColMajor）。
• 正确映射矩阵维度，你的计算逻辑对应y = y + val * x，需设置alpha=1.0、beta=1.0，并传递正确的矩阵尺寸参数。

内容的提问来源于stack exchange，提问作者Pratyush Das