跳转至

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 ] = \sum_{i=0}^n\sum_{j=0}^m\sum_{k=0}^{l}A[i, k] * B[k, j]+C[i,j]

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

在这里插入图片描述

n,m,l分别表示矩阵的维度,n\times l的A矩阵和l\times m的B矩阵先做矩阵乘法运算,然后在通过和n\times 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[i_0, i_1, ..., i_9],最终我只希望得到一个6维的向量,那么其中有4层的for-loop就可以被reduce掉。可能矩阵的乘法并不能看到他的优点,当我们要去写一个非常简单的卷积的时候,就可以看到reduce带来的优势了。这里用一个数字图像处理中的简单卷积举例子(input feature map的channel是1, output feature map的channel也是1),算法的描述如下所示,input是一个n*n的卷积,卷积核的大小是5*5,output是通过te.compute计算得到。

在这里插入图片描述

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

Output[d_i,d_j] = \sum_{d_i=0}^{5}\sum_{d_j=0}^{5}Input[d_i, d_j] * Filter[d_i,d_j]

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

那么表达式就被更新为:

Output[i,j] = \sum_{i=0}^n\sum_{j=0}^m\sum_{d_i=0}^{5}\sum_{d_j=0}^{5}Input[i + d_i, j + d_j] * Filter[d_i,d_j]

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

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

来看下面的例子,比如我有两个数组 A_0[i,j], A_1[i,j] ,那么 B_0[i,j] = A_0[i, j]*3, B_1[i,j] = A_1[i,j] + 5 ,A数组具有相同的维度,长度都为n。那么如果放到C/C++的实现,就是写两层循环循环分别给 B_0,B_1 数组赋值。那么,用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 = \frac{kernel\_size - 1}{2} ),先来讲讲补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。