TVM中將計算算符分成四種

作者：安平博，Xilinx高級工程師；來源：AI加速微信公眾號

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

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

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

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

4 opaque。無法被融合的算符，比如sort。

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

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

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

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

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

TVM中融合流程分為三步：

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

2 建立后支配樹；

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

一 建立DAG圖

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

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

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

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

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

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

二 建立后序支配樹

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

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

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

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

三 融合

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

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

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

1 算符類型是Kopaque的。

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

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

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

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

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

審核編輯：何安