陳天奇：深度學習編譯技術的現狀和未來

傳統的深度學習框架採用人工優化運算元，然後建立運行時圖解釋器來解決內存分配調度等問題。深度學習編譯器技術路線一般指在優化過程中採用了自動或者半自動的代碼生成用以替代人工優化。深度學習編譯器無疑是最近非常熱門的話題。本文主要探討深度學習編譯技術的現狀和未來。

為什麼需要深度學習編譯器

深度學習編譯器的部署目標傳統的深度學習框架也可以做，一個非常自然的問題是為什麼不直接沿用傳統的框架。這是一個編譯器研究者來往往會忽略的問題。深度學習編譯器只有在各種場景超過人工優化的傳統辦法，才有機會真正被採用，到達這一目標之前之前深度學習編譯只是玩具。

從目前的現狀來看，深度學習編譯器TVM已經一定程度上到達了這一目標。在一些部署場景下，深度學習編譯已經到達了可以和傳統框架一拼高下或者超越傳統框架的階段。隨著深度學習自動化編譯研究的進展手寫優化的經驗會被融入到編譯器中，從而把逐步替代傳統的方案。為什麼會如此呢？總結下來的核心是編譯器可以帶來的更多自動化。細化來說分以下幾點：

無限的算力和有限的精力

如果給定一個特定的運算元，一個工程師無疑可以做到非常精細的地步，通過精細地選擇流水線，指令集，預讀，寄存器分配來到達接近peak的效果。深度學習編譯器需要一個大的搜索空間涵蓋手寫優化的方案。一旦搜索空間定義足夠大，從而接近人工。很多人會覺得，一個足夠精力的工程師一定可以超過編譯，這一點無疑是對的。那麼機器（自動編譯）的優勢在什麼地方呢？機器的本質優勢在於其強大的算力，可以針對目標網路的每一層，特定輸入大小場景，進行專門的優化。而人的精力有限，一般會嘗試優化常見的瓶頸運算元，而這種有針對性的優化未必適用於網路的每一層，或者企業相關的應用場景。通過（接近無限）的算力去適配每一個應用場景看到的網路，這是深度學習編譯器比人工路線強的地方。

當比較TVM和傳統方法的時候的時候，我們往往會發現：在標準的benchmark（如imagenet resnet）上，編譯帶來的提升可能只在10%到20%，但是一旦模型相對不標準化（如最近的OctConv，Deformable, 甚至是同樣的resnet不同的輸入尺寸），編譯技術可以帶來的提升會非常巨大。原因也非常簡單，有限的精力使得參與優化的人往往關注有限的公開標準benchmark，但是我們的部署目標往往並非這些benchmark，自動化可以無差別地對我們關心的場景進行特殊優化。接近無限的算力和有限的精力的差別正是為什麼編譯技術一定會越來越重要的原因。

編譯器可以站在不一樣的起跑線上

如果比拼優化一個固定某一層的3x3卷積，往往精心的手工優化是機器未必可以超越的。那麼在一個標準的benchmark如resnet上，為什麼編譯技術還是可以超過人呢？其原因是往往編譯器可以和人站在不一樣的起跑線上。

舉一個簡單的例子，如果優化一個n層的MLP，我們要優化矩陣乘法。從人工優化的角度來說，往往我們會限制優化目標為，優化一個行優先（row-major）存儲的矩陣乘法。但是對於神經網端到端絡本身而言，內部到底是行優先或者是列優先，或者是（NCHWc4, NCHWc8）都是可以的，而且往往對於每一層最好的排布也不一定相同。同樣的，我們往往可以把運算元和後面幾層的各種運算元融合起來，或者針對每一層選擇有利於其的量化方案。還是因為精力有限的緣故，人工優化的庫往往會把全局問題限制在一些子問題（如行優先的矩陣乘法）上。如果編譯器和工程師硬碰硬直接解決同樣的子問題，或許並不能討好。但是自動化後的編譯器可以直接去考慮更大的解決空間，去自動選擇更加高效的數據排布或者運算元融合，從站在了不一樣的起跑線上。這也是為什麼即使在標準benchmark，只要編譯可以做到人工的80%到90%，更好的起跑線帶來的優勢會掩蓋單個部分微弱的劣勢。

編譯器和手工優化結合

最後，編譯的目標並非替代手工優化，而是吸收手工優化的經驗，使得優化更加自動化。一個笨編譯器比不過聰明的腦子。怎麼辦呢，我們需要從實用主義的角度出發給編譯技術提供定製，允許在需要的加入手工優化來助力編譯器。

其中最簡單的辦法當然是直接把一些層offload給類似於cudnn這樣的庫。這也是XLA等在內的工具採取的技術路線。更進一步，TVM允許手工提供微內核（micro kernel）用於優化4x4外積等，但是依然採用自動優化的辦法優化內存排布和loop。達到手工和自動化相結合的目標。類似的，用戶可以通過構造特定的搜索空間模版來加入人工信息。

總結下來，深度學習編譯器之所以可以到達今天的高度，正是得益於深度學習優化工程師的經驗總結。接下來的一段時間，我們應該會看到越來越多優化工程師加入到深度學習編譯器建設中去，使得兩條路線逐漸融為一體。

深度學習編譯和傳統編譯的技術路線差別

在優化需求上深度學習編譯器和傳統編譯器有很大的差別。傳統編譯器注重於優化寄存器使用和指令集匹配，其優化往往偏向於局部。而深度學習編譯器的優化往往需要涉及到全局的改寫，包括之前提到的內存，運算元融合等。目前深度學習框架的圖優化或者高層優化（HLO）部分和傳統編譯的pass比較匹配，這些優化也會逐漸被標準的pass所替代。但是在高層還會有開放的問題，即高層的抽象如何可以做到容易分析又有足夠的表達能力。TVM的Relay，XLA和Glow是三個在這個方向上的例子。

在自動代碼生成上，傳統編譯器的目標是生成比較優化的通用代碼。而深度學習編譯器的目標是生成接近手寫或者更加高效的特定代碼（卷積，矩陣乘法等）。相對的，在一些情況下深度學習編譯器可以花費更多的時間（去尋找這些解決方案。

整數集分析和Polyhedral Model

Polyhedral model是一個將近研究了十年的領域，其核心思想是採用整數集來表示循環迭代的範圍，利用整數集之間的關係來表示迭代變數之間的依賴，從而達到程序分析變換的目的。傳統的poly方法採用了線性約束來表示整數集和整數集之間的關係。Polyhedral方法也是被很多人覺得有希望用於優化深度學習運算元的方法之一。Poly是真的比較重要呢？這個問題可以分兩點來看：

從核心思想上來看，poly的核心思想是整數集抽象和分析。我們也可以稱之為廣義的poly思想。整數集抽象是一個非常值得所有深度學習編譯器採納的抽象。包括TVM，TC，MLIR等在內的各個方案都引入了整數集抽象。具體的技術路線差別會在於如何表示以及整合整數集來進行分析，在這一點上不同的框架的做法有所不同。

當然，如果直接研究傳統的poly文獻，比較狹義上來說的poly包含了一套基於線性約束分析整數集的方法和搜索空間。線性約束空間解法帶來了一些效率上和對於整數集關係的限制。當然採用線性約束求解的好處是可以解決像三角形或者平行四邊形約束的循壞，但是這一類循壞並不常見。反過來說，如何更好的優化規則循環成為更加緊要的問題。因為模型本身的限制，狹義的poly本身不能完全解決搜索空間的問題。

總的來說，直接採用狹義的polyhedral技術並不能完全解決深度學習編譯問題，但是整數集分析（廣義的poly）已經被廣泛地採用於各個方案之中。

包容萬象的搜索空間

搜索空間決定了深度學習編譯器能力的上界，如何設計搜索空間是所有深度學習編譯器需要仔細考慮的話題。很多目前的編譯方案還是採用了有限的規則生成搜索空間。一般來說，這一搜索空間的定義需要大量吸收人工優化經驗並且加以融入。具體的空間包括循環重排，映射部分計算到實際的加速器指令（張量化，tensorzation），流水線優化等。一般很難確定一個完整的解答，以TVM為例，這一搜索空間會在過去和接下來的幾年裡面通過社區不斷迭代，到達越來越好的效果。

用機器學習優化機器學習

搜索空間本身確立之後剩下的問題是如何快速地找到比較好的運算元。我們可以採用機器學習來自動優化運算元。有興趣的同學可以看 AutoTVM：讓AI來編譯優化AI系統底層運算元

編譯器之外的架構

需要指出的是，深度學習編譯器和傳統編譯器的一大區別是出了編譯器之外還有很多周邊的架構用於支持編譯和部署。這些架構包括快速的遠程部署，不同前端的轉化，靈活的運行時設計等。能否有比較好的周邊架構也是會決定深度學習編譯器易用性的關鍵。

AI晶元和編譯器的協同設計

AI晶元需要編譯器嗎？這無疑是一個很有趣的問題。在一定程度上，AI晶元對於編譯器的依賴取決於晶元本身的設計。一般而言，越靈活的晶元對於編譯器的依賴會越大一些。在AI晶元設計之初，大有傳統的CISC風格，把所有的優化在晶元本身解決。但是隨著領域的演化，為了支持更加靈活的需求，AI晶元本身也會在保留張量指令集和特殊內存結構的前提下越來越靈活。相信未來的架構師需要晶元和系統協同設計的，自動化也會越來越多地被應用到專用晶元中去。有興趣的同學可以閱讀 VTA: 開源AI晶元棧

學習與開發深度學習編譯器

深度學習編譯器領域處在起步階段，但是已經有了一定的應用場景。未來這一方向的應用會越來越多，也需要更多的人員參與到學習和開發中。因為深度學習編譯器本身處於研究前沿，最好的學習方式依然是看相關的論文和直接參与社區的開發和討論中去。

最後小廣告：TVM社區堅持合作開發智能，易用，靈活，中立的深度學習編譯器。社區堅持開放和實用主義出發技術路線，以解決最終應用場景為目標和大家一起探討深度學習編譯技術未來和實際場景應用。社區論壇經常有關於深度學習編譯器技術路線，代碼解析，和部署實戰的活躍討論，歡迎更多的同學參與到社區中來。

知乎鏈接：https://zhuanlan.zhihu.com/p/65452090

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！