TVM學習（二）：算符融合

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

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

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

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

3 complex-out。這是計算比較復雜的，比如卷積運算等。

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

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

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

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

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

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

TVM中融合流程分為三步：

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

2 建立后支配樹；

3 應用算符融合算法。

一 建立DAG圖

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

首先TVM中通過如下代碼來遍歷relay樹結構并建立DAG圖。

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

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

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

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

二 建立后序支配樹

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

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

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

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

三 融合

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

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

在runFuse中，有幾種情況是不進行算符融合的：

1 算符類型是Kopaque的。

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

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

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

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

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