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

TVM主要的編譯過(guò)程如下圖：

Import：將tensorflow，onnx，pytorch等構(gòu)建的深度學(xué)習(xí)模型導(dǎo)入，轉(zhuǎn)化成TVM的中間層表示IR。

Lower：將高層IR表示轉(zhuǎn)化成低階TIR表示。

Codegen：內(nèi)存分配和硬件可執(zhí)行程序生成。

圖導(dǎo)入

通過(guò)一個(gè)tensorflow的reception網(wǎng)絡(luò)來(lái)熟悉編譯過(guò)程，其它深度學(xué)習(xí)框架也具有類似過(guò)程。從TVM官網(wǎng)可以下載tensorflow的編譯程序

https://tvm.apache.org/docs/tutorials/frontend/from_tensorflow.html#sphx...。主要代碼如下：

模型的輸入是一個(gè)后綴為pb的文件，它是神經(jīng)網(wǎng)絡(luò)模型圖的protobuf格式存儲(chǔ)文件。Pb是二進(jìn)制形式，pbtxt是文本形式。Import_graph_def函數(shù)是導(dǎo)入pb，graph是tensorflow的圖結(jié)構(gòu)。

From_tensorflow是將tensorflow的圖結(jié)構(gòu)轉(zhuǎn)化成TVM的IR。這個(gè)函數(shù)在文件relay/frontend/tensorflow.py中。函數(shù)的調(diào)用關(guān)系為：

From_tensorflow -> GraphProto.from_tensorflow -> self._get_relay_func。

在get_relay_func中會(huì)遍歷每個(gè)tensorflow的節(jié)點(diǎn)，轉(zhuǎn)換成tvm的IR表示。重點(diǎn)關(guān)注_backtrack_construct函數(shù)。

繼續(xù)深入和算子轉(zhuǎn)化有關(guān)的函數(shù)調(diào)用為：_convert_operator -> convert_map。Convert_map中對(duì)應(yīng)了可支持tensorflow算子到tvm算子的轉(zhuǎn)換關(guān)系。

完成了tensorflow到TVM算子轉(zhuǎn)化后，我們就得到了一個(gè)IRModule。我們可以利用tvm的可視化來(lái)打印出轉(zhuǎn)化后的圖：

Main是主函數(shù)入口，在TVM中以函數(shù)形式反應(yīng)了tensorflow的圖結(jié)構(gòu)。函數(shù)的調(diào)用關(guān)系反應(yīng)了圖的依賴關(guān)系。

編譯

Python中主要代碼位于relay/build_module.py文件中，調(diào)用關(guān)系為build -> BuildModule -> build。在build中通過(guò)字典獲得了C++中的相應(yīng)函數(shù)。

這里不明白如何通過(guò)self.mod[“build”]得到C++中函數(shù)的。_BuildModule()是C++中注冊(cè)到環(huán)境中的一個(gè)函數(shù)。在src/relay/backend/build_module.cc中，

TVM_REGISTER_GLOBAL是將C++函數(shù)注冊(cè)到一個(gè)全局map中。當(dāng)python加載編譯好的動(dòng)態(tài)庫(kù)時(shí)，會(huì)自動(dòng)查詢map中靜態(tài)注冊(cè)的函數(shù)，并添加到python模塊當(dāng)中。

真正build操作位于RelayBuildModule類中，在其中有一個(gè)GetFunction函數(shù)，會(huì)通過(guò)名字查詢要使用的函數(shù)，打包成PackedFunc返回，這個(gè)函數(shù)可能和self.mod[“build”]有關(guān)。PackedFunc是TVM中提供的python的一個(gè)接口，任何函數(shù)都可以封裝成PackedFunc，并給python調(diào)用。更詳細(xì)介紹可看：https://hjchen2.github.io/2020/01/10/TVM-PackedFunc%E5%AE%9E%E7%8E%B0%E6...

繼續(xù)深入代碼，Build -> BuildRelay。這是編譯的主要代碼。其過(guò)程包括optimize，codgen。

Optimize就是執(zhí)行一些優(yōu)化passes，這些passes包括常數(shù)折疊，算符融合等。之后會(huì)調(diào)用graph_codegen->Codegen。Codegen中實(shí)現(xiàn)了內(nèi)存分配和硬件代碼生成。