C++如何优化矩阵乘法 gemm

[Avafly]

最近在用 C 手写模型推理, 其中 gemm 可以说是核心计算, 于是决定以学习为目的自己尝试优化一下.

用 3 个 for 循环可以实现最基本的矩阵乘法, 在我用 simd, blocking, 并行计算这些方法之后, 速度比 naive 版本的快了很多, 但还是会比 openblas 慢不少. 接下来该怎么做有点没头绪了. 我想知道有没有办法能进一步提升? 谢谢

代码地址: https://github.com/Avafly/optimize-gemm

[nagisaushio]

函数签名改成这样试试

void matmul(const float * restrict A, const float * restrict B, float * restrict C,
const int M, const int N, const int K)

[elfive]

如果是 intel 平台，可以考虑使用 Intel Intrinsics Library ，会比 OpenBLAS 快不少。

[elfive]

@elfive 另外还有个 intel MKL 库可以用，这个库有办法在 AMD 平台加速，具体方法可以参考这个博客： https://monsoon-cs.moe/2023-06-19-mkl-on-amd/

[AirCrusher]

可以看下这个 repo ，循序渐进：
https://github.com/yzhaiustc/Optimizing-DGEMM-on-Intel-CPUs-with-AVX512F

[Avafly]

@elfive #2 什么库不重要, 主要是想自己优化 gemm 来学习一下. 实际项目中会都测试一边选性能最好的用的.

[Avafly]

@nagisaushio 这个确实有一些帮助, 不过只能提升一点点大概 0.1GFLOPS 吧, 还是和 openblas, blis 这些有断档的差距. 感觉更多还是算法设计方面的问题, 这部分不知道该怎么做了.

[Avafly]

@AirCrusher 谢谢分享, 这个有点猛汇编都用上了, 我回头看下. 其实后面我看过类似的就是 flame 的教程, 基本上里面的技术都应用到了已经.

[Donaldo]

歪个楼，请问你 README 的矩阵图是用什么画的？

[Avafly]

@Donaldo ppt😂

[Donaldo]

@Avafly #9 👍一样，平常作图全靠 ppt 。

[WonderfulRush]

可以看看这篇文章 里面讲了 cpu 矩阵乘的优化 https://justine.lol/matmul/

[Avafly]

@WonderfulRush
刚看完这篇文章然后看到你的评论...
那个文章挺好的, 但是技术部分讲得有点简略, 而且其实很多提到的技术我已经用了, 比如 blocking, simd 等等.

[toma62299781]

你这个刚好就是 MIT 韩松老师开的 MIT 6.5940 Efficient AI 课程的 lab5 哇，可以去参考一下:
课程主页： https://hanlab.mit.edu/courses/2024-fall-65940
lab5： https://drive.google.com/drive/folders/1MhMvxvLsyYrN-4C6eQG8Zj2JeSuyAOf0

[tankeco]

我感觉你取 index 就做了不少乘法，不确定编译器能不能帮你优化掉，你自己把 index 改成累加的方式试试

[Avafly]

@toma62299781
感谢分享

[Avafly]

@tankeco
是的, 这点我也觉得要花时间想下怎么减少 index.
其实已经优化过一次 index 了, 现在保留的都是为了分块和区分多线程访问空间的, 后面个人感觉这不是影响速度的最大的因素就没继续花心思了.

[dingyaguang117]

因为人家复杂度就不是 O(N^3) ...

[foool]

对比 openblas 中 cblas_sgemm 也是 4 并行度的吗？

[foool]

几个小建议和疑问：
1 先大致理论分析下最大能够达到的 GFLOS 是多少（考虑 CPU 多 port 都可以执行运算）；
2 先用单线程跑到最高速率，排除多线程调度或资源竞争导致的劣化；
3 尝试加上预取指令，perf 看看瓶颈在哪里
4 测试多次，取最优值，看你测试了一次，都会有“冷启动”的问题。
5 omp parallel for schedule(static) 是在编译时就确定代码位于哪个线程了吗，会导致 cache 相关问题吗（ false sharing ）

[dazhangpan]

使用 AMX 指令

[Avafly]

@foool #19
非常感谢你的回复.
1. 最大 GFLOPS 这个我没算, 是以 openblas 的为目标优化的 (试过别的库, 有比 openblas 更快的).
2. 3. 很好的建议, 我回头再优化测试看看.
4. 我是脚本跑 100 次取最优值的.
5. 使用 schedule(static)是因为 for 循环中每次计算量近似才用的, 不过我试过去掉这个, 其实性能基本没区别.

[Avafly]

@dingyaguang117
你是说它们不是传统的 C=AB, 而是用了 Strassen/Winograd 之类的方法减少了复杂度吗?

[dingyaguang117]

@Avafly 不确定 openblas 是否用到了 Strassen/Winograd 算法，感觉想要追求更显著突破可能是需要换算法的

[B1ankCat]

https://doi.org/10.1109/FOCS57990.2023.00130 ，一个思路

[tdxdxoz]

@Avafly 不熟悉你这个问题的背景，但是矩阵乘法本身现在已经能做到 O(n^2.37)左右了。比较实用、常数没那么大的也有 O(n^2.8)左右的算法。

[tdxdxoz]

@Avafly 就是你说的 Strassen/Winograd