TVM源語-Compute篇

【GiantPandaCV導(dǎo)語】使用和魔改TVM也有一段時(shí)間了，其實(shí)很多場景下，都是拿到pytorch的model，然后轉(zhuǎn)成torchscript，通過relay.frontend.from_pytorch導(dǎo)入，然后一步一步在NVIDIA GPU上generate出網(wǎng)絡(luò)中每個(gè)op對應(yīng)的cuda code。但是，當(dāng)我們的場景不在局限在神經(jīng)網(wǎng)絡(luò)的時(shí)候，比如一些由tensor構(gòu)成的密集計(jì)算，就得需要通過tvm的 primitives，也即DSL來定義算法，然后通過AutoTVM或者Ansor來解決問題，當(dāng)然如果要使用Ansor的話，你只需要定義好algorithm是什么樣的，schedule的部分會幫你自動做，當(dāng)然，如果你想得到一個(gè)custom-level的schedule，你不能完全指望Ansor能給你帶來所有，所以關(guān)于tvm primitives的學(xué)習(xí)還是非常重要的。 TVM的設(shè)計(jì)思想是將“compute”和“schedule”進(jìn)行decouple，那么這一片文章就將所有compute有關(guān)的primitives進(jìn)行總結(jié)，下一篇將對schedule有關(guān)的primitives進(jìn)行總結(jié)。

先來從最簡單例子開始，一步一步深入，本篇文章會涉及如下幾個(gè)例子

1. 一維向量的加法 vector_addition
2. 二維矩陣的乘法 gemm
3. 卷積層的實(shí)現(xiàn) conv2d

（一）Vector Addition

先來看第一個(gè)例子，vector_addition。在實(shí)現(xiàn)一個(gè)算法時(shí)，我們需要做的就是將這個(gè)算法的數(shù)學(xué)表達(dá)式寫出來，然后再將其翻譯成為我們熟悉的語言，交給計(jì)算機(jī)去執(zhí)行。

那么vector_addition要做的其實(shí)就是：

,

有了這個(gè)表達(dá)式后。首先需要我們制定數(shù)組的長度為n，然后兩個(gè)數(shù)組A和B，將A和B數(shù)組中對應(yīng)位置元素相加放到數(shù)組C中。來看看在tvm中怎么實(shí)現(xiàn)？

n表示定義的數(shù)組的長度，A，B表示分別開一個(gè)長度為n的數(shù)組，然后通過lambda表達(dá)式將A和B中每個(gè)元素的計(jì)算結(jié)果放進(jìn)C中。關(guān)于te.compute其實(shí)就是你的輸出結(jié)果，第一個(gè)參數(shù)A.shape表示輸出矩陣的shape，lambda i:則可以理解為 for i: 0->n-1，最后通過create_schedule將生成C的過程構(gòu)建出來，這個(gè)構(gòu)建過程其實(shí)就是te.compute做的事情。最后通過tvm.lower將該schedule映射到IR上。我們可以通過print函數(shù)來查看：

是不是和平時(shí)寫的C代碼很像？

（二）GEMM

我們首先寫出GEMM的數(shù)學(xué)表達(dá)式，

我們首先定義維度的矩陣A，維度的矩陣B，維度的矩陣C。來看看TVM的實(shí)現(xiàn)：

n,m,l分別表示矩陣的維度，的A矩陣和的B矩陣先做矩陣乘法運(yùn)算，然后在通過和的C矩陣做相加得到最終的計(jì)算結(jié)果。先看看TVM生成的schedule是什么樣的：

看到第一個(gè)te.compute是做一個(gè)三層的for-loop，也就是我們通常寫兩個(gè)矩陣乘法時(shí)候用到的，不難理解，這里將二維坐標(biāo)的表示拆成了一維坐標(biāo)的形式，其實(shí)不難理解(A[i][j] -> A'[i * width + j])，第二個(gè)te.compute生成的就是對矩陣中每個(gè)對應(yīng)位置的元素的相加。

細(xì)心的同學(xué)可能會發(fā)現(xiàn)，這里出現(xiàn)了一個(gè)新的源語te.reduce_axis。該源語可以說是非常重要的一個(gè)源語，可以幫我們實(shí)現(xiàn)很多算法，特別有必要把這個(gè)reduce拉出來專門講一講。那就先講講reduce這個(gè)操作吧。

我一開學(xué)tvm的時(shí)候，對reduce的認(rèn)識就是“約分”的意思，可能不是非常準(zhǔn)確。就拿矩陣乘法的例子來說， ，可以發(fā)現(xiàn)，在經(jīng)過運(yùn)算后，等號右邊的表達(dá)式有(i, j, k)這三個(gè)維度變成了僅僅只有(i, j)這兩個(gè)維度。當(dāng)然，這樣做的好處是什么？試想有一個(gè)10層for-loop的程序來對一組變量進(jìn)行操作，最終我只希望得到一個(gè)6維的向量，那么其中有4層的for-loop就可以被reduce掉。可能矩陣的乘法并不能看到他的優(yōu)點(diǎn)，當(dāng)我們要去寫一個(gè)非常簡單的卷積的時(shí)候，就可以看到reduce帶來的優(yōu)勢了。這里用一個(gè)數(shù)字圖像處理中的簡單卷積舉例子(input feature map的channel是1， output feature map的channel也是1)，算法的描述如下所示，input是一個(gè)的卷積，卷積核的大小是，output是通過te.compute計(jì)算得到。

來講講上面的寫法，這是一個(gè)非常naive的卷積實(shí)現(xiàn)，不涉及到padding的操作，直接拿著的kernel在一個(gè)的單通道圖像上進(jìn)行濾波，通過數(shù)學(xué)推導(dǎo)，我們可以到針對單一窗口的運(yùn)算結(jié)果：

那么表達(dá)式就被更新為:

因?yàn)樽罱K得到的Output是一個(gè)(n-4) * (n-4)的數(shù)組，那么我們就可以使用reduce來對和進(jìn)行操作。

其實(shí)reduce還是有很多操作需要學(xué)習(xí)的，這里在介紹一下te.compute同時(shí)接受多個(gè)輸入。

來看下面的例子，比如我有兩個(gè)數(shù)組 ,那么 ,  ，A數(shù)組具有相同的維度，長度都為n。那么如果放到C/C++的實(shí)現(xiàn)，就是寫兩層循環(huán)循環(huán)分別給  數(shù)組賦值。那么，用TVM的DSL該怎么實(shí)現(xiàn)呢？

其實(shí)很簡單，看看生成的schedule是什么樣子？

B0,B1的計(jì)算都被統(tǒng)一到兩個(gè)for-loop中了，而不是分開運(yùn)算。當(dāng)然，當(dāng)我們用下面的寫法時(shí)，

那么相對應(yīng)生成的schedule應(yīng)該如下所示：

這種實(shí)現(xiàn)實(shí)際是不高效的，因?yàn)閷τ诰S度相同的for-loop，我們在寫code的時(shí)候，都是盡量將他們放在一起。至于這樣的優(yōu)化是不是適用于所有情況，確實(shí)值得商榷。

(三) 卷積層的實(shí)現(xiàn)

前面在介紹GEMM例子的時(shí)候，我們使用了一個(gè)非常簡單的單通道圖像和濾波器做卷積的例子。然而在深度學(xué)習(xí)中使用卷積的時(shí)候，我們都是使用多個(gè)input channel的input feature map和多個(gè)output channel的feature map，然后對input feature map進(jìn)行padding到合適大小，然后進(jìn)行卷積操作。我們來規(guī)范下conv2d的參數(shù)

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，并且輸入進(jìn)來的維度是64*64的。[512, 256, 3, 3]表示的是卷積核的參數(shù)，output channel是512，input channel是256，必須和input feature map的輸入channel保持一致，然后3乘3表示的是kernel size，pad為1，stride也為1。

OK，有了這些參數(shù)介紹后，我們就可以很容易用TVM的DSL構(gòu)建一套卷積算法的描述了。

卷積第一步要做的就是給input feature map進(jìn)行pad操作，使得其pad后的input feature map在經(jīng)過卷積后，output feature map的尺寸和input feature map的尺寸相同，先來講講補(bǔ)0操作，補(bǔ)0操作在傳統(tǒng)數(shù)字圖像處理中用的也是非常多的。

補(bǔ)0操作，其實(shí)就是在原始的input feature map的上，下，左，右 四個(gè)邊各補(bǔ)了一圈0 (pad=1)，那么原先input feature map中對應(yīng)的Input[0][0]的元素在after padding后就變成了InputPad[1][1]。這樣，我們就可以知道InputPad后哪些element為0，哪些element為1，對應(yīng)生成的schedule如下所示：

補(bǔ)完邊后，接下來就是來做conv2d操作了，由于我們的data layout采用的是 NCHW，所以在用TVM的DSL實(shí)現(xiàn)的過程中，lambda表達(dá)式的循環(huán)順序應(yīng)該是batch->in_channel->height->width。結(jié)合前面講過的一維卷積的例子，針對Filter的三個(gè)維度(out_channel, kernel_size, kernel_size)使用te.reduce_axis操作。

一個(gè)簡單的conv2d算法可以表示成7層for-loop，那么通過三個(gè)reduce_axis操作以后，就會產(chǎn)生剩下的4層for-loop。上圖算法中，B表示batch_size, K表示out_channel, C表示In_channel,Y表示Height, X表示W(wǎng)idth, Fy和Fx分別表示的是kernel_size。那么使用TVM的DSL描述的卷積如下所示：

對應(yīng)的schedule如下所示：

(四)總結(jié)

總結(jié)一下，TVM的DSL其實(shí)包含很多內(nèi)容，和我們平時(shí)寫序列形式的code的style不太一樣，這種寫法更偏向functional programming的模式，當(dāng)然這樣寫的好處也很明顯，通過lambda表達(dá)式和reduce_axis的組合，將for-loop的形式hidden起來，增加大家對于算法的理解，從而讓compiler的后端能更好的優(yōu)化前端通過DSL定義的for-loop。

歡迎關(guān)注GiantPandaCV, 在這里你將看到獨(dú)家的深度學(xué)習(xí)分享，堅(jiān)持原創(chuàng)，每天分享我們學(xué)習(xí)到的新鮮知識。( ? ?ω?? )?

有對文章相關(guān)的問題，或者想要加入交流群，歡迎添加BBuf微信：