让矩阵乘法快四万倍的优化之旅

2,123 阅读8分钟

开启掘金成长之旅!这是我参与「掘金日新计划 · 2 月更文挑战」的第 1 天,点击查看活动详情

关于MIT6.172

MIT6.172课程名为《Performance Engineering of Software Systems》,也即软件系统的性能工程,其主要探讨的是如何在软件层面对软件性能进行优化。具体的优化思路则是基于如cache这样的硬件特性来协助优化,以便更好的使用硬件能力。主题包括性能分析、高性能算法技术、指令级优化、缓存优化、并行编程和构建可扩展系统。

为什么需要性能工程?

当我们提到性能,我们最先想到的会是什么呢?从高维的角度,我们可以认为是低延时、高并发等特性。而如果当我们把视角放到微架构层面,性能就可以从某种程度上认为成在单位时间内运行的指令数。假设单位时间内CPU可以运行一万个指令,那么我们可以用这些指令去做很多事情,例如身份校验,计算和数据存储等功能,这些都是通过指令来实现的。换句话说,性能可以作为计算的货币。我们在计算机上写的每个额外的操作,都需要以性能作为代价去交换。

在过去的一段时间里,CPU的性能提升主要依靠的是提高时钟频率,如下图所示:

image-20230128152643957

可以看到,当到了2004年左右,CPU的时钟频率已经几乎不在上升了;为了能够继续提高CPU的性能,人们将目光转向了增加CPU的核数来提高CPU的性能:

image-20230128153415655

为了提高硬件性能,厂商们使用了诸如多核处理器、多层cache、向量处理器等特性。需要意识到的是,摩尔定律已经快要失效了,我们应该去注重软件的性能,合理的使用硬件特性来提高软件性能。

案例学习:矩阵乘法

在本部分,将会以一个矩阵乘法的例子来直观的体现合理的使用硬件特性对性能的影响。矩阵乘法的实现如下图:

image-20230128154056129

我们在同一台机器上进行代码的运行和结果的测算,配置如下:

image-20230128154236893

时钟频率是2.9GHz,两个芯片,每个芯片由带有9个2路超线程的处理器组成,而每个核每周期可以做8次双精度操作,也即该机器的理想性能为:836G次浮点运算每周期。

python实现

image-20230128154953857

我们通过如上的代码实现4096*4096的矩阵乘法。可以看出,做了一个三重循环,每次循环的维度是n;而在循环内部,做了一次乘法运算一次加法运算,也即整体的算法复杂度为2n^3。该代码执行需要21042秒,约6小时才能执行完成。

代入一下n为4096,则该代码一共执行2n^3=2(2^{12})^3=2^{37}次浮点运算,平均到每一秒约为2^{37}/21042=6.25MFLOPS,对比我们的机器峰值836GFLOPS,可以得出python代码的效率仅为0.00075%。

这非常的符合python这样一门解释型的语言,其运行速度确实是比较慢的。

Java实现

我们尝试用另一门非常火热的语言Java来进行测试。其代码逻辑与python类似。矩阵大小仍然为4096*4096。

从结果上看,Java耗时2738秒,是python运行速度的8.8倍。

考虑到Java的虚拟机特性,其运行速度也在意料之中。

C实现

接着我们保持同样的运行条件,用C这样一门非常底层的语言来进行测试。

从结果上看,C耗时约1156秒,是python耗时的18倍,Java的两倍。

小结

image-20230128160145829

python是一种解释型语言,相对而言的执行效率最慢;Java首先被编译成字节码,然后在虚拟机中被解释执行,速度相对快;C则是直接被编译成了二进制执行,所以速度相对最快。也即语言的不同也有性能的差异,越靠近底层则速度相对越快。

优化Loop

image-20230128160733341

上图是Loop的代码,我们可以看出在保持i、j不动的情况下,k每次替换,都会去访问B[k][j],例如k[0][0]到k[1][0]。而在二维矩阵的存储上,其实际上是顺序存储的,如下图所示:

image-20230128161641176

或者是这样:

image-20230128161731648

那么每一次k变化,都会由当前的row跳到下一个row,例如当k从0变为1的时候,就会从row1跳到row2中;而每个row中实际上存储了4096个数字,这就导致了很差的局部性,每一次的cache都无法命中,从而拖慢了程序运行速度,如下图所示:

image-20230128162114799

尝试对每一种i、j、k的顺序进行测试:

image-20230128162351051

不妨以i、k、j的顺序为例:

image-20230128162426828

在内层循环中,随着j的变化,C这个结果矩阵访问的是C[i][j],只在单独的row内进行访问;A矩阵访问A[i][k],与j无关;B矩阵访问B[k][j],也是在row内,保证了良好的局部性,这使得整体的效率大大提高。

通过工具检查每种顺序的cache miss情况:

image-20230128162621380

符合我们对于cache miss影响程序运行效率的认知。

编译器优化

随着编译器的强化,其能提供越来越多的编译优化选项,我们尝试开启不同程序的优化选项进行测试:

image-20230128162733193

可以看出,相比于不进行优化的版本,O2和O3选项得到了约三倍的性能提升。

并行运行

在前文我们说过,实验机器是一个拥有2块9核处理器的机器,也即我们可以使用18个核来进行运算。当需要并行的时候,我们需要考虑好如何对数据进行分片。

作为一个三重循环,我们对于最内层循环和最外层循环做分片是比较合适的,我们分别进行尝试:

image-20230128163209093

可以看出,对最外层做分片是最适合的,相对而言把该任务划分的最明细。从结果上看也可以看出,并行得到了约18倍的性能提升。

image-20230128163433126

平铺数据

当我们计算一行的矩阵乘法时,我们需要写入C矩阵4096次,读取A矩阵4096次以及读取B矩阵4096*4096次。

image-20230128163811269

而如果我们计算一块C矩阵时,例如我们计算一个64*64的块,我们需要写入C矩阵4096次,读取A矩阵64*4096次,读取B矩阵4096*64次,这显然低于计算一行的读取次数。而读取数据也是需要消耗时间的,还可能会带来cache miss,我们不妨平铺计算。

image-20230128163901837

我们尝试按照不同的值s对C矩阵进行s*s的划分,结果分别如下:

image-20230128164122395

同时我们统计对应的cache miss情况:

image-20230128164156396

可以看到,随着cache miss的提高,整体的运行速度也提升了。

分治

考虑到矩阵乘法的特性:

image-20230128164652650

以及前文对于平铺数据的理解,我们不妨基于2的指数把任务进行分治:

image-20230128164859458

从结果上可以看出,我们可以获得一些提升,但相对幅度并不大:

image-20230128165050373

向量硬件

在前文我们提到,硬件厂商提供了诸多新特性,向量硬件就是其中的一部分,也即计算机体系结构中提到的SIMD(单指令多数据流)模型。而对于我们目前的这个矩阵乘法,就是使用向量硬件非常好的例子,可以成行的进行数据读取和计算。

image-20230128165324429

AVX指令

相对于使用编译器做的向量化处理,我们也可以自己调用硬件厂商提供的向量指令,如intel提供的AVX系列指令。

image-20230128165435377

小结

在以上的过程中,我们自顶向下的利用了硬件的特性,从而来提高代码的效率。在整个流程中,我们都遵循思考、改动、验证的流程来进行代码的优化。整体的效果也非常惊人,在同一台机器上,我们获得了大约四万多倍的性能提升。 这样的改动对于同样的机器如果都适用的话,那么就可以获得惊人的规模效应,极大的提高硬件效率节约成本。