正文 

当前位置： 首页 > 开发 > java教程

tvm巡礼howtooptimizecpu(x86)gemm串讲(代码片段)

just_sort just_sort     2022-11-23     418

关键词：

【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∗在x86的docker中构建tvm的arm环境(代码片段)

...x86-ubuntu与arm-ubuntu能否ping通4.3调用RPC4.4ACL的使用5.arm版的tvm编译和运行时环境5.1构建arm版的tvm编译和运行时环境5.2关于ubuntu18.04升级glibc掉入坑里这件事5.3验证安 查看详情

在x86的docker中构建tvm的arm环境(代码片段)

...x86-ubuntu与arm-ubuntu能否ping通4.3调用RPC4.4ACL的使用5.arm版的tvm编译和运行时环境5.1构建arm版的tvm编译和运行时环境5.2关于ubuntu18.04升级glibc掉入坑里这件事5.3验证安装是否成功结束语前言  本篇文章介绍如何在x86的docker中构建tvm的... 查看详情

职业规划：专属程序员的巡礼之年

百度百科：巡礼之年是由三组钢琴独奏组曲组成的一套专辑，由法兰兹·李斯特（FranzLiszt）作曲。一般被认为是李斯特最杰出的作品，为李斯特音乐风格的集大成之作。第三组曲特别被认为代表了李斯特的晚期风格。文章起名为... 查看详情

足球队巡礼-英冠​女王公园巡游者qpr

足球队巡礼历史文章，《足球队巡礼-英超西汉姆联》这次我们了解一支可能有些小众的足球队，来自英冠的女王公园巡游者足球俱乐部，QueensParkRangersF.C.，俗称QPR。他是一家位于英格兰大伦敦哈默史密斯-富勒姆... 查看详情

关于tvm

偶然间对arm中mali显示核心感兴趣，找到的TVM。将了解到的信息做个备忘。TVM是TensorVirtualMachine的所写？官网上，TVM定义自己为一种 IntermediateRepresentation，使用python来做机器设计，然后通过TVM，部署到我们的目标硬件。 比... 查看详情

Amazon S3 - 匿名 TVM 与身份 TVM

】AmazonS3-匿名TVM与身份TVM【英文标题】：AmazonS3-AnonymousTVMvsIdentityTVM【发布时间】：2013-05-1804:14:48【问题描述】：我正在尝试创建一个类似于facebook应用的社交应用。用户将照片和视频上传到服务器。我正在考虑使用AmazonS3来存储... 查看详情

tvm安装

因为现在NNVM的代码都转移到了TVM中，NNVM代码也不再进行更新，因此选择安装的是TVM。gitclone--recursivehttps://github.com/dmlc/tvmsudoapt-getupdatesudoapt-getinstall-ypythonpython-devpython-setuptoolsgcclibtinfo-devzlib1g-devmkdirbuildcp 查看详情

tvm架构设计

TVM架构设计本文面向希望了解TVM体系结构和/或积极参与项目开发的开发人员。主要内容如下：             示例编译流程概述了TVM将模型的高级概念转换为可部署模块的步骤。 &n... 查看详情

tvm安装（二）：pycharm远程调试

...试环境网上已经有很多答案，这里不重复，但在远程调试TVM时会有一个坑：pycharm不能读取远程机器的环境变量，从而导致cuda出错或各种错误，比如：Checkfailed:(bf!=nullptr)isfalse:target.build.cudaisnotenabled因此需要在pycharm中加入系统变... 查看详情

tvm1介绍/ubuntu安装/向量加法转置举例

文章目录​​什么是TVM​​​​一.TVM安装​​​​1.1需要安装的软件​​​​1.2安装​​​​1.3python环境​​​​1.4C++test​​​​1.5python链接tvm​​​​1.6测试(显示下面表示成功了,因为这里只需要CPU代码优化,因此只使用了LLVM... 查看详情

关于tvm的点滴记录

关于TVM的点滴记录                 查看详情

tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境(代码片段)

文章目录前言1.安装TVM1.1下载源码1.2创建虚拟环境及安装依赖库1.3编译TVM源码1.4验证安装是否成功2.配置vscode3.安装FFINavigator结束语前言  本篇文章介绍一下tvm在linux环境下的安装与编译，以及如何使用vscode来配置tvm的远程... 查看详情

tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境(代码片段)

文章目录前言1.安装TVM1.1下载源码1.2创建虚拟环境及安装依赖库1.3编译TVM源码1.4验证安装是否成功2.配置vscode3.安装FFINavigator结束语前言  本篇文章介绍一下tvm在linux环境下的安装与编译，以及如何使用vscode来配置tvm的远程... 查看详情

tvm概述(代码片段)

TVMTVM是陈天奇领导的一个DL加速框架项目。它处于DL框架（如tensorflow、pytorch）和硬件后端（如CUDA、OpenCL）之间，兼顾了前者的易用性和后者的执行效率。官网：https://tvm.apache.org/代码：https://github.com/apa... 查看详情

足球队巡礼-英超西汉姆联

我很喜欢足球，更喜欢和足球相关的历史，体育场就是见证历史的重要标志，疫情之前，经常游荡朝圣国内和国外的各大球场，无论是北京工人体育场，还是沈阳五里河体育场，又或是曼彻斯特老特拉... 查看详情

足球队巡礼-英超西汉姆联

我很喜欢足球，更喜欢和足球相关的历史，体育场就是见证历史的重要标志，疫情之前，经常游荡朝圣国内和国外的各大球场，无论是北京工人体育场，还是沈阳五里河体育场，又或是曼彻斯特老特拉... 查看详情

tvm图优化graphoptimization

TVM图优化GraphOptimization                           Codegen &n 查看详情

tvm概述(代码片段)

TVMTVM是陈天奇领导的一个DL加速框架项目。它处于DL框架（如tensorflow、pytorch）和硬件后端（如CUDA、OpenCL）之间，兼顾了前者的易用性和后者的执行效率。官网：https://tvm.apache.org/代码：https://github.com/apa... 查看详情