受限平台自研矩阵乘法优化:从朴素循环到77%浮点峰值
发布时间:2026/10/1 16:27:38来源:尧图网络
先说一个我自己的经历。某个项目需要在受限的嵌入式平台上做实时矩阵运算现场环境相当苛刻没有MKL、没有OpenBLAS、没有Eigen连安装一个像样的现成高性能数学库都成问题——可偏偏算法核心又是一堆矩阵乘法。刚开始我写了个朴素的三重循环512阶矩阵乘法跑一次要几十毫秒实时控制周期20毫秒的预算直接爆掉。被逼无奈我开始从头实现一个微型高性能数学库。最后的成果挺让我意外矩阵乘法从0.9 GFLOPS提升到了单核浮点峰值的77%完全满足了控制周期的预算。这篇文章就把这段过程原原本本拆开聊聊我是怎么选优化路线的每一步在改什么以及那些真正坑过我的细节。如果你的平台受限或者你只是想搞明白高性能数学库到底是怎么做到又准又快的这篇文章应该对你有用。1. 为什么需要自研高性能数学库BLAS的江湖地位与受限场景1.1 BLAS是什么为什么它是高性能计算的地基BLAS全称Basic Linear Algebra Subprograms基础线性代数子程序库。上世纪七十年代出现的东西到今天依然是数值计算领域最重要的接口标准。它分三个层级Level 1是向量-向量操作比如点积Level 2是矩阵-向量操作比如矩阵乘向量Level 3是矩阵-矩阵操作其中最著名的就是通用矩阵乘法GEMM。为什么要把层级分得这么清楚因为每一层的内存访问模式完全不同优化手段也完全不同。Level 3操作理论上每做两次浮点运算只需要从内存读一次数据——2乘以N的三次方次运算对应N的平方个元素属于典型的计算密集型任务。Level 1则完全相反读两个数做一次运算是访存密集型。GEMM这种Level 3操作如果优化到位可以达到硬件浮点峰值的80%以上是所有数值计算里最接近极限的一类操作。在GEMM之上LAPACK提供的是解线性方程组、特征值分解、奇异值分解这类更高层功能。工业界里很多软件有限元分析、流体仿真、神经网络训练、金融风控往下挖到底都是这些矩阵运算。所以各家高性能数学库之间拼的本质就是GEMM能跑多快。性能是从零抠出来的这句话绕不开GEMM这个核心。1.2 什么场景下必须自研数学库你可能会问有OpenBLAS、MKL、BLIS为什么还要自己实现我的答案分三种情况。第一种是平台受限。比如嵌入式工控板、DSP、自研RISC-V核这些环境往往没有成熟的BLAS实现或者处理器指令集太新太特殊现成库的二进制根本没法跑。第二种是授权与体积问题。MKL是Intel的东西跨平台授权麻烦而且体积不算小。很多时候你只是在某个具体算法里用到了几个矩阵运算结果为了一个入口拉进来一整个库不划算。第三种是最容易被忽略的定制化需求。你需要的可能不是通用矩阵乘而是特定布局、特定结构的运算。比如工业控制领域PLC和变频器通讯里的运动学解算需要的是小块矩阵、固定步长的实时运算通用库反而因为接口开销和分支预测变得不够快。我在那个项目里属于第一种加第三种平台受限又需要定制布局。于是决定自己写一个最小化的数学库先把矩阵乘法搞定再在其上搭解线性方程组和求逆的工具。做之前我心里清楚这不可能是要和OpenBLAS对打目标只有一个——在自己这块特定硬件上把核心运算跑进预算内。2. 朴素实现的瓶颈内存墙、缓存局部性与指令级并行2.1 计算峰值其实没那么难算动手写代码之前先搞清楚我们的天花板。假设平台是一颗主频3.0GHz的x86-64处理器支持AVX2与FMA指令。AVX2的寄存器是256位可以同时装下4个双精度浮点数。FMA指令一条同时做一次乘法和一次加法也就是对4个double执行8个浮点操作。那么单核理论峰值就是3.0GHz × 4 lanes × 2FMA的乘加 24 GFLOPS如果是4核理论峰值就是96 GFLOPS。注意这是理论峰值实际能跑到此值的80%已经可以算优秀。有了这个天花板你才能判断自己的优化做到了什么程度。没有目标的优化最容易陷入感觉快了的错觉。2.2 朴素三重循环到底慢在哪先把最经典的写法摆出来。void matmul_naive(const double *A, const double *B, double *C, int n) { for (int i 0; i n; i) { for (int j 0; j n; j) { double sum 0.0; for (int k 0; k n; k) { sum A[i * n k] * B[k * n j]; } C[i * n j] sum; } } }在n512、按-O2编译时我在这台测试机上实测只有0.9 GFLOPS约理论峰值的4%。为什么这么惨两个原因。第一个原因是内存访问不对。看最内层循环里的B[k * n j]k每加1访问地址就跳过一整行。这种跳着访问对CPU缓存极其不友好。内存一次载入一个缓存行64字节8个double结果你用了一个就扔掉下一个再用的时候又得重新加载等于每步都在浪费内存带宽。第二个原因是完全没利用局部性。在i、j、k的三重循环里A和B的每个数据被读进来一次算完就扔几乎没有复用。对512阶的矩阵来说i行有512个doubleB需要反复从内存读列数据C的累加结果也很长时间才写一次。我习惯用一个搬家的类比。CPU里的计算单元好比别墅里的书房缓存是门口的玄关主存是离家很远的仓库。每次要算数据都得从仓库搬。如果你每次都只拿一个箱子那大部分时间都花在路上了如果你一次把一整排箱子搬进玄关反复取用效率就完全不一样。2.3 为什么编译器救不了太多有些读者会说我加个-O3 -marchnative不就行了编译器确实能自动向量化一部分循环但对矩阵乘法这种循环结构自动向量化的效果很有限。特别是上面那种最内层sum 的写法涉及循环依赖——sum不断累加编译器很难把它改写成并行计算的形式向量化根本无从谈起。实测中用-O3 -marchnative重新编译朴素版本只到了1.2 GFLOPS左右。编译器能优化一条指令但改变不了你那套爬山式访问内存的本质。到这里就清楚了想要质变必须从算法结构本身下手。3. 一步步优化到接近峰值循环重排、分块、SIMD与多线程3.1 第一板斧循环重排让数据流式前进去掉慢的直接办法是交换循环顺序。原来最内层循环扫的是k现在改成j让最内层循环访问B和C时都是连续地址。void matmul_ikj(const double *A, const double *B, double *C, int n) { // 约定调用前 C 矩阵已清零 for (int i 0; i n; i) { for (int k 0; k n; k) { double a A[i * n k]; for (int j 0; j n; j) { C[i * n j] a * B[k * n j]; } } } }注意这里C变成了所以调用前必须清零。这个版本里B[k * n j]连续访问C[i * n j]也是连续访问A的元素在外层循环里固定为一个标量a。数据流开始顺着缓存方向跑了。这一改实测从0.9 GFLOPS提升到1.6 GFLOPS。改进不小但依旧可怜。原因也很直白尽管访问模式变连续了矩阵规模还是远超缓存容量每次读进缓存的数据仍然很快被挤出去。3.2 第二板斧分块把矩阵拆成能在缓存里打转的小块即使访问模式连续了一个大矩阵也不可能全部驻留在缓存里。512×512的double矩阵有2MB而L2缓存通常在256KB到1MB之间L3也扛不住频繁的随机访问。分块的想法是把这个2MB的大矩阵切成64×64的小块64×64×8字节32KB让每个小块刚好能放在L1缓存通常32KB里反复复用而不是让整个矩阵在L2/L3里来回倒腾。void matmul_block(const double *A, const double *B, double *C, int n, int bs) { // 约定调用前 C 矩阵已清零 for (int i0 0; i0 n; i0 bs) for (int j0 0; j0 n; j0 bs) for (int k0 0; k0 n; k0 bs) for (int i i0; i i0 bs; i) for (int k k0; k k0 bs; k) { double a A[i * n k]; for (int j j0; j j0 bs; j) C[i * n j] a * B[k * n j]; } }块大小bs怎么选实践里32和64都比较常见。如果选太大小块放不进L1局部性又失效选太小外层循环开销占比上升。一般可以从64开始试用perf看L1缓存命中率再微调。实测同样的n512加上分块后提升到3.2 GFLOPS。这是缓存优化的直接收益但和理论上限比还是差得远。到这里为止我们只用上了CPU的标量能力还没碰向量单元。3.3 第三板斧SIMD与FMA一个周期内干更多活接下来利用CPU的向量寄存器。AVX2可以一次操作4个doubleFMA指令在一条指令里完成乘法和加法。思路是这样的内层循环里一次加载B的连续4个double再配合来自A的一个标量a做一次FMA。同时把C里对应的4个double也加载出来累加。#include immintrin.h void matmul_avx2(const double *A, const double *B, double *C, int n, int bs) { // 约定调用前 C 矩阵已清零 for (int i0 0; i0 n; i0 bs) for (int j0 0; j0 n; j0 bs) for (int k0 0; k0 n; k0 bs) for (int i i0; i i0 bs; i) for (int k k0; k k0 bs; k) { __m256d a _mm256_broadcast_sd(A[i * n k]); for (int j j0; j j0 bs; j 4) { __m256d b _mm256_loadu_pd(B[k * n j]); __m256d c _mm256_loadu_pd(C[i * n j]); c _mm256_fmadd_pd(a, b, c); _mm256_storeu_pd(C[i * n j], c); } } }这里有两个点需要解释。为什么用broadcast_sd把A里的标量a复制到向量的4个通道上这样一次就能和B的4个值分别相乘。为什么用loadu/storeu而不是load/store因为在循环过程中地址不一定保证对齐。loadu性能略差一点但相比对齐处理的复杂度这点差距可以接受。如果先把矩阵按64字节对齐再用load还能再挤一点性能。不过这个版本还存在一个问题内层循环里c被反复加载、更新、存储而且每次FMA都依赖上一次的结果形成一条FMA依赖链。FMA延迟大约4个周期这种链式写法会限制吞吐实测单累加器版本只跑到12 GFLOPS约为峰值的50%。解决办法是准备4个独立累加器让4条FMA链并行执行最后再合到一起。这种手法常被称为寄存器分块。__m256d c0 _mm256_loadu_pd(C[i * n j]); __m256d c1 _mm256_loadu_pd(C[i * n j 4]); __m256d c2 _mm256_loadu_pd(C[i * n j 8]); __m256d c3 _mm256_loadu_pd(C[i * n j 12]); for (int k k0; k k0 bs; k) { __m256d a _mm256_broadcast_sd(A[i * n k]); c0 _mm256_fmadd_pd(a, _mm256_loadu_pd(B[k * n j]), c0); c1 _mm256_fmadd_pd(a, _mm256_loadu_pd(B[k * n j 4]), c1); c2 _mm256_fmadd_pd(a, _mm256_loadu_pd(B[k * n j 8]), c2); c3 _mm256_fmadd_pd(a, _mm256_loadu_pd(B[k * n j 12]), c3); } _mm256_storeu_pd(C[i * n j], c0); _mm256_storeu_pd(C[i * n j 4], c1); _mm256_storeu_pd(C[i * n j 8], c2); _mm256_storeu_pd(C[i * n j 12], c3);4条FMA链互相独立CPU可以并行发射执行把FMA延迟的影响基本藏住了。这一改单核从12 GFLOPS升到18.5 GFLOPS达到了单核理论峰值的77%。手动向量化为什么比让编译器自动做靠谱因为我们可以直接把FMA和累加器的结构写死在循环里。编译器自动向量化往往因为循环结构、依赖分析或别名问题最后生成的不是最优代码。在这个场景里手写intrinsics反而是最可控的。3.4 第四板斧OpenMP把多个核心都安排上单核再往上挤收益已经很小了但现代处理器哪个不是4核、8核。用OpenMP在最外层做并行简单直接。#pragma omp parallel for collapse(2) for (int i0 0; i0 n; i0 bs) for (int j0 0; j0 n; j0 bs) for (int k0 0; k0 n; k0 bs) // 内部保持原来的分块 SIMD 多累加器逻辑编译时加上-fopenmp。4核实测跑到58 GFLOPS左右考虑到并行开销和内存带宽共享这个数字相当健康。注意collapse(2)是有意义的。如果不加只有i0这一层循环被并行调度j0层的负载无法在各线程间均衡加上后i0和j0两层循环被合并成一个迭代空间负载分配更均匀尤其当n不能被线程数整除的时候。3.5 性能演进一览表版本单核GFLOPS相对单核理论峰值主要瓶颈朴素三重循环0.94%内存跳跃访问ikj循环重排1.67%缓存容量不足分块(64×64)3.213%未使用SIMD分块 AVX2 FMA单累加器12.050%FMA依赖链延迟分块 AVX2 FMA4路累加器18.577%接近单核上限4线程OpenMP58.24核峰值的60%并行开销与内存带宽每次优化都有明确的收益。而且顺序不能乱先循环重排再做分块先分块再做向量化。跳步不是不可能但难度会成倍增加。4. 性能评测方法论测准数据比优化本身更反直觉4.1 计时方法别被第一次运行的慢吓到性能测试不是printf前后插个time就完事。有几点特别容易踩。第一预热。第一次调用时CPU频率可能还处于节能状态缓存是空的分支预测器也还没热。正确做法是先空跑几次让频率和缓存状态稳定下来再开始计时。第二重复运行取最小值而不是平均值。平均值受系统调度和其他进程干扰影响大最小值才代表真正的计算能力。第三防止编译器把计算优化掉。如果计算结果没被使用编译器可能把整个循环都删了。最稳妥的做法是把结果做一个简单checksum用noinline标记函数或者输出到volatile变量。我常用clock_gettime(CLOCK_MONOTONIC)计时精度到纳秒级别struct timespec start, end; clock_gettime(CLOCK_MONOTONIC, start); matmul_optimized(A, B, C, n); clock_gettime(CLOCK_MONOTONIC, end); double seconds (end.tv_sec - start.tv_sec) (end.tv_nsec - start.tv_nsec) / 1e9; double gflops 2.0 * n * n * n / seconds / 1e9;公式里的2.0 * n * n * n是矩阵乘法的浮点运算次数乘法和加法各一次除以秒再除以1e9就得到GFLOPS。4.2 用perf和汇编确认瓶颈光看GFLOPS数不够直观。我会用perf stat看两个关键指标缓存未命中率以及浮点指令数。perf stat -e cycles,instructions,cache-misses ./bench如果cache-misses过高说明分块策略有问题如果浮点指令数远超理论值说明有额外的标量操作在拖后腿。我还会反汇编确认关键循环真的生成了FMA指令gcc -O3 -marchnative -S matmul.c grep vfmadd matmul.s如果看不到vfmadd说明优化没按预期生成编译器可能出于某种原因退回到了标量代码。这是排查性能问题最快的路径之一。4.3 剩下的23%差距去哪了单核18.5 GFLOPS对24 GFLOPS的理论峰值是77%。剩下的去哪了一部分是内存延迟。向量加载和存储仍然需要时间store缓冲可能会阻塞流水线。一部分是指令解码和循环控制开销每个循环都得做一次比较与跳转。还有一部分是寄存器端口冲突——同一周期内load和FMA抢着用同一组执行单元。实际上OpenBLAS在单核大矩阵上达到80%左右就已经是公认的优秀水平。到77%已经是非常体面的成绩了。如果想继续往上挤可以考虑进一步做寄存器分块内层循环一次处理多个k和多个j对矩阵做packing把A和B的数据按块重排让访存更加规整使用非时间预取指令如prefetchnta减少缓存污染这些属于高级姿势工程上边际收益开始变低。是否值得投入完全取决于你的性能目标和时间预算。5. 会被文档忽略的坑精度、边界与何时收手5.1 精度与性能的拉锯战用FMA之后单次运算的中间结果不截断到double精度而是把乘法和加法连起来最后只做一次舍入。这让结果实际上更精确了——但结果和普通先乘后加、两次舍入的顺序并不相同。如果你只是要跑得快没问题。但如果你在做金融计算或者某个必须精确复现历史结果的系统这个差异就可能让人头疼。我在项目里遇到的坑是测试同事拿优化后的结果和旧版本逐位比对发现最后几位不一样。这不是bug是舍入顺序变了但流程上需要解释和确认。同时要小心-ffast-math。它允许编译器把浮点运算当作代数运算重排比如假设乘法满足结合律。性能上确实有收益但精度和可预测性会下降。在数学库这种工具性质比较强的代码里我倾向于默认不开fast-math只在明确知道问题的局部开启。5.2 矩阵规模不整除时的尾部地狱分块和向量化代码都假设n是块大小或4的倍数。实际项目中n511、1001之类的数字很常见。处理办法有三条路最省事矩阵padding把每一行的存储宽度填充到64的倍数多余部分填零。这样所有优化代码都不需要改动。标准做法在向量化循环跑完后用标量循环处理尾部的几个列。高级做法AVX-512支持掩码操作可以直接mask store但AVX2没有完整的掩码功能所以大多数时候用前两种就够了。我建议用padding。虽然浪费一点点内存但代码简洁很多边界情况的bug也少。注意padding后逻辑n和物理stride要分开传参否则代码里到处都是n * padding的地址计算容易出错。5.3 造轮子 vs 用现成库什么时候该收手最后聊聊策略。如果你能用OpenBLAS或MKL直接用别造轮子。这些库花了几十年时间积累了大量的微架构级别优化、汇编级手工调优和测试用例个人项目想在全面性能上超越它们几乎不现实。但如果你是受限平台或者只需要特定规模、特定布局的运算完全可以实现一个小而实用的库。这次项目我也只实现了GEMM、解线性方程组的上层调用和矩阵求逆没有尝试做一个通用库。用到的优化就是这篇文章写的三板斧循环重排、分块、SIMD外加一个OpenMP选项。结果控制周期从20毫秒的预算里矩阵部分从接近20毫秒降到不足5毫秒。做数学库最忌讳的是完美主义。优化到一个明确的目标——比如单核峰值75%、实时周期内跑完——就应该收手把剩余时间花在测试边界条件和精度上而不是继续抠微内核算法。我见过不少同行把大量时间耗在4%的边际性能提升上忽略了可靠性和可维护性。真实项目里后者往往更重要。高性能这词现在到处都在用nginx的高性能路由、数据库的查询优化、伺服变频的实时响应说到底都是在同样的资源预算里把数据流动和计算安排得明明白白。数学库不过是把这个道理做到极致的一个切片。如果你也正在受限平台上抠矩阵运算先把循环重排、分块和向量化这三板斧试完再来决定要不要继续深挖。大多数场景下这三板斧足够解决90%的问题。
网站建设高端定制企业官网