跳转至

TVM 源语 = Compute 篇

【GiantPandaCV 导语】使用和魔改 TVM 也有一段时间了,其实很多场景下,都是拿到 pytorch 的 model,然后转成 torchscript,通过 relay.frontend.from_pytorch 导入,然后一步一步在 NVIDIA GPU 上 generate 出网络中每个 op 对应的 cuda code。但是,当我们的场景不在局限在神经网络的时候,比如一些由 tensor 构成的密集计算,就得需要通过 tvm 的 primitives,也即 DSL 来定义算法,然后通过 AutoTVM 或者 Ansor 来解决问题,当然如果要使用 Ansor 的话,你只需要定义好 algorithm 是什么样的,schedule 的部分会帮你自动做,当然,如果你想得到一个 custom-level 的 schedule,你不能完全指望 Ansor 能给你带来所有,所以关于 tvm primitives 的学习还是非常重要的。 TVM 的设计思想是将 “compute” 和“schedule”进行 decouple,那么这一片文章就将所有 compute 有关的 primitives 进行总结,下一篇将对 schedule 有关的 primitives 进行总结。

先来从最简单例子开始,一步一步深入,本篇文章会涉及如下几个例子

  1. 一维向量的加法 vector_addition
  2. 二维矩阵的乘法 gemm
  3. 卷积层的实现 conv2d

(一)Vector Addition

先来看第一个例子,vector_addition。在实现一个算法时,我们需要做的就是将这个算法的数学表达式写出来,然后再将其翻译成为我们熟悉的语言,交给计算机去执行。

那么 vector_addition 要做的其实就是:

C[i]=A[i]+B[i],

有了这个表达式后。首先需要我们制定数组的长度为 n,然后两个数组 A 和 B,将 A 和 B 数组中对应位置元素相加放到数组 C 中。来看看在 tvm 中怎么实现?

在这里插入图片描述

n 表示定义的数组的长度,A,B 表示分别开一个长度为 n 的数组,然后通过 lambda 表达式将 A 和 B 中每个元素的计算结果放进 C 中。关于 te.compute 其实就是你的输出结果,第一个参数 A.shape 表示输出矩阵的 shape,lambda i: 则可以理解为 for i: 0->n-1,最后通过 create_schedule 将生成 C 的过程构建出来,这个构建过程其实就是 te.compute 做的事情。最后通过 tvm.lower 将该 schedule 映射到 IR 上。我们可以通过 print 函数来查看:

在这里插入图片描述

是不是和平时写的 C 代码很像?

在这里插入图片描述

(二)GEMM

我们首先写出 GEMM 的数学表达式,

D[i,j]=ni=0mj=0lk=0A[i,k]B[k,j]+C[i,j]

我们首先定义维度N×L 的矩阵 A,维度L×M 的矩阵 B,维度N×M 的矩阵 C。来看看 TVM 的实现:

在这里插入图片描述

n,m,l 分别表示矩阵的维度,n×l 的 A 矩阵和l×m 的 B 矩阵先做矩阵乘法运算,然后在通过和n×m 的 C 矩阵做相加得到最终的计算结果。先看看 TVM 生成的 schedule 是什么样的:

在这里插入图片描述

看到第一个 te.compute 是做一个三层的 for-loop,也就是我们通常写两个矩阵乘法时候用到的,不难理解,这里将二维坐标的表示拆成了一维坐标的形式,其实不难理解 (A[i][j] -> A'[i * width + j]),第二个 te.compute 生成的就是对矩阵中每个对应位置的元素的相加。

细心的同学可能会发现,这里出现了一个新的源语 te.reduce_axis。该源语可以说是非常重要的一个源语,可以帮我们实现很多算法,特别有必要把这个 reduce 拉出来专门讲一讲。那就先讲讲 reduce 这个操作吧。

我一开学 tvm 的时候,对 reduce 的认识就是 “约分” 的意思,可能不是非常准确。就拿矩阵乘法的例子来说, C[i,j]+=A[i,k]B[k,j],可以发现,在经过运算后,等号右边的表达式有 (i, j, k) 这三个维度变成了仅仅只有 (i, j) 这两个维度。当然,这样做的好处是什么?试想有一个 10 层 for-loop 的程序来对一组变量进行操作A[i0,i1,...,i9],最终我只希望得到一个 6 维的向量,那么其中有 4 层的 for-loop 就可以被 reduce 掉。可能矩阵的乘法并不能看到他的优点,当我们要去写一个非常简单的卷积的时候,就可以看到 reduce 带来的优势了。这里用一个数字图像处理中的简单卷积举例子 (input feature map 的 channel 是 1, output feature map 的 channel 也是 1),算法的描述如下所示,input 是一个nn 的卷积,卷积核的大小是55,output 是通过 te.compute 计算得到。

在这里插入图片描述

来讲讲上面的写法,这是一个非常 naive 的卷积实现,不涉及到 padding 的操作,直接拿着55 的 kernel 在一个nn 的单通道图像上进行滤波,通过数学推导,我们可以到针对单一窗口的运算结果:

Output[di,dj]=5di=05dj=0Input[di,dj]Filter[di,dj]

当窗口滑动起来后,就得去改变 (i, j) 的值了,我们只需要在 input[di,dj] 的基础上添加坐标 (i, j) 就行。

那么表达式就被更新为:

Output[i,j]=ni=0mj=05di=05dj=0Input[i+di,j+dj]Filter[di,dj]

因为最终得到的 Output 是一个 (n-4) * (n-4) 的数组,那么我们就可以使用 reduce 来对didj 进行操作。

其实 reduce 还是有很多操作需要学习的,这里在介绍一下 te.compute 同时接受多个输入。

来看下面的例子,比如我有两个数组 A0[i,j],A1[i,j] , 那么 B0[i,j]=A0[i,j]3, B1[i,j]=A1[i,j]+5 ,A 数组具有相同的维度,长度都为 n。那么如果放到 C/C++ 的实现,就是写两层循环循环分别给 B0B1 数组赋值。那么,用 TVM 的 DSL 该怎么实现呢?

在这里插入图片描述

其实很简单,看看生成的 schedule 是什么样子?

在这里插入图片描述

B0,B1 的计算都被统一到两个 for-loop 中了,而不是分开运算。当然,当我们用下面的写法时,

在这里插入图片描述

那么相对应生成的 schedule 应该如下所示:

在这里插入图片描述

这种实现实际是不高效的,因为对于维度相同的 for-loop,我们在写 code 的时候,都是尽量将他们放在一起。至于这样的优化是不是适用于所有情况,确实值得商榷。

(三) 卷积层的实现

前面在介绍 GEMM 例子的时候,我们使用了一个非常简单的单通道图像和滤波器做卷积的例子。然而在深度学习中使用卷积的时候,我们都是使用多个 input channel 的 input feature map 和多个 output channel 的 feature map,然后对 input feature map 进行 padding 到合适大小,然后进行卷积操作。我们来规范下 conv2d 的参数

data layout:NCHW

input feature map:[128, 256, 64, 64]

filter: [512, 256, 3, 3, 1, 1] (pad: 1,stride:1)

解释下,[128, 256, 64, 64] 表示的是,输入的特征图的 batch 为 128,input channel 是 256,并且输入进来的维度是 64*64 的。[512, 256, 3, 3] 表示的是卷积核的参数,output channel 是 512,input channel 是 256,必须和 input feature map 的输入 channel 保持一致,然后 3 乘 3 表示的是 kernel size,pad 为 1,stride 也为 1。

OK,有了这些参数介绍后,我们就可以很容易用 TVM 的 DSL 构建一套卷积算法的描述了。

卷积第一步要做的就是给 input feature map 进行 pad 操作,使得其 pad 后的 input feature map 在经过卷积后,output feature map 的尺寸和 input feature map 的尺寸相同pad=kernel_size12),先来讲讲补 0 操作,补 0 操作在传统数字图像处理中用的也是非常多的。

在这里插入图片描述

补 0 操作,其实就是在原始的 input feature map 的上,下,左,右 四个边各补了一圈 0 (pad=1),那么原先 input feature map 中对应的 Input[0][0]的元素在 after padding 后就变成了 InputPad[1][1],以此类推,在 y 方向和 x 方向的 [1, 64] 出对应的就是原先的 Input(64 * 64)。这样,我们就可以知道 InputPad 后哪些 element 为 0,哪些 element 为 1,对应生成的 schedule 如下所示:

在这里插入图片描述

补完边后,接下来就是来做 conv2d 操作了,由于我们的 data layout 采用的是 NCHW,所以在用 TVM 的 DSL 实现的过程中,lambda 表达式的循环顺序应该是 batch->in_channel->height->width。结合前面讲过的一维卷积的例子,针对 Filter 的三个维度 (out_channel, kernel_size, kernel_size) 使用 te.reduce_axis 操作。

在这里插入图片描述

一个简单的 conv2d 算法可以表示成 7 层 for-loop,那么通过三个 reduce_axis 操作以后,就会产生剩下的 4 层 for-loop。上图算法中,B 表示 batch_size, K 表示 out_channel, C 表示 In_channel,Y 表示 Height, X 表示 Width, Fy 和 Fx 分别表示的是 kernel_size。那么使用 TVM 的 DSL 描述的卷积如下所示:

在这里插入图片描述

对应的 schedule 如下所示:

在这里插入图片描述

(四) 总结

总结一下,TVM 的 DSL 其实包含很多内容,和我们平时写序列形式的 code 的 style 不太一样,这种写法更偏向 functional programming 的模式,当然这样写的好处也很明显,通过 lambda 表达式和 reduce_axis 的组合,将 for-loop 的形式 hidden 起来,增加大家对于算法的理解,从而让 compiler 的后端能更好的优化前端通过 DSL 定义的 for-loop。


本文总阅读量71