陳天奇等人提出AutoTVM：讓AI來編譯優化AI系統底層算子

Alexnet 橫空出世的兩年之前，我選擇了深度學習作為大學畢設方向。當時的深度學習并沒有現在那麼火熱，也沒有現有的 GPU 卷積實作，我在實驗室的 GTX 470 上面第一次手寫 GPU 卷積用以支援卷積 RBM，當時針對 CPU 十多倍的加速比讓我興奮不已。雖然最終那段研究經曆并沒有結果，但是計算本身對于機器學習的推動作用卻深深地印在了我的腦海裡。

深度學習系統支是推動進步的引擎，而在引擎的核心，則是像 cuDNN 這樣的針對硬體的高效算子實作。每一個高效的算子庫的背後都是大量工程的工程努力。雖然我自己已經從事了幾年深度學習系統的工作，如何實作高效的硬體代碼始終對于我來說是一個頭疼的問題。最近的深度學習系統趨勢開始進行更多的高層優化：包括探索從不同的資料表示，精度和算子融合等各個方面，而這些優化勢必會帶來更多的算子版本需要去優化，這無疑是一個非常大的工程問題。我們還需要把這個問題乘上包括移動晶片，移動 GPU 和 AI 加速器等等後端硬體的種類。針對每一個平台都去實作一遍底層的優化代價無疑會非常大。

我們希望利用機器學習本身去解決這個問題，以深度學習編譯器 TVM 為基礎自動化地優化深度學習算子性能。寫到這裡，第一個大家會問的問題是為什麼機器有可能會比手工做的優化要好。其實具體原因也很簡單，一個人的精力是有限的，我們可以針對幾個特定的場景去做一些優化。而機器有無限的精力可以針對每一種類型，資料表示，算子融合和實際的輸入尺寸做特定的調整。像算子融合，這樣的優化必須要利用自動代碼生成才可以做到。是以機器在一定程度上可以取巧。

如何可以讓機器比對手寫優化性能

思路非常直接：

1. 建立一個足夠大的搜尋空間，保證可能的人工手寫優化全部包含在這個搜尋空間裡面 
2. 快速地搜尋這個這個空間，擷取優化的實作

這兩個想法很簡單，其實在兩年之前我們基本上就有了這一個解決問題的藍圖，但是執行都非常困難。第一個問題是如何建立足夠大的搜尋空間，使得其可以包含人可能達到的手寫優化。這個在之前 TVM 的介紹裡面我們提到過，我們通過參考改進 Halide 的 schedule 程式設計源語，加入了對于 GPU，加速器的抽象支援，通過一年的努力使得搜尋空間基本可以比對手寫優化的性能。值得指出的是手工的優化可以非常多樣性，并不是所有的優化都可以被簡單的一個固定 pattern 總結。這也是之前很多自動代碼生成架構無法達到比較好的效果的一個原因。我們也需要不斷地總結抽象的優化規律，加入到 TVM 的原語中。另外在必要的時候我們也可以通過 tensorize 這個抽象可以插入一些手工優化的 micro kernel，但是依然使用架構來做外層排程，達到最好的效果。

有個足夠好的搜尋空間，剩下的問題是如何在幾十億的可能性裡面去選擇比較好的實作。這裡有幾種常見的做法。傳統的高性能計算庫會采用 Auto Tuning，也就是把可能的參數都嘗試一遍。這樣做的潛在問題是空間太大之後枚舉開銷過高。另外一種常見的做法是類似于資料庫的做法，針對程式建立一個代價估價函數，然後利用估價函數來搜尋。這個做法可能碰到的主要問題是估價函數不一定可以估計準确，并且針對每個新的硬體特性必須要重新設計估價函數。

我們的利用機器學習來學習程式空間的代價估價函數。具體地說，探索程式在一開始會随機地選取一些設定，直接到硬體上面去運作生成的代碼，再通過得到的回報資料來更新我們的程式代價估計函數。這裡面比較有趣的一點是模型的可遷移性。因為真正的深度學習系統需要優化許多不一樣輸入類型，輸入形狀的算子。一個可遷移的模型可以通過學習已經看到過的算子優化記錄來預測新的目标的代價，導緻最後的搜尋時間可以大幅降低。

其實在這個整個優化的過程中機器學習算法隻是起到了重要但是隻是一部分的作用。在一些如 ARM 我們的先驗模版比較強的情況下，随機搜尋就可以達到非常不錯的效果。但是重要的是我們需要建立叢集化，自動化，python 可以對接的架構來做這些實驗。這裡特别提一下 TVM 最近引入的 RPC tracker，支援了這一個功能，才可以使得我們可以在多種裝置上進行實驗。