機器之心專欄
選自：TVM社區
Alexnet 橫空出世的兩年之前，我選擇了深度學習作為本科畢設方向。當時的深度學習並沒有現在那麼火熱，也沒有現有的 GPU 卷積實現，我在實驗室的 GTX 470 上面第一次手寫 GPU 卷積用以支持卷積 RBM，當時針對 CPU 十多倍的加速比讓我興奮不已。雖然最終那段研究經歷並沒有結果，但是計算本身對於機器學習的推動作用卻深深地印在了我的腦海裡。
深度學習系統支是推動進步的引擎，而在引擎的核心，則是像 cuDNN 這樣的針對硬件的高效算子實現。每一個高效的算子庫的背後都是大量工程的工程努力。雖然我自己已經從事了幾年深度學習系統的工作，如何實現高效的硬件代碼始終對於我來說是一個頭疼的問題。最近的深度學習系統趨勢開始進行更多的高層優化：包括探索從不同的數據表示，精度和算子融合等各個方面，而這些優化勢必會帶來更多的算子版本需要去優化，這無疑是一個非常大的工程問題。我們還需要把這個問題乘上包括移動芯片，移動 GPU 和 AI 加速器等等後端硬件的種類。針對每一個平臺都去實現一遍底層的優化代價無疑會非常大。
我們希望利用機器學習本身去解決這個問題，以深度學習編譯器 TVM 為基礎自動化地優化深度學習算子性能。寫到這裡，第一個大家會問的問題是為什麼機器有可能會比手工做的優化要好。其實具體原因也很簡單，一個人的精力是有限的，我們可以針對幾個特定的場景去做一些優化。而機器有無限的精力可以針對每一種類型，數據表示，算子融合和實際的輸入尺寸做特定的調整。像算子融合，這樣的優化必須要利用自動代碼生成才可以做到。所以機器在一定程度上可以取巧。
如何可以讓機器匹配手寫優化性能
思路非常直接：
建立一個足夠大的搜索空間，保證可能的人工手寫優化全部包含在這個搜索空間裡面
快速地搜索這個這個空間，獲取優化的實現
這兩個想法很簡單，其實在兩年之前我們基本上就有了這一個解決問題的藍圖，但是執行都非常困難。第一個問題是如何建立足夠大的搜索空間，使得其可以包含人可能達到的手寫優化。這個在之前 TVM 的介紹裡面我們提到過，我們通過參考改進 Halide 的 schedule 編程源語，加入了對於 GPU，加速器的抽象支持，通過一年的努力使得搜索空間基本可以匹配手寫優化的性能。值得指出的是手工的優化可以非常多樣性，並不是所有的優化都可以被簡單的一個固定 pattern 總結。這也是之前很多自動代碼生成框架無法達到比較好的效果的一個原因。我們也需要不斷地總結抽象的優化規律，加入到 TVM 的原語中。另外在必要的時候我們也可以通過 tensorize 這個抽象可以插入一些手工優化的 micro kernel，但是依然使用框架來做外層調度，達到最好的效果。
有個足夠好的搜索空間，剩下的問題是如何在幾十億的可能性裡面去選擇比較好的實現。這裡有幾種常見的做法。傳統的高性能計算庫會採用 Auto Tuning，也就是把可能的參數都嘗試一遍。這樣做的潛在問題是空間太大之後枚舉開銷過高。另外一種常見的做法是類似於數據庫的做法，針對程序建立一個代價估價函數，然後利用估價函數來搜索。這個做法可能碰到的主要問題是估價函數不一定可以估計準確，並且針對每個新的硬件特性必須要重新設計估價函數。
我們的利用機器學習來學習程序空間的代價估價函數。具體地說，探索程序在一開始會隨機地選取一些設定，直接到硬件上面去運行生成的代碼，再通過得到的反饋數據來更新我們的程序代價估計函數。這裡面比較有趣的一點是模型的可遷移性。因為真正的深度學習系統需要優化許多不一樣輸入類型，輸入形狀的算子。一個可遷移的模型可以通過學習已經看到過的算子優化記錄來預測新的目標的代價，導致最後的搜索時間可以大幅降低。
其實在這個整個優化的過程中機器學習算法只是起到了重要但是隻是一部分的作用。在一些如 ARM 我們的先驗模版比較強的情況下，隨機搜索就可以達到非常不錯的效果。但是重要的是我們需要建立集群化，自動化，python 可以對接的框架來做這些實驗。這裡特別提一下 TVM 最近引入的 RPC tracker，支持了這一個功能，才可以使得我們可以在多種設備上進行實驗。
論文鏈接
對於具體實驗效果和方法介紹有興趣的同學可以看論文裡面的結果
機器學習：https://arxiv.org/abs/1805.08166
系統： https://arxiv.org/abs/1802.04799
未來展望
在可見的未來裡面我們還是需要高性能優化的知識不斷充實完善 TVM。手寫算子優化和自動優化可能會在比較長的一段時並存下去。但是我們個人相信更多的自動化和半自動話也會成為主流。文章中提到的搜索空間，RPC 集群支持都已經正式加入到 TVM 中去。自動化模塊也會在下一個版本和大家見面。
圍繞這個問題，這個方向還留下包括模型表示，遷移等在內還有很多未知的問題等待著大家去探索，也歡迎更多的小夥伴加入 TVM 社區。
後記
本次機器學習相關工作主要成員來自華盛頓大學，上海交大和復旦大學。我本人十分感謝交大 acm 學習經歷對於這次工作的影響：十年之前的這個時候我在交大 acm 班正好在做編譯原理大作業。本次論文的主力之一，第二作者鄭憐憫也來自交大，是今年編譯原理的助教。我們當時編譯原理課程主現在依然掛在李沐在交大的主頁上 http://bcmi.sjtu.edu.cn/home/limu/tiger/ 。
✄------------------------------------------------
廣告&商務合作：[email protected]
閱讀更多 坤艮機器之心 的文章
