TVM 巡礼How to optimize cpu(x86) gemm串讲

Posted just_sort

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了TVM 巡礼How to optimize cpu(x86) gemm串讲相关的知识，希望对你有一定的参考价值。

【GiantPandaCV导语】最近在整理一些编译器方面的基础知识翻译，回顾了一下TVM的Schedule然后想起自己1年前做的一些GEMM相关的实验和探索没做什么总结。所以基于TVM的三个教程也即TVM的三代优化来做对之前的学习做一个简单的总结，在本篇文章中我原创的内容主要集中在对各个Schedule的单独详解以及引入了TIR的伪代码描述帮助读者更好的理解TVM Schedule。还有深入解读了Blocking技术以及Auto Schedule搜出的高性能程序，简要分析了micro Kernel的汇编。另外我制作了一个控制Schedule变量进行测试得到的GFLOPS柱状图，方便大家对比各代优化的效果。
感慨一下，我认为当前做优化的人以及独特的优化技巧都绝不在少数。如果基于通用的自动代码生成就可以达到任意硬件GFLOPS的80-90%，那么深度学习编译器可能就真的“收敛”了。但毕竟基于Ansor在Jetson Nano上搜一下GEMM的优化程序，性能非常差是我实验过的一个事实。无论是编译器还是深度学习框架对某些常见场景有过拟合是可以理解的，这并不妨碍Ansor是一个非常优秀的工作。我个人觉得我们对深度编译器的看法不能过于激进，目前深度学习编译器在代码生成上上仍然处于一个半自动状态，应该没有任何一家做业务的公司可以放心把各种要上线的业务的性能交给编译器来保证，所以应当稳中求快好一点。和性能快慢相比，我个人现阶段最在意的事情是代码是否鲁棒，很多情况下代码写正确且稳定运行的难度不小于一些优化的奇淫技巧。综上，请辩证的看待本篇文章中隐含的个人想法，包括对此提出批评。
如果你读完本文觉得对你有一点用的话欢迎点一个赞，整理这篇文章陆续花了我2天时间。本文相关代码都在：https://github.com/BBuf/tvm_mlir_learn ，欢迎大家star我这个仓库一起学习TVM和MLIR，目前已经有近400个Star了。如果本篇文章点赞能过百，后面会考虑出一篇解读TVM CUDA优化教程的类似文章。请多多点赞，thanks!

0x0. 前言

本文主要梳理一下在21年接触到优化gemm的知识，做一个学习总结。行文的顺序大概为：

  • 介绍本文依赖的硬件环境和本文要完成的任务。
  • 回顾gflops的计算方法并计算本地机器的GFLOPS。
  • RoofLine模型。
  • How to optimize GEMM on CPU 教程讲解。(https://github.com/apache/tvm/blob/main/gallery/how_to/optimize_operators/opt_gemm.py) 并细致梳理每一种优化方法的作用以及它们在IR中的表示。
  • Optimizing Operators with Schedule Templates and AutoTVM 教程讲解。（https://github.com/apache/tvm/blob/main/gallery/tutorial/autotvm_matmul_x86.py）
  • Optimizing Operators with Auto-scheduling 教程讲解。（https://github.com/apache/tvm/blob/main/gallery/tutorial/auto_scheduler_matmul_x86.py）
  • 总结。

先放一张总览图：

0x1. 本文依赖的硬件环境以及本文要完成的任务

本文的相关实验以及测试数据在同一硬件环境以及Ubuntu操作系统中完成，下图为机器的配置硬件信息如下：

注意到此CPU的核心为64，主频的最大值最小值分别为3900和1000Mhz，以及L1d，L2，L3 cache的大小分别为32K，1024K，22528K。

接下来说一下本文要完成的任务是什么？非常简单，本文的任务就是实现 C = A ∗ B C=A*B C=A∗B，其中 C C C的维度是 [ m , n ] [m, n] [m,n]， A A A的维度是 [ m , k ] [m,k] [m,k]， B B B的维度是 [ k , n ] [k, n] [k,n]，那么矩阵乘法的原始实现就是：

// gemm C = A * B + C
void MatrixMultiply(int m, int n, int k, float *a, float *b, float *c)

    for(int i = 0; i < m; i++)
        for (int j=0; j<n; j++ )    
            for (int p=0; p<k; p++ )      
                C(i, j) = C(i, j) + A(i, p) * B(p, j);

为了让这个问题更加简单，在本文我们约定m，n，k的值都为1024。

0x2. 回顾GFLOPS的计算方法并计算本地机器的GFLOPS

FLOPs：注意s小写，是FLoating point OPerations的缩写（s表复数），意指浮点运算数，理解为计算量。可以用来衡量模型的复杂度。针对神经网络模型的复杂度评价，应指的是FLOPs，而非FLOPS。
FLOPS：注意全大写，是floating point operations per second的缩写，意指每秒浮点运算次数，理解为计算速度。是一个衡量硬件性能的指标。比如nvidia官网上列举各个显卡的算力（Compute Capability）用的就是这个指标。

在讲解TVM的几个教程之前，我们需要先来计算一下机器的GFLOPS并以此判断后面几种TVM教程中生成的GEMM代码的性能相比于硬件浮点峰值大概到达了什么水平。

浮点峰值一般是计算单位时间内乘法和加法的最大吞吐量，单位为GFLOPS或者TFLOPS，表示每一秒可以计算的乘法和加法的总次数。我们可以根据硬件的结构计算一个理论的浮点峰值，并且也可以基于汇编指令来实测硬件的浮点峰值。由于对这一款CPU的架构了解比较有限，我们这里就直接来实测浮点峰值，后续的讲解基本也只和实测的浮点峰值有关。

基于高叔叔的https://github.com/pigirons/cpufp 进行测试，这里的基本原理是在循环内安排尽可能多的无数据依赖的乘加汇编指令掩盖因寄存器依赖浪费的时间周期，具体见：https://zhuanlan.zhihu.com/p/28226956。我们执行如下命令进行测试：

git clone git@github.com:pigirons/cpufp.git
sh build.sh
./cpufp num_threads

分别测试num_threads=1，2，4时的GFLOPS，也即单核心，双核心以及四核心的GFLOPS。结果如下：

Thread(s): 1
avx512_vnni int8 perf: 267.7062 GFLOPS.
avx512f fp32 perf: 66.9300 GFLOPS.
avx512f fp64 perf: 33.4678 GFLOPS.
fma fp32 perf: 73.3017 GFLOPS.
fma fp64 perf: 36.6564 GFLOPS.
avx fp32 perf: 36.6528 GFLOPS.
avx fp64 perf: 18.3270 GFLOPS.
sse fp32 perf: 22.3124 GFLOPS.
sse fp64 perf: 11.1580 GFLOPS.

Thread(s): 2
avx512_vnni int8 perf: 535.5030 GFLOPS.
avx512f fp32 perf: 133.8812 GFLOPS.
avx512f fp64 perf: 66.9439 GFLOPS.
fma fp32 perf: 146.6221 GFLOPS.
fma fp64 perf: 73.3141 GFLOPS.
avx fp32 perf: 73.3229 GFLOPS.
avx fp64 perf: 36.6583 GFLOPS.
sse fp32 perf: 44.6286 GFLOPS.
sse fp64 perf: 22.3151 GFLOPS.

Thread(s): 4
avx512_vnni int8 perf: 1070.9130 GFLOPS.
avx512f fp32 perf: 267.7571 GFLOPS.
avx512f fp64 perf: 133.8734 GFLOPS.
fma fp32 perf: 293.1998 GFLOPS.
fma fp64 perf: 146.6209 GFLOPS.
avx fp32 perf: 146.6289 GFLOPS.
avx fp64 perf: 73.3129 GFLOPS.
sse fp32 perf: 89.2652 GFLOPS.
sse fp64 perf: 44.6303 GFLOPS.

容易观察到GFLOPs和CPU核心数是一个正比例的关系。本文的优化只关注在单核心上，后面的程序都使用 os.environ['TVM_NUM_THREADS']=str(1) 将程序绑定在一个CPU核心上。

0x3. RoofLine模型

上面已经介绍了GFLOPs（1GFLOPs=10^9FLOPs）和GFLOPS的概念，这里要简单引入一下计算密度和RoofLine模型。所谓计算密度指的是单位访存量下所需的计算量，单位是 FLOPs/Bytes。对于本文的任务来说，FLOPs即浮点的计算次数为 2 × 1024 × 1024 × 1024 = 2.147 G F L O P s 2 \\times 1024 \\times 1024 \\times 1024 = 2.147GFLOPs 2×1024×1024×1024=2.147GFLOPs，这里的Bytes指的是访存量，在这个例子中访存量是 3 ∗ 1024 ∗ 1024 ∗ s i z e o f ( f l o a t ) = 12 M B = 0.0117 G B 3*1024*1024*sizeof(float)=12MB=0.0117GB 3∗1024∗1024∗sizeof(float)=12MB=0.0117GB，所以这里的计算密度为 2.147 G F L O P s / 0.0117 G B = 183.5 F L O P s / B y t e s 2.147GFLOPs/0.0117GB=183.5FLOPs/Bytes 2.147GFLOPs/0.0117GB=183.5FLOPs/Bytes。而RoofLine模型是一个用来评估程序在硬件上能达到的性能上界的模型，可用下图表示：

注意到我们计算出的计算密度183.5FLOPs/Bytes是远大于单核心的fma fp32 perf: 73.3017 GFLOPS的，所以很显然我们的算子是一个计算密集形算子，那么计算速度的上限就是峰值计算速度不会受制于带宽。那么我们就可以放心的进行接下来的讲解了。

指的一提的是理论上的RoofLine模型和硬件真实的RoofLine模型还有一定的Gap，以及对于矩阵乘来说某些参数的改变可能会让这个算子从计算密集型朝着访存密集形发生改变。推荐商汤田子宸兄的这篇《深度学习模型大小与模型推理速度的探讨》文章，里面对RoofLine模型做了更加详细的解释以及思考。https://zhuanlan.zhihu.com/p/411522457

Whatever，我们现在知道本文的 1024 × 1024 1024\\times 1024 1024×1024规模的矩阵乘是一个计算密集型的算子，那我们就去看看TVM针对这种访存密集型算子相比浮点峰值可以优化到什么水平吧。

0x4. How to optimize GEMM on CPU 教程讲解。

这一篇教程全面讲解了TVM里面对schedule的优化方法，这些优化基本都来自于我以前介绍过的 https://github.com/flame/how-to-optimize-gemm 这个仓库，比如分块，矢量化，循环重排，Array Packing等等技巧。TVM这个教程将所有的优化技巧基本都用上了，然后直接给出了最终的结果。我觉得这对不太了解TVM的人跨度稍大。所以接下来我将逐一对这些技巧进行介绍并且逐步应用于Naive的TVM张量表达式并展示每一种优化技巧对应的TIR的变化。

这一节教程涉及到的优化技巧为：

  • Blocking（分块）
  • Vectorization（向量化）
  • Loop Permutation（调整分块的计算顺序）
  • Array Packing（数据重排）
  • Write cache for blocks（为写操作数据块创建连续内存缓存）
  • Parallel（并发运算）

接下来细说。

0x4.1 Blocking（分块）

Blocking通常也叫Tiling，是优化循环一种常见并且很有效的策略。我们先说一下Blocking是要解决什么问题然后我们来看是怎么解决的。看如下的例子：

int *A = new float [N];
int *B = new float[M];
for(int i = 0; i < N; i++)
	for(int j = 0; j < M; j++) 
		A[i] += B[j];

如果这个例子中M是比较大的，那么整个算法一定是不高效的。原因在于计算机在访问数组时并不是逐个从内存中访问的，而是以CacheLine为单位去访问，CacheLine的大小一般为64K。也就是说当我们访问B[0]的时候，实际上B[0]-B[15]已经在同一个CacheLine中进而被放进了Cache中，这样访问的速度就会更快。但由于CPU的Cache时有限的，当Cache满了之后如果还要新的CacheLine要进来，那么Cache就要把旧的CacheLine驱逐出去。

比如这里M比较大，那么我们在访问B[200]的时候，B[0]对应的CacheLine已经被驱逐出缓存了，这样随着i的递增下次要重新访问B[0]就不得不再次从内存中去加载。因此整个算法由于需要循环访问的数据在缓存中频繁的换入换出，导致性能比较差。

为了解决这个问题，引入了Blocking或者叫Tiling。Blocking就是在总的数据量无法放进Cache的情况下，把总数据分成一个小块去访问，让每个Tile都可以在Cache中。具体的做法就是把一个内层循环分解成 outer loop * inner loop。然后把 outer loop 移到更外层去，从而确保 inner loop 的数据访问一定能在Cache中。

对于下面的这段代码：

for(int i = 0; i < N; i++)
	for(int j = 0; j < M; j++) 
		A[i] += B[j];

当M比较大的时候，我们可以简单计算一下它Cache Miss的次数，对于A来说，Cache Miss的次数为 N / B N / B N/B，其中B的Cache Miss次数为 M / B ∗ N M / B * N M/B∗N，乘以N是因为i每次递增的时候，由于M很大B开头的一部分数据已经被驱逐出缓存了。这样总的Cache Miss次数就为 N / B + M / B ∗ N = N / B ∗ ( M + 1 ) N / B + M / B * N=N / B * (M + 1) N/B+M/B∗N=N/B∗(M+1)。

按照Blocking的解决方案，我们把内层循环for(int j = 0; j < M; j++)按T的大小来进行分解。那么上面的代码变成：

for (int j_o = 0; j_o < M; j_o += T)
	for (int i = 0; i < N; i++) 
		for (int j_i = j_o; j_i < j_o + T; j_i++)
			A[i] += B[j_i];

可以发现这里的最内层j_i循环始终可以满足在缓存中，总的CacheMiss为： N / B ∗ M / T N / B * M / T N/B∗M/T + T / B T / B T/B * M / T M / T M/T = N / B ∗ M / T + M / B = M / B × ( N / T + 1 ) N / B * M / T + M / B = M / B \\times (N / T + 1) N/B∗M/T+M/B=M/B×(N/T+1)。从数量级上看，经过一次Block的优化，Cache Miss的次数减少了T倍。

很明显我们发现循环 i i i现在变成了inner loop，我们同样对它进行分解。那么伪代码变成：

for (int i_o = 0; i_o < N; i_o += T)
	for (int j_o = 0; j_o < M; j_o += T) 
		for (int i_i = i_o; i_i < i_o + T; i_i++)
			for (int j_i = j_o; j_i < j_o + T; j_i++)
				A[i_i] += B[j][i];

最终得到Cache Miss数量级相比于原始的循环将会减少 T 2 T^2 T2倍。这就是Blocking或者Tiling的基本原理，通过分块降低数据访问的Cache Miss从而提升性能。接下来，我们基于TVM的教程看一下对矩阵乘算子进行Tiling得到了什么结果。

我们先补充一下TVM以默认Schedule运行矩阵乘的TIR，对应的实验代码在https://github.com/BBuf/tvm_mlir_learn/blob/main/optimize_gemm/optimize_matmul_in_gemm/tvm_without_tune_default_schedule.py：

@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = "from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True
  buffers = A: Buffer(A_2: Pointer(float32), float32, [1048576], []),
             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])
  buffer_map = A_1: A, B_1: B, C_1: C 
  for (m: int32, 0, 1024) 
    for (n: int32, 0, 1024) 
      C[((m*1024) + n)] = 0f32
      for (k: int32, 0, 1024) 
        let cse_var_2 = (m*1024)
        let cse_var_1 = (cse_var_2 + n)
        C[cse_var_1] = (C[cse_var_1: int32] + (A[(cse_var_2: int32 + k)]*B[((k*1024) + n)]))

可以看到这就是简单的三重for循环嵌套，没有任何优化技巧，自然这份代码的效率也是比较低的，可以当作本文的BaseLine。我根据它的执行时间计算了GFLOPS，结果如下：

0.687 GFLOPS

然后我们看一下加了Blocking优化对应的TIR，代码在https://github.com/BBuf/tvm_mlir_learn/blob/main/optimize_gemm/optimize_matmul_in_gemm/tvm_without_tune_only_blocking.py：

@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = "from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True
  buffers = A: Buffer(A_2: Pointer(float32), float32, [1048576], []),
             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])
  buffer_map = A_1: A, B_1: B, C_1: C 
  for (m.outer: int32, 0, 32) 
    for (n.outer: int32, 0, 32) 
      for (m.inner.init: int32, 0, 32) 
        for (n.inner.init: int32, 0, 32) 
          C[((((m.outer*32768) + (m.inner.init*1024)) + (n.outer*32)) + n.inner.init)] = 0f32
        
      
      for (k.outer: int32, 0, 256) 
        for (k.inner: int32, 0, 4) 
          for (m.inner: int32, 0, 32) 
            for (n.inner: int32, 0, 32) 
              let cse_var_3 = (n.outer*32)
              let cse_var_2 = ((m.outer*32768) + (m.inner*1024))
              let cse_var_1 = ((cse_var_2 + cse_var_3) + n.inner)
              C[cse_var_1] = (C[cse_var_1: int32] + (A[((cse_var_2: int32 + (k.outer*4)) + k.inner)]*B[((((k.outer*4096) + (k.inner*1024)) + cse_var_3: int32) + n.inner)]))

可以看到和我们介绍的Blocking的原理差不多，通过将i, j, k三重循环分别分解为outer和inner两重新的循环对原始计算进行分块，减少Cache Miss。从TIR我们看到矩阵A被分成 32 ∗ 4 32*4 32∗以上是关于TVM 巡礼How to optimize cpu(x86) gemm串讲的主要内容，如果未能解决你的问题，请参考以下文章

How to optimize large state Flink job?

How to Configure Nginx for Optimized Performance

TVM 图优化Graph Optimization

moment 多语言 包问题

How to wipe silicon to CPU 如何给CPU正确涂抹硅脂

How to Limit NodeRunner.exe High Memory, CPU Usage