TVM學(xué)習(xí)（二）：算符融合

作者：安平博，Xilinx高級(jí)工程師；來(lái)源：AI加速微信公眾號(hào)

算符融合將多個(gè)計(jì)算單元揉進(jìn)一個(gè)計(jì)算核中進(jìn)行，減少了中間數(shù)據(jù)的搬移，節(jié)省了計(jì)算時(shí)間。TVM中將計(jì)算算符分成四種：

1 injective。一一映射函數(shù)，比如加法，點(diǎn)乘等。

2 reduction。輸入到輸出具有降維性質(zhì)的，比如sum。

3 complex-out。這是計(jì)算比較復(fù)雜的，比如卷積運(yùn)算等。

4 opaque。無(wú)法被融合的算符，比如sort。

根據(jù)以上對(duì)算符的不同類型，TVM提供了三種融合規(guī)則：

從一定角度看，這種融合實(shí)際上是數(shù)據(jù)計(jì)算pipeline化，即兩次計(jì)算中間數(shù)據(jù)不再經(jīng)歷store-load的過(guò)程，而是直接給到下一個(gè)計(jì)算單元完成計(jì)算。

在走入fuse ops代碼之前，還需要了解一些算法基礎(chǔ)知識(shí)。算符融合中應(yīng)用了支配樹(shù)算法。在一個(gè)有向無(wú)環(huán)圖中，對(duì)于一個(gè)節(jié)點(diǎn)n來(lái)說(shuō)，從初始節(jié)點(diǎn)s出發(fā)到達(dá)n的所有路徑都經(jīng)歷一個(gè)節(jié)點(diǎn)m，那么m就是n的支配點(diǎn)。而距離n最近的支配點(diǎn)被稱作立即支配點(diǎn)。以r為樹(shù)根，將所有立即支配點(diǎn)按照支配關(guān)系連接起來(lái)就形成了支配樹(shù)。立即后支配點(diǎn)是從一個(gè)點(diǎn)n出發(fā)所有到終止節(jié)點(diǎn)的路徑中通過(guò)的最近節(jié)點(diǎn)，形成的支配樹(shù)是后支配樹(shù)。

在DAG中，對(duì)于一個(gè)點(diǎn)，所有能到達(dá)它的點(diǎn)在支配樹(shù)中的LCA，就是它支配樹(shù)中的父親。為什么算符融合要建立在后支配樹(shù)的基礎(chǔ)上呢？我猜測(cè)可能是因?yàn)閷?duì)于兩個(gè)可融合算符在DAG中位置分為兩種，一種是父子關(guān)系，那么可以直接執(zhí)行算符融合算法；另外一種是它們之間是后支配關(guān)系。對(duì)于具有后支配關(guān)系的兩個(gè)節(jié)點(diǎn)（n->m），就要判斷未來(lái)路徑上的節(jié)點(diǎn)是否都能夠和點(diǎn)m發(fā)生融合，如果可以，那么n也可以和m發(fā)生融合。比如下圖：

Conv2d要和elemwise add融合，必須判斷它的三個(gè)op是否能和elemwise add融合。

TVM中融合流程分為三步：

1 遍歷relay樹(shù)，建立DAG用于后支配樹(shù)分析；

2 建立后支配樹(shù)；

3 應(yīng)用算符融合算法。

一 建立DAG圖

算符融合代碼在src/relay/transforms/fuse_ops.cc中。其中算符融合也應(yīng)用在常量折疊中。

首先TVM中通過(guò)如下代碼來(lái)遍歷relay樹(shù)結(jié)構(gòu)并建立DAG圖。

VisitExpr可以遞歸的調(diào)用在類IndexedforwardGraph中定義的VisitExpr_函數(shù)，通過(guò)深度優(yōu)先搜索遍歷relay樹(shù)，并且建立DAG圖。深度優(yōu)先搜索是從exit節(jié)點(diǎn)作為根節(jié)點(diǎn)反向搜鎖的，因此搜索樹(shù)是一個(gè)后序搜索樹(shù)。Outputs中保存了一個(gè)節(jié)點(diǎn)的輸入的邊，在構(gòu)建后序支配樹(shù)會(huì)通過(guò)這些輸入邊求取LCA。那么在這個(gè)搜索樹(shù)基礎(chǔ)上應(yīng)用支配樹(shù)算法，就能夠得到一個(gè)后序支配樹(shù)了。在這個(gè)類中針對(duì)不同節(jié)點(diǎn)類型重寫(xiě)visitExpr_函數(shù)，節(jié)點(diǎn)類型有FunctionNode，ConstantNode, CallNode, TuppleNode等。我們來(lái)看CallNode的訪問(wèn)函數(shù)定義：

在最后還會(huì)遞歸調(diào)用ExprVisitor::VisitExpr_函數(shù)，最終將深度優(yōu)先搜索到的節(jié)點(diǎn)按照葉節(jié)點(diǎn)起始順序一次加入DAG圖中。只有ConstantNode的訪問(wèn)函數(shù)中不再調(diào)用VisitExpr_，因?yàn)槌Ａ抗?jié)點(diǎn)應(yīng)該不存在葉節(jié)點(diǎn)了。在callNode中會(huì)將其輸入加入到DAG中，同時(shí)遍歷和輸入以及其op連接的節(jié)點(diǎn)，ExprVisitor中對(duì)CallNode訪問(wèn)函數(shù)定義為：

因?yàn)镋xprVisitor是被IndexForwardGraph繼承的，而VisitExpr_是虛擬函數(shù)，this就會(huì)指向IndexForwardGraph實(shí)例，最終就會(huì)調(diào)用這個(gè)類中定義的VisitExpr_函數(shù)，實(shí)現(xiàn)遞歸的遍歷relay樹(shù)。

這里要關(guān)注一下OpPatternKind，它定義了算子類型，是不同融合算法使用的依據(jù)。其定義在include/tvm/relay/op_attr_types.h文件中。

二 建立后序支配樹(shù)

接下來(lái)看后序支配樹(shù)的構(gòu)建。構(gòu)建函數(shù)是PostDom。因?yàn)楦?jié)點(diǎn)（DAG圖的出口）在post_dfs_order中最后，所以從根節(jié)點(diǎn)開(kāi)始尋找每個(gè)節(jié)點(diǎn)出點(diǎn)的LCA，這個(gè)LCA就是后序支配點(diǎn)。

GetNode函數(shù)是獲得支配點(diǎn)，構(gòu)建支配樹(shù)。在GetNode中，首先初始化根節(jié)點(diǎn)，然后求每個(gè)節(jié)點(diǎn)的輸入節(jié)點(diǎn)的LCA，即是這個(gè)節(jié)點(diǎn)的支配點(diǎn)。

LeastComonAncestor函數(shù)中主要代碼是：

通過(guò)兩兩求節(jié)點(diǎn)的LCA，來(lái)求取所有節(jié)點(diǎn)的LCA。程序會(huì)將計(jì)算圖中的末節(jié)點(diǎn)深度設(shè)置為1。然后向上逐層增加，那么LCA的共同祖先是相同的，深度也一定是一致。遍歷所有的節(jié)點(diǎn)，就得到一個(gè)后向支配樹(shù)。節(jié)點(diǎn)的pattern指向他的LCA。在計(jì)算支配點(diǎn)的pattern的時(shí)候，會(huì)依據(jù)pattern的定義，選擇pattern值最大的作為L(zhǎng)CA的pattern。這塊不是太深入理解?？赡苁瞧涠x的從最小值到最大值pattern可以向下進(jìn)行融合，比如kElemWise=0, kInjective=2, 那么前者就能融合到KInjective中。

三 融合

完成了DAG和postDominator tree構(gòu)建后，就開(kāi)始融合操作。TVM中定義了group結(jié)構(gòu)體，用于表示融合后的圖結(jié)構(gòu)。Group結(jié)構(gòu)體如下：

如果某些算符可以融合，那么就通過(guò)這個(gè)結(jié)構(gòu)體中的parent，master_ref將這些節(jié)點(diǎn)建立連接關(guān)系。Group首先進(jìn)行初始化和DAG相同的圖。然后分別遍歷dag，postDominator tree，以及group圖中節(jié)點(diǎn)，來(lái)判斷算符是否能被融合。Dag中和postDom中對(duì)應(yīng)相同index的節(jié)點(diǎn)分別是被支配點(diǎn)和支配點(diǎn)。主要融合函數(shù)是以下兩個(gè)函數(shù)：

在runFuse中，有幾種情況是不進(jìn)行算符融合的：

1 算符類型是Kopaque的。

2 該節(jié)點(diǎn)不存在支配點(diǎn)。

3 能夠融合的節(jié)點(diǎn)超過(guò)了一定數(shù)量。

融合操作算法基本上是考察當(dāng)前節(jié)點(diǎn)到其支配點(diǎn)所有路徑上的節(jié)點(diǎn)是否都符合融合規(guī)則，如果符合就進(jìn)行融合，不符合就不融合。函數(shù)CheckPath就是用于考察src到sink路徑是否能夠融合的。

融合分成了三個(gè)phase，每個(gè)phase處理不同可融合類型。這里我沒(méi)有深入研究。當(dāng)判斷支配樹(shù)的前后節(jié)點(diǎn)可以融合后，就通過(guò)函數(shù)commitFuse執(zhí)行融合操作。

完成融合之后，會(huì)遍歷節(jié)點(diǎn)創(chuàng)建新的graph。