如何在C++中优化高频调用的(5,5)×(5,25)小矩阵乘法？

问题背景 #

我需要优化一个执行3次(5,5)×(5,25)小矩阵乘法的函数mxm(A,B1,B2,B3)：A为(5,5)形状，B1-B3为(5,25)形状，输出C1-C3满足C1=A*B1、C2=A*B2、C3=A*B3。该函数在HPC程序中会被调用数十亿次，目前已实现Eigen库版本和手动循环展开版本，测试显示循环展开版本耗时仅为Eigen版本的1/4左右，希望找到更优的优化方案。

现有代码 #

main.cpp

#include <iostream>
#include <time.h>
#include "test.hpp"

int main(){
    Eigen::Array<float,5,5,1> A;
    Eigen::Array<float,5,25,1> B1,B2,B3,C1,C2,C3,D1,D2,D3;

    A.setRandom(); B1.setRandom(); 
    B2.setRandom(); B3.setRandom();

    clock_t tic,toc;
    const int nn = 500000;

    tic = clock();
    for(int i = 0; i < nn; i ++){
        mxm5_eigen(A.data(),B1.data(),B2.data(),B3.data(),C1.data(),
            C2.data(),C3.data());
    }
    toc = clock();
    std::cout << "Eigen time " << (double)(toc-tic) / CLOCKS_PER_SEC << "\n"; 

    tic = clock();
    for(int i = 0; i < nn; i ++){
        mxm5(A.data(),B1.data(),B2.data(),B3.data(),D1.data(),
            D2.data(),D3.data());
    }
    toc = clock();
    std::cout << "unrolled loop time " << (double)(toc-tic) / CLOCKS_PER_SEC << "\n";
}

test.hpp

#include <Eigen/Core>

void inline mxm5_eigen(const float *__restrict Ain,const float *__restrict B1in,
          const float *__restrict B2in,const float *__restrict B3in,
          float *__restrict C1in,float *__restrict C2in,float *__restrict C3in)
{
    const int n1 = 5, n3 = 25;
    Eigen::Map<const Eigen::Matrix<float,5,5,1>> A(Ain);
    Eigen::Map<const Eigen::Matrix<float,5,25,1>> B1(B1in),B2(B2in),B3(B3in);
    Eigen::Map<Eigen::Matrix<float,5,25,1>> C1(C1in),C2(C2in),C3(C3in);

    C1.noalias() = A * B1;
    C2.noalias() = A * B2;
    C3.noalias() = A * B3;
}

// c[i,j] = sum_k A[i,k] * B[k,j]
void inline mxm5(const float *__restrict A,const float *__restrict B1,
          const float *__restrict B2,const float *__restrict B3,
          float *__restrict C1,float *__restrict C2,float *__restrict C3)
{
    const int  n1 = 5, n3 = 25;
    for(int i = 0; i < n1; i ++){
        #pragma GCC ivdep
        for(int j = 0; j < n3; j ++){
            int idx = i * n3 + j;
            C1[idx] = A[i * n1 + 0] * B1[0 * n3 + j] + 
                      A[i * n1 + 1] * B1[1 * n3 + j] +
                      A[i * n1 + 2] * B1[2 * n3 + j] +
                      A[i * n1 + 3] * B1[3 * n3 + j] +
                      A[i * n1 + 4] * B1[4 * n3 + j];

            C2[idx] = A[i * n1 + 0] * B2[0 * n3 + j] + 
                      A[i * n1 + 1] * B2[1 * n3 + j] +
                      A[i * n1 + 2] * B2[2 * n3 + j] +
                      A[i * n1 + 3] * B2[3 * n3 + j] +
                      A[i * n1 + 4] * B2[4 * n3 + j];

            C3[idx] = A[i * n1 + 0] * B3[0 * n3 + j] + 
                      A[i * n1 + 1] * B3[1 * n3 + j] +
                      A[i * n1 + 2] * B3[2 * n3 + j] +
                      A[i * n1 + 3] * B3[3 * n3 + j] +
                      A[i * n1 + 4] * B3[4 * n3 + j];                   
        }
    }
}

编译命令与测试结果 #

编译命令：

g++ main.cpp -O3 -march=native -ffast-math -I/path/to/eigen/

测试结果：

Eigen time 0.078125
unrolled loop time 0.03125

进阶优化方法 #

1. 数据布局优化：合并批量矩阵，提升内存连续性 #

当前B1/B2/B3是独立的(5,25)矩阵，内存上分散存储，导致SIMD指令无法一次性加载足够的连续数据。可以将三个B矩阵合并为一个**(5, 75)**的连续矩阵（按列优先存储），对应C矩阵也合并为(5,75)，这样：

• B的第k行第3*j + b列对应原Bb+1的第k行第j列（b=0,1,2）
• C的第i行第3*j + b列对应原Cb+1的第i行第j列

这种布局能让CPU一次性加载更多连续元素，最大化SIMD指令的利用率。

2. 寄存器复用：减少重复内存访问 #

当前代码在每个j循环中都会重复读取A的行元素，虽然编译器会做部分优化，但手动将A的行加载到寄存器变量中，能进一步减少内存访问开销：

// 在i循环内，先加载A的第i行到寄存器变量
float a0 = A[i*5 + 0], a1 = A[i*5 + 1], a2 = A[i*5 + 2], a3 = A[i*5 + 3], a4 = A[i*5 + 4];

之后在j循环中直接使用这些变量，避免每次计算都从内存读取A的元素。

3. 手动SIMD向量化：利用CPU向量指令 #

针对x86平台的AVX2/AVX-512指令集，手动编写向量运算代码，一次性处理多个j的计算。例如，用AVX2的256位向量（可存储8个float），同时计算2-4个j位置的C1/C2/C3值：

#include <immintrin.h>

void inline mxm5_avx2(const float *__restrict A, const float *__restrict B, 
          float *__restrict C) {
    const int n1 = 5, n3 = 25, batch = 3;
    for(int i = 0; i < n1; i++){
        // 加载A的第i行到寄存器，广播为向量
        __m256 a0 = _mm256_set1_ps(A[i*5 + 0]);
        __m256 a1 = _mm256_set1_ps(A[i*5 + 1]);
        __m256 a2 = _mm256_set1_ps(A[i*5 + 2]);
        __m256 a3 = _mm256_set1_ps(A[i*5 + 3]);
        __m256 a4 = _mm256_set1_ps(A[i*5 + 4]);

        for(int j = 0; j < n3 * batch; j += 8){
            // 加载B的5行，每行8个连续元素
            __m256 b0 = _mm256_loadu_ps(&B[0 * n3 * batch + j]);
            __m256 b1 = _mm256_loadu_ps(&B[1 * n3 * batch + j]);
            __m256 b2 = _mm256_loadu_ps(&B[2 * n3 * batch + j]);
            __m256 b3 = _mm256_loadu_ps(&B[3 * n3 * batch + j]);
            __m256 b4 = _mm256_loadu_ps(&B[4 * n3 * batch + j]);

            // 计算向量乘法并累加
            __m256 c = _mm256_add_ps(_mm256_mul_ps(a0, b0), _mm256_mul_ps(a1, b1));
            c = _mm256_add_ps(c, _mm256_mul_ps(a2, b2));
            c = _mm256_add_ps(c, _mm256_mul_ps(a3, b3));
            c = _mm256_add_ps(c, _mm256_mul_ps(a4, b4));

            // 存储结果到C
            _mm256_storeu_ps(&C[i * n3 * batch + j], c);
        }
    }
}

注：这里假设合并后的B矩阵是连续存储的，若原B1/B2/B3是分散的，可以先做一次内存拷贝合并，或者直接在程序中使用合并后的布局。

4. 编译器指令优化 #

• 用#pragma omp simd simdlen(N)替代#pragma GCC ivdep：omp simd兼容性更好，能明确指定向量长度N（比如8对应AVX2，16对应AVX-512），帮助编译器生成更优的向量化代码。
• 编译时明确指定指令集：比如-mavx2或-mavx512f，比-march=native更可控，确保编译器生成目标指令集的代码。
• 增加-funroll-loops：让编译器自动展开更多循环，减少循环控制开销。

5. 精度折中优化（可选） #

如果程序允许降低数值精度，可以考虑：

• 使用float16半精度：AVX-512支持FP16指令，能将内存带宽提升一倍，计算吞吐量也更高。
• 使用-ffast-math的进阶选项：比如-ffp-contract=fast允许编译器进行浮点收缩优化，进一步提升计算效率。

内容的提问来源于stack exchange，提问作者Suwen Jun