聊聊：TVM和深度學習編譯器

• 2019 年 11 月 12 日
• 筆記

本文轉載自：知乎文章和商湯泰坦公開課

作者 | HONG

北大電腦系統結構博士

商湯科技高級研究員

第一次在知乎上發文章，緣由是上周末看到藍色和楊軍兩位大佬對於TVM和XLA的非常有價值的討論，還看到了高叔叔（多年計算庫開發經驗）關於自動程式碼生成和計算庫bare-metal式優化的對比，感覺非常有意思，所以也想談談自己的看法。

我的背景是GPU系統結構和演算法加速，從去年開始接觸TVM有小一年的時間。

這兩篇文章里已經對TVM有了非常好的討論，我想從下面幾個方面談一下：

1. 程式碼生成與自動程式碼生成

2. TVM和XLA的IR區別

3. TVM可能的應用場景

4. 目前來看自動程式碼生成技術有待提高的地方

01

程式碼生成與自動程式碼生成

程式碼生成是傳統編譯器後端模組之一，解決的問題是將優化後的機器無關IR，經過指令選擇、暫存器分配和求值順序等步驟，生成二進位機器程式碼。另外，傳統編譯器里將語法解析器到三地址表示轉化的過程，也稱為中間程式碼生成。

那麼我們來看一下TVM里的程式碼生成做了什麼。TVM的輸入為：tensor、一系列lambda表達式和相應的schedule，然後解析器生成中間表達，中間表達經過一系列編譯優化，最後通過程式碼生成器產生相應的源程式碼或機器程式碼。

事實上，對於NVIDIA GPU而言，有兩種方式輸出程式碼。一種是直接生成CUDA源碼；另一種是通過LLVM生成PTX code，再經過cuda runtime driver編譯成cubin程式碼。

我們對比了這兩種方式的性能差異，發現第一種方式，也就是僅生成源碼的方式，性能甚至會比第二種高出一個數量級。

所以從這個意義上來看，我覺得TVM充當的角色應該是通過將計算和調度分離，在一個大的搜索空間內快速完成程式碼優化，而程式碼編譯的工作還是應該交給設備提供商。

從程式碼生成的角度來看TVM，我覺得這裡的程式碼生成更多的是將一個只有計算沒有調度的程式碼，也就是 lambda表達式，通過一系列優化轉化成一個擁有了一個好的調度的程式碼（例如輸出的CUDA程式碼）。

當然這樣做的前提也和Halide的一樣，Halide自己的定位式DSL，TVM也需要有深度學習的計算多是Tensor類這種嵌套循環計算模式的大前提。

從這個意義上來講，我們將自動程式碼生成拆解成兩部分來看。先看「自動」，我理解自動的意思是將schedule自動加到計算上，而不用人來手工進行優化操作；「程式碼生成」的意思是一種src2src的程式碼轉化。

02

TVM與XLA的IR區別

關鍵詞：TVM的IR

TVM的IR主要針對嵌套循環優化，因此核心部分很清晰。下面節選了部分：

由parser生成IR之後，scheduler和lower部分不斷對IR進行更新，生成的IR已經非常像一段源程式碼了，最後經過字元串拼接轉換成CUDA源程式碼。

因此，TVM IR的特點是專註在嵌套循環部分，有意保留嵌套循環之間的資訊，然後不斷更新嵌套循環關係以及與之相關的memory操作。

上述部分是TVM偏運算元端的IR，偏計算圖端的部分是Relay（NNVM的後繼），Relay部分提供了DAG和A-Normal兩種類型表達計算圖的方式，其中A-Normal是lambda表達式計算續體傳遞風格（CPS）的一種管理性源碼規約，分別供偏好於深度學習和計算語言的人員使用，兩者是基本等價的。

可以看到，TVM對於圖部分的IR和運算元部分的IR，有明顯的分層。

關鍵詞：XLA的IR

XLA的IR為HLO IR，HLO為High Level Optimizer的縮寫。顧名思義，這一層的IR主要描述的是設備無關優化，而設備相關優化會藉助LLVM後端來完成。

HLO IR相比TVM IR最大的區別是：

HLO IR中既表示DAG，又表示加減乘除這些細節的運算，以及相關的輔助功能，比如layout相關的reshape。TVM IR的分層標準是計算圖和運算元，而HLO IR的分層標準是設備無關和設備相關。

XLA IR在優化中，會將一些具名運算元節點（BatchNormalization）直接替換為包含計算細節（+-*/），同時插入一些相關的add、multiply和maximum等節點；或者將另外的具名運算元（Conv）替換為cuDNN API，並且插入相應的call、reshape等節點。接下來，會做一些fusion和dse等優化操作。

更細緻的討論，我們稍後會以學習筆記的形式共享出來。以下是一些相關的優化總結：

03

TVM可能的應用場景

我們將TVM可能的應用場景分成了兩種，一種是離線支援敏捷運算元優化，通過TVM自動程式碼生成的功能，部分解放優化人員手動添加優化的負擔。

具體地，給定一個計算，優化人員進行workload characterization分析完計算訪存pattern之後，可以擬定幾種可能的優化方案，如果TVM的schedule中包含，那麼可以直接應用TVM來生成包含該優化的程式碼。然後藉助profiler來分析程式碼的瓶頸，進一步修正或者更新優化方案，再試圖藉助TVM的自動生成能力，如此循環往複。到最後，如果TVM程式碼的可讀性足夠的話，優化人員可以在其上做手動優化；如果可讀性不好的話，TVM也可以幫助優化人員快速確定優化方案。

另外一種場景是，運行時程式碼生成。

當研究員在試驗新的運算元時，由於沒有可供調用的計算庫，同時缺乏CUDA等高性能程式碼開發經驗，因此很多情況下是通過裸跑python程式碼來調用新運算元的。PyTorch已經在今年5月份開始將TVM接入到PyTorch中實現運行時程式碼優化編譯，目前已經可以運行了。

具體的方式如下圖（來源於PyTorch/TVM的github）：

左圖為PyTorch本來的編譯優化流程，通過script和trace兩種方式，進入到PyTorch IR優化流程中，IR進行優化之後，經由執行引擎來進行執行。

右圖為TVM接入的方式，通過修改PyTorch工程中script和trace的部分，將TVM IR以及相關優化勾到PyTorch IR上，因此首先需要一個Lower部分，做PyTorch IR到TVM IR之間的解析轉化，然後通過TVM IR來做相應優化，最後在編譯到PyTorch IR上。

這樣做的方式很像MLIR中將不同IR連通起來，合作優化的思路，不過這樣連通兩個IR已經有非常大的工程量，如果不同IR的數據結構和實現細節差異較大，多個IR的連通會非常困難，MLIR的部分還沒有非常深入的研究，暫時不做過多評價。

04

自動程式碼生成技術有待提高的地方

自動程式碼生成技術目前仍然有許多需要完善的地方，以下幾點是我個人的思考：

a. 目前TVM生成程式碼在多數情況下需要已知Tensor的形狀，而這使得TVM被拿來做推理引擎以及JIT來使用。這個局限背後的原因是，在做各種各樣的嵌套循環優化時會對循環的邊界進行多次操作，如果循環邊界未知的話，多次操作之後，對於循環邊界的推斷將變得非常麻煩，目前的區間分析技術，還不能完全勝任。因此，支援未知Tensor形狀的程式碼生成，對於拓寬自動程式碼生成技術的應用非常重要，卻又非常難做。

b. TVM需要人為來寫schedule。這個問題在很多文章里都有提及，Halide和TVM也都做了很多嘗試，auto-tuning更多的是去explore對應schedule的參數，如何高效地auto-schedule還是個開放問題。

c. polyhedron model是一個做auto-schedule的思路。不過從Tensor Comprehesion的例子來看，想把它用好還非常困難，在閱讀相關源程式碼時，當時的contributor對裡面的實現也並不滿意，甚至會發出sigh的感嘆。不過由於目前TVM對不完美循環，或者循環變數有依賴的情況沒有辦法處理，這種情況也可以考慮使用polyhedron model來做。

d. 進一步提高程式碼可讀性。

以上是我的一些思考，關於XLA IR、Glow IR和TVM IR到底哪一種更好，或者該如何協同使用，現有深度學習框架中的IR如何借鑒使用這些工作，都是非常開放的問題，最終會收斂到什麼樣子，現在還不能輕下判斷。更多的深度學習編譯器的實踐、體會和學習筆記，供大家批評指正！