作者:安平博,Xilinx高級工程師;來源:AI加速微信公眾號
Import:將tensorflow,onnx,pytorch等構建的深度學習模型導入,轉化成TVM的中間層表示IR。
Lower:將高層IR表示轉化成低階TIR表示。
Codegen:內存分配和硬件可執行程序生成。
圖導入
通過一個tensorflow的reception網絡來熟悉編譯過程,其它深度學習框架也具有類似過程。從TVM官網可以下載tensorflow的編譯程序
https://tvm.apache.org/docs/tutorials/frontend/from_tensorflow.html#sphx...。主要代碼如下:
模型的輸入是一個后綴為pb的文件,它是神經網絡模型圖的protobuf格式存儲文件。Pb是二進制形式,pbtxt是文本形式。Import_graph_def函數是導入pb,graph是tensorflow的圖結構。
From_tensorflow是將tensorflow的圖結構轉化成TVM的IR。這個函數在文件relay/frontend/tensorflow.py中。函數的調用關系為:
From_tensorflow -> GraphProto.from_tensorflow -> self._get_relay_func。
在get_relay_func中會遍歷每個tensorflow的節點,轉換成tvm的IR表示。重點關注_backtrack_construct函數。
繼續深入和算子轉化有關的函數調用為:_convert_operator -> convert_map。Convert_map中對應了可支持tensorflow算子到tvm算子的轉換關系。
完成了tensorflow到TVM算子轉化后,我們就得到了一個IRModule。我們可以利用tvm的可視化來打印出轉化后的圖:
Main是主函數入口,在TVM中以函數形式反應了tensorflow的圖結構。函數的調用關系反應了圖的依賴關系。
編譯
Python中主要代碼位于relay/build_module.py文件中,調用關系為build -> BuildModule -> build。在build中通過字典獲得了C++中的相應函數。
這里不明白如何通過self.mod[“build”]得到C++中函數的。_BuildModule()是C++中注冊到環境中的一個函數。在src/relay/backend/build_module.cc中,
TVM_REGISTER_GLOBAL是將C++函數注冊到一個全局map中。當python加載編譯好的動態庫時,會自動查詢map中靜態注冊的函數,并添加到python模塊當中。
真正build操作位于RelayBuildModule類中,在其中有一個GetFunction函數,會通過名字查詢要使用的函數,打包成PackedFunc返回,這個函數可能和self.mod[“build”]有關。PackedFunc是TVM中提供的python的一個接口,任何函數都可以封裝成PackedFunc,并給python調用。更詳細介紹可看:https://hjchen2.github.io/2020/01/10/TVM-PackedFunc%E5%AE%9E%E7%8E%B0%E6...
繼續深入代碼,Build -> BuildRelay。這是編譯的主要代碼。其過程包括optimize,codgen。
Optimize就是執行一些優化passes,這些passes包括常數折疊,算符融合等。之后會調用graph_codegen->Codegen。Codegen中實現了內存分配和硬件代碼生成。
審核編輯:何安
-
TVM
+關注
關注
0文章
19瀏覽量
3652
發布評論請先 登錄
相關推薦
評論