AVX2/FMA3实现运行时可变尺寸极小三矩阵乘积C=B'AB的C++优化

你当前采用的固定尺寸矩阵乘内核、两次通用矩阵乘法+临时转置存储的方案，在极小尺寸可变矩阵场景下性能不及串行是必然结果，核心问题如下：

• 固定步长8的内层循环完全不适配小矩阵：如果矩阵维度小于8，或者维度不是8的整数倍，现有代码要么直接越界出错，要么需要添加大量边界判断分支，分支预测失败的开销已经完全盖过SIMD指令的并行收益
• 现有内核是编译期固定尺寸N的实现，和运行时可变尺寸的需求不匹配，强行适配会引入大量冗余判断和无效计算
• 两次通用矩阵乘法+B转置临时存储的流程，在矩阵元素仅为个位数到几十的场景下，函数调用、内存拷贝、重复访存的成本远高于计算本身——C=B'*A*B属于典型的二次型结构，完全可以融合计算流程，不需要存储中间结果
• 现有内核每次内层循环都执行load、mul、add、store操作，没有做寄存器级数据复用，数据反复在内存和寄存器之间搬运，访存开销吃掉了SIMD的计算加速收益
• 没有针对数据类型适配FMA3指令：FMA3原生支持浮点融合乘加，硬套整数场景不会获得收益，反而可能因为指令选择不当降低效率

针对总元素量少、尺寸运行时可变的场景，不要套用大矩阵GEMM的分块、pack优化逻辑，这类逻辑的预热开销在矩阵元素少于几百个时完全无法收回成本，优化方向如下：

• 融合全计算流程：跳过temp=A*B的中间结果存储，直接展开C的元素计算公式C[i][j] = sum_k sum_l B[k][i] * A[k][l] * B[l][j]，全程在寄存器中完成计算，不需要额外临时空间
• 按尺寸分路径处理：
  • 当n、m维度均不超过4时，直接走标量计算路径，开启O2/O3优化后编译器会自动展开短循环，甚至生成最优的SIMD指令序列，性能优于手写硬编码AVX2
  • 当维度在5~16区间时，使用AVX2指令做向量计算，不需要复杂循环分块，仅处理向量宽度对齐的部分，剩余元素用标量收尾即可
• 取消单独的B转置操作：不需要提前把转置后的B存入临时数组，计算时直接按索引取B的对应元素即可，省掉转置的内存拷贝开销
• 按数据类型选指令：浮点场景直接使用FMA3融合乘加指令一次完成乘加操作，整数场景使用常规mul+add组合即可（FMA3无原生整数融合乘加指令）

以下实现针对int32类型极小矩阵，支持运行时可变尺寸，无临时矩阵开销：

#include <immintrin.h>
#include <cstring>

// 计算C = B^T * A * B
// 参数说明：
// A: n*n 方阵，行优先存储
// B: n*m 矩形矩阵，行优先存储
// C: m*m 结果方阵，行优先存储
// 适用场景：n*m < 256的极小尺寸矩阵
void small_quad_form_int32(const int* A, const int* B, int* C, int n, int m) {
    // 初始化结果矩阵为0
    memset(C, 0, m * m * sizeof(int));

    // 极小尺寸走标量路径，编译器自动展开
    if (n <= 4 && m <= 4) {
        for (int i = 0; i < m; i++) {
            for (int k = 0; k < n; k++) {
                const int bki = B[k * m + i];
                for (int l = 0; l < n; l++) {
                    const int coeff = bki * A[k * n + l];
                    for (int j = 0; j < m; j++) {
                        C[i * m + j] += coeff * B[l * m + j];
                    }
                }
            }
        }
        return;
    }

    // 中等偏小尺寸走AVX2路径
    for (int i = 0; i < m; i++) {
        for (int k = 0; k < n; k++) {
            const __m256i vec_bki = _mm256_set1_epi32(B[k * m + i]);
            for (int l = 0; l < n; l++) {
                const __m256i vec_akl = _mm256_set1_epi32(A[k * n + l]);
                const __m256i vec_coeff = _mm256_mullo_epi32(vec_bki, vec_akl);
                int j = 0;
                // 每次处理8个连续元素
                for (; j + 8 <= m; j += 8) {
                    const __m256i vec_blj = _mm256_loadu_si256((const __m256i*)&B[l * m + j]);
                    __m256i vec_c = _mm256_loadu_si256((const __m256i*)&C[i * m + j]);
                    vec_c = _mm256_add_epi32(vec_c, _mm256_mullo_epi32(vec_coeff, vec_blj));
                    _mm256_storeu_si256((__m256i*)&C[i * m + j], vec_c);
                }
                // 处理剩余不足8个的元素
                const int scalar_coeff = _mm256_cvtsi256_si32(vec_coeff);
                for (; j < m; j++) {
                    C[i * m + j] += scalar_coeff * B[l * m + j];
                }
            }
        }
    }
}

如果是单精度浮点计算，把整数intrinsic替换为对应浮点指令，核心乘加部分直接用_mm256_fmadd_ps即可利用FMA3的性能优势。

• 编译时必须开启-O2或-O3优化等级，附加-mavx2 -mfma编译选项，未开优化的手写SIMD性能远低于编译优化后的标量代码
• 不要为小矩阵引入大矩阵GEMM的pack、预取、多级分块逻辑，这类操作的固定开销在小矩阵场景下无法摊薄
• 如果矩阵元素数值范围允许，可改用int16甚至int8类型，进一步提升SIMD的并行宽度
• 矩阵存储尽量保持行优先连续，避免跨步访存带来的开销

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