代码性能优化——分块技巧

代码性能优化——分块技巧

同样的算法、同样的计算量，一些高手写的程序可能比初学者的快百倍。像MKL、Openblas这些专业的计算库，更是将性能优化到了极致。

那么他们是怎么做到的呢？本篇我们就从分块技巧开始来管中窥豹一下。

代码实例

开门见山，以一个实用的程序作为例子：向量乘矩阵

如图所示,M个元素的向量乘以M行N列的矩阵，生成N个元素的向量。
基本代码如下：

#define M (2048)
#define N (1024)

static void naive(float vec_in[M], float mat[M][N], float vec_out[N]){
    for(int n = 0; n < N; n++){      // 每次循环算出一个输出
        for(int m = 0; m < M; m++){  // M个元素乘累加
            vec_out[n] += vec_in[m] * mat[m][n];
        }
    }
}

程序耗时：27.45 ms

这里假设vec_out内存的初始值都已清零，故可以直接累加。

接下来我们就对他进行分块改造。设每块的元素个数为BLK_SZ。首先，把内层循环等价拆分为两层循环：

#define M (2048)
#define N (1024)
#define BLK_SZ (32)

static void blocking(float vec_in[M], float mat[M][N], float vec_out[N]){
    for(int n = 0; n < N; n++){
        for(int m0 = 0; m0 < M / BLK_SZ; m0++){  // 外层循环，分为M/BLK_SZ个块
            for(int m1 = 0; m1 < BLK_SZ; m1++){  // 内层循环，块内计算
                vec_out[n] += vec_in[m0 * BLK_SZ + m1] * mat[m0 * BLK_SZ + m1][n];
            }
        }
    }
}

再把m0的循环提升到最外层，这样只是改变了运算的顺序，不影响结果。

#define M (2048)
#define N (1024)
#define BLK_SZ (32)

static void blocking(float vec_in[M], float mat[M][N], float vec_out[N]){
    for(int m0 = 0; m0 < M / BLK_SZ; m0++){    // 移动到最外层
        for(int n = 0; n < N; n++){
            for(int m1 = 0; m1 < BLK_SZ; m1++){
                vec_out[n] += vec_in[m0 * BLK_SZ + m1] * mat[m0 * BLK_SZ + m1][n];
            }
        }
    }
}

程序耗时：6.54 ms

以上程序均使用gcc -O3编译，可以看到程序的速度变为原来的4倍多。

理论讲解

为什么同样的计算量耗时却差异很大？奥秘在于缓存（cache）。

CPU执行指令，要先从外部的主存将数据取到寄存器。由于主存速度很慢，通常在中间增加了比主存更小但更快的cache，三者的性能关系如上图。

cache用于自动缓存一份常用的数据。CPU取数时先看看cache中有没有，有则能快速取到（cache hit)，没有则只能从主存中慢速的取出了(cache miss)。

这就和平时收拾东西是一样的道理。我们的桌面是空间有限，为此要把常用的东西放到桌面以快速取用，不常用的放到更远的抽屉、柜子。如此才能效率最高。

为了保证cache中的数据为常用数据，计算机使用如下的策略向其中搬运：

• 如果CPU读取的数据不在cache中，则从主存将该数据及其周围的一块数据一起搬入cache中。
• 如果cache容量不足，则新数据替换一块旧的数据。

使用该策略是因为它很好的适应了程序的特性，即局部性原理：

• 时间局部性：如果一个存储器位置被引用了一次，那么程序很可能在不远的将来再次重复引用它。
• 空间局部性：如果一个存储器位置被引用了一次，那么程序很可能在不远的将来引用附近的一个存储器位置。

cache中搬入了刚刚用过的数据和周围的一块数据，很可能程序接下来会被用到。比如循环重复使用同一个数据时，再比如顺序访问一个数组时。如此可以尽量减少访问主存的次数，提高性能。

不仅如此，这样从主存一次搬运一块连续的数据，也比分成多次独立的搬运更快。（突发传输）

实例解析

结合理论，再回到我们最初的例子。

为简单起见，我们只考虑输入向量的cache使用情况。假设从主存向cache一次搬运 4 格数据，cache最多存储 8 格。

上图列举了各次cache miss的和为此搬入的数据。下面具体分步骤看一下：

• 最开始cache为空。
• 计算位置1时，CPU发现该数据cache中没有（cache miss)，故而从主存中搬运一块数据（编号1～4）到cache。由于要从主存取数，本次耗时较大。
• 计算位置2、3、4时，由于数据已经在cache中了，故可以快速取到（cache hit）。
• 计算位置5时，再次发生了cache miss。为此又搬运了一块数据到cache（编号5～8）
• 计算位置6、7、8时，cache hit。
• 计算位置9时，再次发生了cache miss。并且此时cache已经存满。为此，替换一块数据（1～4）存储9～12。
• 计算位置10、11、12，cache hit

至此，我们完成了向量和矩阵第 1 列的运算，得到的第 1 个输出值。发生了 3 次cache miss。

接下来输入向量要和矩阵第 2 列运算了。然而此时有一个尴尬的问题：位置1～4的数据已经从cache中替换出去了，此时使用又需要从主存中搬运。最差的情况下每列运算都cache miss 3 次。而这，就是程序耗时大的根源。

既然和矩阵每一列计算时，都需要用向量位置1～4的数据，那我们可不可以趁他们还在cache中时直接把这些计算全算完了？

当然可以! 这就是实例代码中的分块所做的事情：

如图所示，当第一次用到输入向量位置1的数据时，因cache miss将1～4搬入cache。此时将这 4 格数据分别和矩阵的每一列作乘累加运算，并保存到输出向量。当然此时输出的位置只累加了部分数据，没关系，后面计算接着累加即可。

        for(int n = 0; n < N; n++){
            for(int m1 = 0; m1 < BLK_SZ; m1++){
                vec_out[n] += vec_in[m0 * BLK_SZ + m1] * mat[m0 * BLK_SZ + m1][n];
            }
        }

完成上述运算，位置1～4的数据已经不会再用到了。因此不必在意它是否会被从cache替换掉了。

后面的运算与此类似，将5～8和9～12分别看作一块数据用于计算：

在使用9～12的数据计算时，1～4同样被替换出了cache，但由于1～4需要的计算已经完成，故不会有性能损失。

可以看到完成所有运算一共才发生3次cache miss，减少了从主存取数的时间，进而提升了性能。

在例子中，我们分块的大小为32个数。这个大小不一定和cache搬运的块大小相同，并且和具体程序相关。在实践中通常会做一定的实验选取最佳值。

总结

总体来说，分块技巧就是利用了cache的特性进行了加速。避免数据过大造成重复的搬入搬出。

其将数据拆分成可被cache存下的小块，并趁其在cache中时完成尽量多的计算。

其实cache还有很多可以深究的点，比如cache存储的多种策略（直接映射、组相连、全相连）、cache数据替换策略（LFU、LRU）、多级cache等。本文也忽略了输入矩阵和输出向量的cache miss。但该例子已足够我们学习分块技巧的核心思想了。

为了简单，本文例子还不能完全体现出分块技巧的威力。下面列举一个矩阵乘矩阵的例子，感兴趣的读者可以自行研究。

#define M (1024)
#define N (1024)
#define K (1024)

static void naive_mat(float mat_in1[M][K],  float mat_in2[K][N], float mat_out[M][N]){
    for(int m = 0; m < M; m++){
        for(int n = 0; n < N; n++){
            for(int k = 0; k < K; k++){
                mat_out[m][n] += mat_in1[m][k] * mat_in2[k][n];
            }
        }
    }
}

程序耗时：1663.65 ms

#define M (1024)
#define N (1024)
#define K (1024)
#define BLK_SZ (32)
#define BLK_SZ_K (8)

static void blocking_mat(float mat_in1[M][K],  float mat_in2[K][N], float mat_out[M][N]){
    for(int m0 = 0; m0 < M / BLK_SZ; m0++){
        for(int n0 = 0; n0 < N / BLK_SZ; n0++){
            for(int k0 = 0; k0 < K / BLK_SZ_K; k0++){
                for(int m1 = 0; m1 < BLK_SZ; m1++){
                    for(int n1 = 0; n1 < BLK_SZ; n1++){
                        for(int k1 = 0; k1 < BLK_SZ_K; k1++){
                            mat_out[m0 * BLK_SZ + m1][n0 *BLK_SZ + n1] += mat_in1[m0 * BLK_SZ +m1][k0 * BLK_SZ_K +k1] * mat_in2[k0 * BLK_SZ_K +k1][n0 * BLK_SZ + n1];
                        }
                    }
                }
            }
        }
    }
}

程序耗时：144.27 ms
性能提升超11倍。

所以性能提升百倍似乎也不是很难嘛，只要再想办法提升10倍就可以了～

除本文讲的分块技巧外，其实还有很多程序优化手段。比如数据重排、循环展开、SIMD指令优化、汇编优化、多线程、量化以及现在前沿的自动优化+代码生成（TVM、Halide）等。在硬件上，还可以选用算力更强的硬件（DSP、GPU、FPGA、ASIC）来加速，如今火热的神经网络加速器就是其中之一。

那么，敬请期待更多相关介绍吧～